Correctly mirror upstream websocket subprotocols back to clients

Previously the first subprotocol was always used automatically,
regardless of the upstream protocol. That was almost always what you
wanted (most cases just use one subprotocol anyway) but could easily
be wrong in some advanced cases. We now specifically match what the
upstream server sends, when forwarding.

When not forwarding, we stick to the same behaviour, replying with the
first suggested protocol, if one is provided.

Author: pimterry

src/rules/websockets/websocket-handlers.ts

@@ -65,6 +65,10 @@ export interface WebSocketHandler extends WebSocketHandlerDefinition {
     ): Promise<void>;
 }
 
+interface InterceptedWebSocketRequest extends http.IncomingMessage {
+    upstreamWebSocketProtocol?: string | false;
+}
+
 interface InterceptedWebSocket extends WebSocket {
     upstreamWebSocket: WebSocket;
 }
@@ -188,7 +192,17 @@ export class PassThroughWebSocketHandler extends PassThroughWebSocketHandlerDefi
     private initializeWsServer() {
         if (this.wsServer) return;
 
-        this.wsServer = new WebSocket.Server({ noServer: true });
+        this.wsServer = new WebSocket.Server({
+            noServer: true,
+            // Mirror subprotocols back to the client:
+            handleProtocols(protocols, request: InterceptedWebSocketRequest) {
+                return request.upstreamWebSocketProtocol
+                    // If there's no upstream socket, default to mirroring the first protocol. This matches
+                    // WS's default behaviour - we could be stricter, but it'd be a breaking change.
+                    ?? protocols.values().next().value
+                    ?? false; // If there were no protocols specific and this is called for some reason
+            },
+        });
         this.wsServer.on('connection', (ws: InterceptedWebSocket) => {
             pipeWebSocket(ws, ws.upstreamWebSocket);
             pipeWebSocket(ws.upstreamWebSocket, ws);
@@ -355,6 +369,9 @@ export class PassThroughWebSocketHandler extends PassThroughWebSocketHandlerDefi
         } as WebSocket.ClientOptions & { lookup: any, maxPayload: number });
 
         upstreamWebSocket.once('open', () => {
+            // Used in the subprotocol selection handler during the upgrade:
+            (req as InterceptedWebSocketRequest).upstreamWebSocketProtocol = upstreamWebSocket.protocol || false;
+
             this.wsServer!.handleUpgrade(req, incomingSocket, head, (ws) => {
                 (<InterceptedWebSocket> ws).upstreamWebSocket = upstreamWebSocket;
                 incomingSocket.emit('ws-upgrade', ws);

test/integration/websockets.spec.ts

@@ -75,7 +75,19 @@ nodeOnly(() => {
 
             // Real server that echoes every message
             wsPort = await portfinder.getPortPromise();
-            wsServer = new WebSocket.Server({ port: wsPort });
+            wsServer = new WebSocket.Server({
+                port: wsPort,
+                handleProtocols: (protocols, request) => {
+                    const protocolIndex = request.headers['echo-ws-protocol-index'];
+                    if (protocolIndex !== undefined) {
+                        return [...protocols.values()][
+                            parseInt(protocolIndex as string)
+                        ];
+                    } else {
+                        return false;
+                    }
+                }
+            });
 
             wsServer.on('connection', (ws, request) => {
                 if (request.headers['echo-headers']) {
@@ -146,17 +158,13 @@ nodeOnly(() => {
             it("forwards the incoming requests's headers", async () => {
                 mockServer.forAnyWebSocket().thenPassThrough();
 
-                const ws = new WebSocket(
-                    `ws://localhost:${wsPort}`,
-                    ['subprotocol-a', 'subprotocol-b'],
-                    {
-                        agent: new HttpProxyAgent(`http://localhost:${mockServer.port}`),
-                        headers: {
-                            'echo-headers': 'true',
-                            'Funky-HEADER-casing': 'Header-Value'
-                        }
+                const ws = new WebSocket(`ws://localhost:${wsPort}`, {
+                    agent: new HttpProxyAgent(`http://localhost:${mockServer.port}`),
+                    headers: {
+                        'echo-headers': 'true',
+                        'Funky-HEADER-casing': 'Header-Value'
                     }
-                );
+                });
 
                 const response = await new Promise<Buffer>((resolve, reject) => {
                     ws.on('message', resolve);
@@ -176,7 +184,43 @@ nodeOnly(() => {
                     [ 'Sec-WebSocket-Version', '13' ],
                     [ 'Connection', 'Upgrade' ],
                     [ 'Upgrade', 'websocket' ],
-                    [ 'Sec-WebSocket-Extensions', 'permessage-deflate; client_max_window_bits' ],
+                    [ 'Sec-WebSocket-Extensions', 'permessage-deflate; client_max_window_bits' ]
+                ]);
+            });
+
+
+            it("forwards the incoming requests' & resulting response's subprotocols", async () => {
+                mockServer.forAnyWebSocket().thenPassThrough();
+
+                const ws = new WebSocket(
+                    `ws://localhost:${wsPort}`,
+                    ['subprotocol-a', 'subprotocol-b'], // Request two sub protocols
+                    {
+                        agent: new HttpProxyAgent(`http://localhost:${mockServer.port}`),
+                        headers: {
+                            'echo-headers': 'true',
+                            'echo-ws-protocol-index': 1 // Server should select index 1 (2nd)
+                        }
+                    }
+                );
+
+                const response = await new Promise<Buffer>((resolve, reject) => {
+                    ws.on('message', resolve);
+                    ws.on('error', reject);
+                });
+
+                // The server's selected subprotocol should be mirrored back to the client:
+                expect(ws.protocol).to.equal('subprotocol-b');
+
+                ws.close(1000);
+
+                const protocolHeaders = JSON.parse(response.toString()).filter(([key]: [key: string]) =>
+                    // The key is random, so we don't check it here.
+                    key == 'Sec-WebSocket-Protocol'
+                );
+
+                // Server should have seen both requested protocols:
+                expect(protocolHeaders).to.deep.equal([
                     [ 'Sec-WebSocket-Protocol', 'subprotocol-a,subprotocol-b' ]
                 ]);
             });