From a9885d83ef4f337568d554ede7dcf6459f03eeec Mon Sep 17 00:00:00 2001
From: Hayden Gascoigne <haydengascoigne@gmail.com>
Date: Fri, 29 May 2020 10:09:23 -0500
Subject: [PATCH] Support custom websocket sub protocols

---
 lib/requestHandler.js | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/lib/requestHandler.js b/lib/requestHandler.js
index 6a6c6c034..7c58ffddb 100644
--- a/lib/requestHandler.js
+++ b/lib/requestHandler.js
@@ -215,6 +215,7 @@ function fetchRemoteResponse(protocol, options, reqData, config) {
 function getWsReqInfo(wsReq) {
   const headers = wsReq.headers || {};
   const host = headers.host;
+  const wsProtocol = headers['sec-websocket-protocol'] || '';
   const hostName = host.split(':')[0];
   const port = host.split(':')[1];
 
@@ -248,7 +249,8 @@ function getWsReqInfo(wsReq) {
     hostName: hostName,
     port: port,
     path: path,
-    protocol: isEncript ? 'wss' : 'ws'
+    protocol: isEncript ? 'wss' : 'ws',
+    wsProtocol: wsProtocol
   };
 }
 /**
@@ -712,7 +714,7 @@ function getWsHandler(userRule, recorder, wsClient, wsReq) {
     const serverInfo = getWsReqInfo(wsReq);
     const serverInfoPort = serverInfo.port ? `:${serverInfo.port}` : '';
     const wsUrl = `${serverInfo.protocol}://${serverInfo.hostName}${serverInfoPort}${serverInfo.path}`;
-    const proxyWs = new WebSocket(wsUrl, '', {
+    const proxyWs = new WebSocket(wsUrl, serverInfo.wsProtocol, {
       rejectUnauthorized: !self.dangerouslyIgnoreUnauthorized,
       headers: serverInfo.noWsHeaders
     });