pass through all the valid headers when proxing the WebSocket, and adds related test cases

test/server/server.js

@@ -7,7 +7,6 @@ const https = require('https');
 const certMgr = require('../../lib/certMgr');
 const fs = require('fs');
 const nurl = require('url');
-const websocket = require('koa-websocket');
 const color = require('colorful');
 const WebSocketServer = require('ws').Server;
 const tls = require('tls');
@@ -65,6 +64,23 @@ function KoaServer() {
     yield next;
   };
 
+  this.logWsRequest = function (wsReq) {
+    const headers = wsReq.headers;
+    const host = headers.host;
+    const isEncript = true && wsReq.connection && wsReq.connection.encrypted;
+    const protocol = isEncript ? 'wss' : 'ws';
+    let key = `${protocol}://${host}${wsReq.url}`;
+    // take proxy data with 'proxy-' + url
+    if (headers['via-proxy'] === 'true') {
+      key = PROXY_KEY_PREFIX + key;
+    }
+
+    self.requestRecordMap[key] = {
+      headers: wsReq.headers,
+      body: ''
+    }
+  };
+
   this.start();
 }
 
@@ -236,11 +252,14 @@ KoaServer.prototype.constructRouter = function () {
   return router;
 };
 
-KoaServer.prototype.constructWsRouter = function () {
-  const wsRouter = KoaRouter();
-  const self = this;
-  wsRouter.get('/test/socket', function *(next) {
-    const ws = this.websocket;
+KoaServer.prototype.createWsServer = function (httpServer) {
+  const wsServer = new WebSocketServer({
+    server: httpServer,
+    path: '/test/socket'
+  });
+  wsServer.on('connection', (ws, wsReq) => {
+    const self = this;
+    self.logWsRequest(wsReq);
     const messageObj = {
       type: 'initial',
       content: 'default message'
@@ -251,10 +270,7 @@ KoaServer.prototype.constructWsRouter = function () {
       printLog('message from request socket: ' + message);
       self.handleRecievedMessage(ws, message);
     });
-    yield next;
-  });
-
-  return wsRouter;
+  })
 };
 
 KoaServer.prototype.getRequestRecord = function (key) {
@@ -277,14 +293,13 @@ KoaServer.prototype.handleRecievedMessage = function (ws, message) {
 KoaServer.prototype.start = function () {
   printLog('Starting the server...');
   const router = this.constructRouter();
-  const wsRouter = this.constructWsRouter();
   const self = this;
   const app = Koa();
-  websocket(app);
 
   app.use(router.routes());
-  app.ws.use(wsRouter.routes());
   this.httpServer = app.listen(DEFAULT_PORT);
+  this.createWsServer(this.httpServer);
+
 
   printLog('HTTP is now listening on port :' + DEFAULT_PORT);
 
@@ -303,7 +318,8 @@ KoaServer.prototype.start = function () {
         server: self.httpsServer
       });
 
-      wss.on('connection', (ws) => {
+      wss.on('connection', (ws, wsReq) => {
+        self.logWsRequest(wsReq);
         ws.on('message', (message) => {
           printLog('received in wss: ' + message);
           self.handleRecievedMessage(ws, message);

test/spec_rule/no_rule_websocket_spec.js

@@ -6,7 +6,7 @@
 const ProxyServerUtil = require('../util/ProxyServerUtil.js');
 const { generateWsUrl, directWs, proxyWs } = require('../util/HttpUtil.js');
 const Server = require('../server/server.js');
-const { printLog, isArrayEqual } = require('../util/CommonUtil.js');
+const { printLog, isArrayEqual, isCommonReqEqual } = require('../util/CommonUtil.js');
 
 testWebsocket('ws');
 testWebsocket('wss');
@@ -26,6 +26,11 @@ function testWebsocket(protocol, masked = false) {
       'Send the message with default option4'
     ];
 
+    const websocketHeaders = {
+      referer: 'https://www.anyproxy.io/websocket/test',
+      origin: 'www.anyproxy.io'
+    }
+
     beforeAll((done) => {
       jasmine.DEFAULT_TIMEOUT_INTERVAL = 200000;
       printLog('Start server for no_rule_websocket_spec');
@@ -47,11 +52,11 @@ function testWebsocket(protocol, masked = false) {
     it('Default websocket option', done => {
       const directMessages = []; // set the flag for direct message, compare when both direct and proxy got message
       const proxyMessages = [];
-      let directHeaders;
-      let proxyHeaders;
+      let directResHeaders;
+      let proxyResHeaders;
 
-      const ws = directWs(url);
-      const proxyWsRef = proxyWs(url);
+      const ws = directWs(url, websocketHeaders);
+      const proxyWsRef = proxyWs(url, websocketHeaders);
       ws.on('open', () => {
         ws.send(testMessageArray[0], masked);
         for (let i = 1; i < testMessageArray.length; i++) {
@@ -74,13 +79,13 @@ function testWebsocket(protocol, masked = false) {
         }
       });
 
-      ws.on('headers', (headers) => {
-        directHeaders = headers;
+      ws.on('upgrade', (res) => {
+        directResHeaders = res.headers;
         compareMessageIfReady();
       });
 
-      proxyWsRef.on('headers', (headers) => {
-        proxyHeaders = headers;
+      proxyWsRef.on('upgrade', (res) => {
+        proxyResHeaders = res.headers;
         compareMessageIfReady();
       });
 
@@ -114,12 +119,13 @@ function testWebsocket(protocol, masked = false) {
         const targetLen = testMessageArray.length;
         if (directMessages.length === targetLen
           && proxyMessages.length === targetLen
-          && directHeaders && proxyHeaders
+          && directResHeaders && proxyResHeaders
         ) {
           expect(isArrayEqual(directMessages, testMessageArray)).toBe(true);
           expect(isArrayEqual(directMessages, proxyMessages)).toBe(true);
-          expect(directHeaders['x-anyproxy-websocket']).toBeUndefined();
-          expect(proxyHeaders['x-anyproxy-websocket']).toBe('true');
+          expect(directResHeaders['x-anyproxy-websocket']).toBeUndefined();
+          expect(proxyResHeaders['x-anyproxy-websocket']).toBe('true');
+          expect(isCommonReqEqual(url, serverInstance)).toBe(true);
           done();
         }
       }

test/util/CommonUtil.js

@@ -120,6 +120,7 @@ function isCommonResHeaderEqual(directHeaders, proxyHeaders, requestUrl) {
 *
 */
 function isCommonReqEqual(url, serverInstance) {
+  console.info('==> trying to get the url ', url);
   try {
     let isEqual = true;
 
@@ -139,6 +140,27 @@ function isCommonReqEqual(url, serverInstance) {
     delete directReqObj.headers['transfer-encoding'];
     delete proxyReqObj.headers['transfer-encoding'];
 
+    // delete the headers that should not be passed by AnyProxy
+    delete directReqObj.headers.connection;
+    delete proxyReqObj.headers.connection;
+
+    // delete the headers related to websocket establishment
+    const directHeaderKeys = Object.keys(directReqObj.headers);
+    directHeaderKeys.forEach((key) => {
+      // if the key matchs 'sec-websocket', delete it
+      if (/sec-websocket/ig.test(key)) {
+        delete directReqObj.headers[key];
+      }
+    });
+
+    const proxyHeaderKeys = Object.keys(proxyReqObj.headers);
+    proxyHeaderKeys.forEach((key) => {
+      // if the key matchs 'sec-websocaket', delete it
+      if (/sec-websocket/ig.test(key)) {
+        delete proxyReqObj.headers[key];
+      }
+    });
+
     isEqual = isEqual && directReqObj.url === proxyReqObj.url;
     isEqual = isEqual && isObjectEqual(directReqObj.headers, proxyReqObj.headers, url);
     isEqual = isEqual && directReqObj.body === proxyReqObj.body;

test/util/HttpUtil.js

@@ -187,17 +187,20 @@ function doUpload(url, method, filepath, formParams, headers = {}, isProxy) {
   return requestTask;
 }
 
-function doWebSocket(url, isProxy) {
+function doWebSocket(url, headers = {}, isProxy) {
   let ws;
   if (isProxy) {
+    headers['via-proxy'] = 'true';
     const agent = new HttpsProxyAgent(SOCKET_PROXY_HOST);
     ws = new WebSocket(url, {
       agent,
-      rejectUnauthorized: false
+      rejectUnauthorized: false,
+      headers
     });
   } else {
     ws = new WebSocket(url, {
-      rejectUnauthorized: false
+      rejectUnauthorized: false,
+      headers
     });
   }
 
@@ -252,12 +255,12 @@ function directOptions(url, headers = {}) {
   return directRequest('OPTIONS', url, {}, headers);
 }
 
-function proxyWs(url) {
-  return doWebSocket(url, true);
+function proxyWs(url, headers) {
+  return doWebSocket(url, headers, true);
 }
 
-function directWs(url) {
-  return doWebSocket(url);
+function directWs(url, headers) {
+  return doWebSocket(url, headers);
 }
 
 /**