add basic support for websocket proxy

Author: codingfishman

bin/anyproxy

@@ -1,6 +1,7 @@
 #!/usr/bin/env node
 
-'use strict'
+'use strict';
+
 const program = require('commander'),
   color = require('colorful'),
   packageInfo = require('../package.json'),
@@ -17,6 +18,7 @@ program
   .option('-i, --intercept', 'intercept(decrypt) https requests when root CA exists')
   .option('-s, --silent', 'do not print anything into terminal')
   .option('-c, --clear', 'clear all the certificates and temp files')
+  .option('--ws-intercept', 'intercept websocket')
   .option('--ignore-unauthorized-ssl', 'ignore all ssl error')
   .parse(process.argv);
 
@@ -62,7 +64,8 @@ if (program.clear) {
       webInterface: {
         enable: true,
         webPort: program.web,
-      },        
+      },
+      wsIntercept: program.wsIntercept,
       forceProxyHttps: program.intercept,
       dangerouslyIgnoreUnauthorized: !!program.ignoreUnauthorizedSsl,
       silent: program.silent
@@ -107,7 +110,7 @@ if (program.clear) {
     } catch (e) {}
     logUtil.printLog(errorTipText, logUtil.T_ERR);
     try {
-      proxyServer && proxyServer.close();        
+      proxyServer && proxyServer.close();
     } catch (e) {}
     process.exit();
   });

lib/httpsServerMgr.js

@@ -9,12 +9,12 @@ const async = require('async'),
   certMgr = require('./certMgr'),
   logUtil = require('./log'),
   util = require('./util'),
+  wsServerMgr = require('./wsServerMgr'),
   co = require('co'),
   constants = require('constants'),
   asyncTask = require('async-task-mgr');
 
 const createSecureContext = tls.createSecureContext || crypto.createSecureContext;
-
 //using sni to avoid multiple ports
 function SNIPrepareCert(serverName, SNICallback) {
   let keyContent,
@@ -80,7 +80,6 @@ function createHttpsServer(config) {
         key: keyContent,
         cert: crtContent
       }, config.handler).listen(config.port);
-
       resolve(server);
     });
   });
@@ -131,6 +130,7 @@ class httpsServerMgr {
     this.instanceDefaultHost = '127.0.0.1';
     this.httpsAsyncTask = new asyncTask();
     this.handler = config.handler;
+    this.wsHandler = config.wsHandler
   }
 
   getSharedHttpsServer(hostname) {
@@ -159,12 +159,15 @@ class httpsServerMgr {
             });
           }
 
+          wsServerMgr.getWsServer({
+            server: httpsServer,
+            connHandler: self.wsHandler
+          });
 
