From 17ac3a298583c6b5660448afc4b53cdc87bdb675 Mon Sep 17 00:00:00 2001
From: Junxiao Shi <sunnylandh@gmail.com>
Date: Mon, 27 Jun 2022 18:16:03 -0400
Subject: [PATCH] Pass WebSocket subprotocol to upstream (#256)

---
 README.md         | 1 +
 index.js          | 7 ++++++-
 test/websocket.js | 7 ++++---
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 720b789..797dc1d 100644
--- a/README.md
+++ b/README.md
@@ -165,6 +165,7 @@ A few things are missing:
 1. forwarding headers as well as `rewriteHeaders`. Note: Only cookie headers are being forwarded
 2. request id logging
 3. support `ignoreTrailingSlash`
+4. forwarding more than one subprotocols. Note: Only the first subprotocol is being forwarded
 
 Pull requests are welcome to finish this feature.
 
diff --git a/index.js b/index.js
index 2b618e3..73936fc 100644
--- a/index.js
+++ b/index.js
@@ -90,6 +90,11 @@ function setupWebSocketProxy (fastify, options, rewritePrefix) {
       return
     }
 
+    const subprotocols = []
+    if (source.protocol) {
+      subprotocols.push(source.protocol)
+    }
+
     let optionsWs = {}
     if (request.headers.cookie) {
       const headers = { cookie: request.headers.cookie }
@@ -100,7 +105,7 @@ function setupWebSocketProxy (fastify, options, rewritePrefix) {
 
     const url = createWebSocketUrl(request)
 
-    const target = new WebSocket(url, optionsWs)
+    const target = new WebSocket(url, subprotocols, optionsWs)
 
     fastify.log.debug({ url: url.href }, 'proxy websocket')
     proxyWebSockets(source, target)
diff --git a/test/websocket.js b/test/websocket.js
index 53eac08..49bdd17 100644
--- a/test/websocket.js
+++ b/test/websocket.js
@@ -8,9 +8,10 @@ const { createServer } = require('http')
 const { promisify } = require('util')
 const { once } = require('events')
 const cookieValue = 'foo=bar'
+const subprotocolValue = 'foo-subprotocol'
 
 test('basic websocket proxy', async (t) => {
-  t.plan(3)
+  t.plan(4)
 
   const origin = createServer()
   const wss = new WebSocket.Server({ server: origin })
@@ -18,6 +19,7 @@ test('basic websocket proxy', async (t) => {
   t.teardown(origin.close.bind(origin))
 
   wss.on('connection', (ws, request) => {
+    t.equal(ws.protocol, subprotocolValue)
     t.equal(request.headers.cookie, cookieValue)
     ws.on('message', (message) => {
       t.equal(message.toString(), 'hello')
@@ -38,7 +40,7 @@ test('basic websocket proxy', async (t) => {
   t.teardown(server.close.bind(server))
 
   const options = { headers: { cookie: cookieValue } }
-  const ws = new WebSocket(`ws://localhost:${server.server.address().port}`, options)
+  const ws = new WebSocket(`ws://localhost:${server.server.address().port}`, [subprotocolValue], options)
 
   await once(ws, 'open')
 
@@ -47,7 +49,6 @@ test('basic websocket proxy', async (t) => {
   stream.write('hello')
 
   const [buf] = await once(stream, 'data')
-
   t.equal(buf.toString(), 'hello')
 
   await Promise.all([