test websocket message hooks

Author: guox191

lib/handlers/wsHandler.js

@@ -233,4 +233,4 @@ module.exports = function getWsHandler(userRule, recorder, wsClient, wsReq) {
     logUtil.debug(e.stack);
     console.error(e);
   }
-}
+}

test/server/server.js

@@ -80,8 +80,11 @@ function KoaServer() {
 
     self.requestRecordMap[key] = {
       headers: wsReq.headers,
-      body: ''
+      body: '',
+      messages: [],
     }
+
+    return self.requestRecordMap[key];
   };
 
   this.start();
@@ -280,18 +283,20 @@ KoaServer.prototype.createWsServer = function (httpServer) {
     path: '/test/socket'
   });
   wsServer.on('connection', (ws, wsReq) => {
-    const self = this;
-    self.logWsRequest(wsReq);
-    const messageObj = {
+    const logRecord = this.logWsRequest(wsReq);
+
+    ws.send(JSON.stringify({
       type: 'initial',
       content: 'default message'
-    };
+    }));
 
-    ws.send(JSON.stringify(messageObj));
     ws.on('message', message => {
       printLog('message from request socket: ' + message);
-      self.handleRecievedMessage(ws, message);
+      this.handleRecievedMessage(ws, message);
+      logRecord.messages.push(message);
     });
+
+    ws.on('error', e => console.error('error happened in websocket server', e));
   })
 };
 
@@ -322,7 +327,6 @@ KoaServer.prototype.start = function () {
   this.httpServer = app.listen(DEFAULT_PORT);
   this.createWsServer(this.httpServer);
 
-
   printLog('HTTP is now listening on port :' + DEFAULT_PORT);
 
   certMgr.getCertificate('localhost', (error, keyContent, crtContent) => {
@@ -336,19 +340,7 @@ KoaServer.prototype.start = function () {
       }, app.callback());
 
       // create wss server
-      const wss = new WebSocketServer({
-        server: self.httpsServer
-      });
-
-      wss.on('connection', (ws, wsReq) => {
-        self.logWsRequest(wsReq);
-        ws.on('message', (message) => {
-          printLog('received in wss: ' + message);
-          self.handleRecievedMessage(ws, message);
-        });
-      });
-
-      wss.on('error', e => console.error('error happened in wss:%s', e));
+      this.createWsServer(self.httpsServer);
 
       self.httpsServer.listen(HTTPS_PORT);
 

test/spec_rule/rule/rule_replace_ws_message.js

@@ -0,0 +1,26 @@
+module.exports = {
+  *summary() {
+    return 'The rule to replace websocket message';
+  },
+
+  *beforeSendWsMessageToClient(requestDetail) {
+    const message = requestDetail.data;
+    try {
+      const messageObject = JSON.parse(message);
+      if (messageObject.type === 'onMessage') {
+        messageObject.content = 'replaced by beforeSendWsMessageToClient';
+        return {
+          data: JSON.stringify(messageObject),
+        }
+      }
+    } catch (err) { /* ignore error */ }
+
+    return null;
+  },
+
+  *beforeSendWsMessageToServer() {
+    return {
+      data: 'replaced by beforeSendWsMessageToServer',
+    };
+  },
+};

test/spec_rule/rule_replace_ws_message_spec.js

@@ -0,0 +1,69 @@
+const async = require('async');
+const ProxyServerUtil = require('../util/ProxyServerUtil.js');
+const TestServer = require('../server/server.js');
+const { printLog } = require('../util/CommonUtil.js');
+const { proxyWs, generateWsUrl } = require('../util/HttpUtil.js');
+const rule = require('./rule/rule_replace_ws_message');
+
+describe('Rule to replace the websocket message', () => {
+  let testServer = null;
+  let proxyServer = null;
+
+  beforeAll((done) => {
+    jasmine.DEFAULT_TIMEOUT_INTERVAL = 50000;
+    printLog('Start server for rule_replace_ws_message_spec');
+
+    testServer = new TestServer();
+    proxyServer = ProxyServerUtil.proxyServerWithRule(rule);
+
+    setTimeout(done, 2000);
+  });
+
+  afterAll(() => {
+    testServer && testServer.close();
+    proxyServer && proxyServer.close();
+    printLog('Close server for rule_replace_ws_message_spec');
+  });
+
+  it('should replace websocket message from server', (done) => {
+    async.mapSeries([
+      { scheme: 'ws', masked: false },
+      { scheme: 'ws', masked: true },
+      { scheme: 'wss', masked: false },
+      { scheme: 'wss', masked: true },
+    ], (unit, callback) => {
+      const url = generateWsUrl(unit.scheme, '/test/socket');
+      const wsClient = proxyWs(url);
+
+      wsClient.on('open', () => {
+        wsClient.send('test', unit.masked);
+      });
+  
+      wsClient.on('message', (message) => {
+        // test beforeSendWsMessageToServer
+        const requestRecord = testServer.getProxyRequestRecord(url);
+        expect(requestRecord.messages[0]).toBe('replaced by beforeSendWsMessageToServer');
+
+        try {
+          const result = JSON.parse(message);
+          if (result.type === 'onMessage') {
+            // test beforeSendWsMessageToClient
+            expect(result.content).toBe('replaced by beforeSendWsMessageToClient');
+            callback();
+          }
+        } catch (err) { /* ignore error */ }
+      });
+  
+      wsClient.on('error', (err) => {
+        printLog('Error happened in proxy websocket');
+        callback(err);
+      });
+    }, (err) => {
+      if (err) {
+        done.fail(err);
+      } else {
+        done();
+      }
+    });
+  });
+});

test/util/CommonUtil.js

@@ -6,7 +6,7 @@ const color = require('colorful');
 
 function _isDeepEqual(source, target) {
   // if the objects are Array
-  if (source.constructor === Array && target.constructor === Array) {
+  if (Array.isArray(source) && Array.isArray(target)) {
     if (source.length !== target.length) {
       return false;
     }