-          httpsServer.on('upgrade', (req, socket, head) => {
-            const reqHost = req.headers.host || 'unknown host';
-            logUtil.printLog(`wss:// is not supported when intercepting https. This request will be closed by AnyProxy. You may either exclude this domain in your rule file, or stop all https intercepting. (${reqHost})`, logUtil.T_ERR);
-            socket.end();
+          httpsServer.on('upgrade', (req, cltSocket, head) => {
+            logUtil.debug('will let WebSocket server to handle the upgrade event');
           });
+
           const result = {
             host: finalHost,
             port: instancePort,

lib/log.js

@@ -25,6 +25,7 @@ function printLog(content, type) {
   if (!ifPrint) {
     return;
   }
+
   const timeString = util.formatDate(new Date(), 'YYYY-MM-DD hh:mm:ss');
   switch (type) {
     case LogLevelMap.tip: {
@@ -62,6 +63,7 @@ function printLog(content, type) {
     }
 
     case LogLevelMap.debug: {
+      console.log(color.cyan(`[AnyProxy Log][${timeString}]: ` + content));
       return;
     }
 
@@ -73,6 +75,27 @@ function printLog(content, type) {
 }
 
 module.exports.printLog = printLog;
+
+module.exports.debug = (content) => {
+  printLog(content, LogLevelMap.debug);
+};
+
+module.exports.info = (content) => {
+  printLog(content, LogLevelMap.tip);
+};
+
+module.exports.warn = (content) => {
+  printLog(content, LogLevelMap.warn);
+};
+
+module.exports.error = (content) => {
+  printLog(content, LogLevelMap.error);
+};
+
+module.exports.ruleError = (content) => {
+  printLog(content, LogLevelMap.rule_error);
+};
+
 module.exports.setPrintStatus = setPrintStatus;
 module.exports.setLogLevel = setLogLevel;
 module.exports.T_TIP = LogLevelMap.tip;

lib/recorder.js

@@ -4,11 +4,13 @@
 const Datastore = require('nedb'),
   path = require('path'),
   fs = require('fs'),
+  logUtil = require('./log'),
   events = require('events'),
   iconv = require('iconv-lite'),
   proxyUtil = require('./util');
 
 const BODY_FILE_PRFIX = 'res_body_';
+const WS_MESSAGE_FILE_PRFIX = 'ws_message_';
 const CACHE_DIR_PREFIX = 'cache_r';
 function getCacheDir() {
   const rand = Math.floor(Math.random() * 1000000),
@@ -85,6 +87,10 @@ class Recorder extends events.EventEmitter {
     }
   }
 
+  emitUpdateLatestWsMessage(id, message) {
+    this.emit('updateLatestWsMsg', message);
+  }
+
   updateRecord(id, info) {
     if (id < 0) return;
     const self = this;
@@ -98,6 +104,27 @@ class Recorder extends events.EventEmitter {
     self.emitUpdate(id, finalInfo);
   }
 
+  /**
+  * This method shall be called at each time there are new message
+  *
+  */
+  updateRecordWsMessage(id, message) {
+    const cachePath = this.cachePath;
+    if (id < 0) return;
+    try {
+      const recordWsMessageFile = path.join(cachePath, WS_MESSAGE_FILE_PRFIX + id);
+      fs.appendFile(recordWsMessageFile, JSON.stringify(message) + ',', () => {});
+    } catch (e) {
+      console.error(e);
+      logUtil.error(e.message + e.stack);
+    }
+
+    this.emitUpdateLatestWsMessage(id, {
+      id: id,
+      message: message
+    });
+  }
+
   updateExtInfo(id, extInfo) {
     const self = this;
     const db = self.db;
@@ -138,6 +165,10 @@ class Recorder extends events.EventEmitter {
     fs.writeFile(bodyFile, info.resBody, () => {});
   }
 
+  /**
+  * get body and websocket file
+  *
+  */
   getBody(id, cb) {
     const self = this;
     const cachePath = self.cachePath;
@@ -159,6 +190,7 @@ class Recorder extends events.EventEmitter {
   getDecodedBody(id, cb) {
     const self = this;
     const result = {
+      method: '',
       type: 'unknown',
       mime: '',
       content: ''
@@ -170,6 +202,9 @@ class Recorder extends events.EventEmitter {
         return;
       }
 
+      // also put the `method` back, so the client can decide whether to load ws messages
+      result.method = doc[0].method;
+
       self.getBody(id, (error, bodyContent) => {
         if (error) {
           cb(error);
@@ -212,6 +247,44 @@ class Recorder extends events.EventEmitter {
     });
   }
 
+  /**
+  * get decoded WebSoket messages
+  *
+  */
+  getDecodedWsMessage(id, cb) {
+    const self = this;
+    const cachePath = self.cachePath;
+
+    if (id < 0) {
+      cb && cb([]);
+    }
+
+    const wsMessageFile = path.join(cachePath, WS_MESSAGE_FILE_PRFIX + id);
+    fs.access(wsMessageFile, fs.F_OK || fs.R_OK, (err) => {
+      if (err) {
+        cb && cb(err);
+      } else {
+        fs.readFile(wsMessageFile, 'utf8', (error, content) => {
+          if (error) {
+            cb && cb(err);
+          }
+
+          try {
+            // remove the last dash "," if it has, since it's redundant
+            // and also add brackets to make it a complete JSON structure
+            content = `[${content.replace(/,$/, '')}]`;
+            const messages = JSON.parse(content);
+            cb(null, messages);
+          } catch (e) {
+            console.error(e);
+            logUtil.error(e.message + e.stack);
+            cb(e);
+          }
+        });
+      }
+    });
+  }
+
   getSingleRecord(id, cb) {
     const self = this;
     const db = self.db;