add basic support for websocket proxy

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 {
             });
           }
 
-
-          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();
+          wsServerMgr.getWsServer({
+            server: httpsServer,
+            connHandler: self.wsHandler
           });
+
+          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;

lib/requestHandler.js

@@ -11,6 +11,7 @@ const http = require('http'),
   Stream = require('stream'),
   logUtil = require('./log'),
   co = require('co'),
+  WebSocket = require('ws'),
   HttpsServerMgr = require('./httpsServerMgr'),
   Readable = require('stream').Readable;
 
@@ -175,6 +176,34 @@ function fetchRemoteResponse(protocol, options, reqData, config) {
   });
 }
 
+/**
+* get request info from the ws client, includes:
+ host
+ port
+ path
+ protocol  ws/wss
+
+ @param @required wsClient the ws client of WebSocket
+*
+*/
+function getWsReqInfo(wsClient) {
+  const upgradeReq = wsClient.upgradeReq || {};
+  const header = upgradeReq.headers || {};
+  const host = header.host;
+  const hostName = host.split(':')[0];
+  const port = host.split(':')[1];
+
+  // TODO 如果是windows机器，url是不是全路径？需要对其过滤，取出
+  const path = upgradeReq.url || '/';
+
+  const isEncript = true && upgradeReq.connection && upgradeReq.connection.encrypted;
+  return {
+    hostName: hostName,
+    port: port,
+    path: path,
+    protocol: isEncript ? 'wss' : 'ws'
+  };
+}
 /**
  * get a request handler for http/https server
  *
@@ -452,10 +481,10 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
     const host = req.url.split(':')[0],
       targetPort = req.url.split(':')[1];
     let shouldIntercept;
+    let interceptWsRequest = false;
     let requestDetail;
     let resourceInfo = null;
     let resourceInfoId = -1;
-
     const requestStream = new CommonReadableStream();
 
     /*
@@ -468,14 +497,17 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
     co(function *() {
       // determine whether to use the man-in-the-middle server
       logUtil.printLog(color.green('received https CONNECT request ' + host));
-      if (reqHandlerCtx.forceProxyHttps) {
-        shouldIntercept = true;
-      } else {
-        requestDetail = {
-          host: req.url,
-          _req: req
-        };
-        shouldIntercept = yield userRule.beforeDealHttpsRequest(requestDetail);
+      requestDetail = {
+        host: req.url,
+        _req: req
+      };
+      // the return value in default rule is null
+      // so if the value is null, will take it as final value
+      shouldIntercept = yield userRule.beforeDealHttpsRequest(requestDetail);
+
+      // otherwise, will take the passed in option
+      if (shouldIntercept === null) {
+        shouldIntercept = reqHandlerCtx.forceProxyHttps;
       }
     })
     .then(() =>
@@ -494,7 +526,13 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
             try {
               const chunkString = chunk.toString();
               if (chunkString.indexOf('GET ') === 0) {
-                shouldIntercept = false; //websocket
+                shouldIntercept = false; // websocket, do not intercept
+
+                // if there is '/do-not-proxy' in the request, do not intercept the websocket
+                // to avoid AnyProxy itself be proxied
+                if (reqHandlerCtx.wsIntercept && chunkString.indexOf('GET /do-not-proxy') !== 0) {
+                  interceptWsRequest = true;
+                }
               }
             } catch (e) {
               console.error(e);
@@ -531,10 +569,19 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
     .then(() => {
       // determine the request target
       if (!shouldIntercept) {
-        return {
+        // server info from the original request
+        const orginServer = {
           host,
-          port: (targetPort === 80) ? 443 : targetPort,
+          port: (targetPort === 80) ? 443 : targetPort
         }
+
+        const localHttpServer = {
+          host: 'localhost',
+          port: reqHandlerCtx.httpServerPort
+        }
+
+        // for ws request, redirect them to local ws server
+        return interceptWsRequest ? localHttpServer : orginServer;
       } else {
         return httpsServerMgr.getSharedHttpsServer(host).then(serverInfo => ({ host: serverInfo.host, port: serverInfo.port }));
       }
@@ -594,6 +641,158 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
   }
 }
 
+/**
+* get a websocket event handler
+  @param @required {object} wsClient
+*/
+function getWsHandler(userRule, recorder, wsClient) {
+  const self = this;
+  try {
+    let resourceInfoId = -1;
+    const resourceInfo = {
+      wsMessages: [] // all ws messages go through AnyProxy
+    };
+    const clientMsgQueue = [];
+    const serverInfo = getWsReqInfo(wsClient);
+    const wsUrl = `${serverInfo.protocol}://${serverInfo.hostName}:${serverInfo.port}${serverInfo.path}`;
+    const proxyWs = new WebSocket(wsUrl, '', {
+      rejectUnauthorized: !self.dangerouslyIgnoreUnauthorized
+    });
+
+    if (recorder) {
+      Object.assign(resourceInfo, {
+        host: serverInfo.hostName,
+        method: 'WebSocket',
+        path: serverInfo.path,
+        url: wsUrl,
+        req: wsClient.upgradeReq || {},
+        startTime: new Date().getTime()
+      });
+      resourceInfoId = recorder.appendRecord(resourceInfo);
+    }
+
+    /**
+    * store the messages before the proxy ws is ready
+    */
+    const sendProxyMessage = (event) => {
+      const message = event.data;
+      if (proxyWs.readyState === 1) {
+        // if there still are msg queue consuming, keep it going
+        if (clientMsgQueue.length > 0) {
+          clientMsgQueue.push(message);
+        } else {
+          proxyWs.send(message);
+        }
+      } else {
+        clientMsgQueue.push(message);
+      }
+    }
+
+    /**
+    * consume the message in queue when the proxy ws is not ready yet
+    * will handle them from the first one-by-one
+    */
+    const consumeMsgQueue = () => {
+      while (clientMsgQueue.length > 0) {
+        const message = clientMsgQueue.shift();
+        proxyWs.send(message);
+      }
+    }
+
+    /**
+    * When the source ws is closed, we need to close the target websocket.
+    * If the source ws is normally closed, that is, the code is reserved, we need to transfrom them
+    */
+    const getCloseFromOriginEvent = (event) => {
+      const code = event.code || '';
+      const reason = event.reason || '';
+      let targetCode = '';
+      let targetReason = '';
+      if (code >= 1004 && code <= 1006) {
+        targetCode = 1000; // normal closure
+        targetReason = `Normally closed. The origin ws is closed at code: ${code} and reason: ${reason}`;
+      } else {
+        targetCode = code;
+        targetReason = reason;
+      }
+
+      return {
+        code: targetCode,
+        reason: targetReason
+      }
+    }
+
+    /**
+    * consruct a message Record from message event
+    * @param @required {event} messageEvent the event from websockt.onmessage
+    * @param @required {boolean} isIn whether the message is to or from server
+    *
+    */
+    const recordMessage = (messageEvent, isToServer) => {
+      const message = {
+        time: Date.now(),
+        message: messageEvent.data,
+        isToServer: isToServer
+      };
+
+      // resourceInfo.wsMessages.push(message);
+      recorder && recorder.updateRecordWsMessage(resourceInfoId, message);
+    };
+
+    proxyWs.onopen = () => {
+      consumeMsgQueue();
+    }
+
+    // this event is fired when the connection is build and headers is returned
+    proxyWs.on('headers', (headers, response) => {
+      resourceInfo.endTime = new Date().getTime();
+      resourceInfo.res = { //construct a self-defined res object
+        statusCode: response.statusCode,
+        headers: headers,
+      };
+
+      resourceInfo.statusCode = response.statusCode;
+      resourceInfo.resHeader = headers;
+      resourceInfo.resBody = '';
+      resourceInfo.length = resourceInfo.resBody.length;
+
+      recorder && recorder.updateRecord(resourceInfoId, resourceInfo);
+    });
+
+    proxyWs.onerror = (e) => {
+      // https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent#Status_codes
+      wsClient.close(1001, e.message);
+      proxyWs.close(1001);
+    }
+
+    proxyWs.onmessage = (event) => {
+      recordMessage(event, false);
+      wsClient.readyState === 1 && wsClient.send(event.data);
+    }
+
+    proxyWs.onclose = (event) => {
+      logUtil.debug(`proxy ws closed with code: ${event.code} and reason: ${event.reason}`);
+      const targetCloseInfo = getCloseFromOriginEvent(event);
+      wsClient.readyState !== 3 && wsClient.close(targetCloseInfo.code, targetCloseInfo.reason);
+    }
+
+    wsClient.onmessage = (event) => {
+      recordMessage(event, true);
+      sendProxyMessage(event);
+    }
+
+    wsClient.onclose = (event) => {
+      logUtil.debug(`original ws closed with code: ${event.code} and reason: ${event.reason}`);
+      const targetCloseInfo = getCloseFromOriginEvent(event);
+      proxyWs.readyState !== 3 && proxyWs.close(targetCloseInfo.code, targetCloseInfo.reason);
+    }
+  } catch (e) {
+    logUtil.debug('WebSocket Proxy Error:' + e.message);
+    logUtil.debug(e.stack);
+    console.error(e);
+  }
+}
+
 class RequestHandler {
 
   /**
@@ -602,6 +801,7 @@ class RequestHandler {
    * @param {object} config
    * @param {boolean} config.forceProxyHttps proxy all https requests
    * @param {boolean} config.dangerouslyIgnoreUnauthorized
+     @param {number} config.httpServerPort  the http port AnyProxy do the proxy
    * @param {object} rule
    * @param {Recorder} recorder
    *
@@ -609,22 +809,36 @@ class RequestHandler {
    */
   constructor(config, rule, recorder) {
     const reqHandlerCtx = this;
+    this.forceProxyHttps = false;
+    this.dangerouslyIgnoreUnauthorized = false;
+    this.httpServerPort = '';
+    this.wsIntercept = false;
+
     if (config.forceProxyHttps) {
       this.forceProxyHttps = true;
     }
+
     if (config.dangerouslyIgnoreUnauthorized) {
       this.dangerouslyIgnoreUnauthorized = true;
     }
+
+    if (config.wsIntercept) {
+      this.wsIntercept = config.wsIntercept;
+    }
+
+    this.httpServerPort = config.httpServerPort;
     const default_rule = util.freshRequire('./rule_default');
     const userRule = util.merge(default_rule, rule);
 
     reqHandlerCtx.userRequestHandler = getUserReqHandler.apply(reqHandlerCtx, [userRule, recorder]);
+    reqHandlerCtx.wsHandler = getWsHandler.bind(this, userRule, recorder);
 
     reqHandlerCtx.httpsServerMgr = new HttpsServerMgr({
-      handler: reqHandlerCtx.userRequestHandler
+      handler: reqHandlerCtx.userRequestHandler,
+      wsHandler: reqHandlerCtx.wsHandler // websocket
     });
 
-    this.connectReqHandler = getConnectReqHandler.apply(reqHandlerCtx, [userRule, recorder, reqHandlerCtx.httpsServerMgr])
+    this.connectReqHandler = getConnectReqHandler.apply(reqHandlerCtx, [userRule, recorder, reqHandlerCtx.httpsServerMgr]);
   }
 }
 

lib/rule_default.js

@@ -1,12 +1,12 @@
 'use strict';
 
 module.exports = {
-  
+
   summary: 'the default rule for AnyProxy',
-  
+
   /**
-   * 
-   * 
+   *
+   *
    * @param {object} requestDetail
    * @param {string} requestDetail.protocol
    * @param {object} requestDetail.requestOptions
@@ -23,8 +23,8 @@ module.exports = {
 
 
   /**
-   * 
-   * 
+   *
+   *
    * @param {object} requestDetail
    * @param {object} responseDetail
    */
@@ -34,21 +34,22 @@ module.exports = {
 
 
   /**
-   * 
-   * 
-   * @param {any} requestDetail 
-   * @returns 
+   * default to return null
+   * the user MUST return a boolean when they do implement the interface in rule
+   *
+   * @param {any} requestDetail
+   * @returns
    */
   *beforeDealHttpsRequest(requestDetail) {
-    return false;
+    return null;
   },
 
   /**
-   * 
-   * 
-   * @param {any} requestDetail 
-   * @param {any} error 
-   * @returns 
+   *
+   *
+   * @param {any} requestDetail
+   * @param {any} error
+   * @returns
    */
   *onError(requestDetail, error) {
     return null;
@@ -56,11 +57,11 @@ module.exports = {
 
 
   /**
-   * 
-   * 
-   * @param {any} requestDetail 
-   * @param {any} error 
-   * @returns 
+   *
+   *
+   * @param {any} requestDetail
+   * @param {any} error
+   * @returns
    */
   *onConnectError(requestDetail, error) {
     return null;

lib/util.js

@@ -4,8 +4,8 @@ const fs = require('fs'),
   path = require('path'),
   mime = require('mime-types'),
   color = require('colorful'),
+  crypto = require('crypto'),
   Buffer = require('buffer').Buffer,
-  configUtil = require('./configUtil'),
   logUtil = require('./log');
 const networkInterfaces = require('os').networkInterfaces();
 
@@ -35,7 +35,7 @@ function getUserHome() {
 module.exports.getUserHome = getUserHome;
 
 function getAnyProxyHome() {
-  const home = configUtil.getAnyProxyHome();
+  const home = path.join(getUserHome(), '/.anyproxy/');
   if (!fs.existsSync(home)) {
     fs.mkdirSync(home);
   }
@@ -306,3 +306,18 @@ module.exports.isIpDomain = function (domain) {
 
   return ipReg.test(domain);
 };
+
+/**
+* To generic a Sec-WebSocket-Accept value
+* 1. append the `Sec-WebSocket-Key` request header with `matic string`
+* 2. get sha1 hash of the string
+* 3. get base64 of the sha1 hash
+*/
+module.exports.genericWsSecAccept = function (wsSecKey) {
+  // the string to generate the Sec-WebSocket-Accept
+  const magicString = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';
+  const targetString = `${wsSecKey}${magicString}`;
+  const shasum = crypto.createHash('sha1');
+  shasum.update(targetString);
+  return shasum.digest('base64');
+}

lib/webInterface.js

@@ -128,6 +128,7 @@ class webInterface extends events.EventEmitter {
               res.json({
                 id: query.id,
                 type: result.type,
+                method: result.meethod,
                 fileName: result.fileName,
                 ref: `/downloadBody?id=${query.id}&download=${isDownload}&raw=${!isDownload}`
               });
@@ -143,6 +144,7 @@ class webInterface extends events.EventEmitter {
               res.json({
                 id: query.id,
                 type: result.type,
+                method: result.method,
                 resBody: result.content
               });
             };
@@ -188,6 +190,22 @@ class webInterface extends events.EventEmitter {
         }
       });
 
+      app.get('/fetchWsMessages', (req, res) => {
+        const query = req.query;
+        if (query && query.id) {
+          recorder.getDecodedWsMessage(query.id, (err, messages) => {
+            if (err) {
+              console.error(err);
+              res.json([]);
+              return;
+            }
+            res.json(messages);
+          });
+        } else {
+          res.json([]);
+        }
+      });
+
       app.get('/fetchCrtFile', (req, res) => {
         res.setHeader('Access-Control-Allow-Origin', '*');
         const _crtFilePath = certMgr.getRootCAFilePath();

lib/wsServer.js

@@ -99,7 +99,11 @@ class wsServer {
 
       wss.broadcast = function (data) {
         if (typeof data === 'object') {
-          data = JSON.stringify(data);
+          try {
+            data = JSON.stringify(data);
+          } catch (e) {
+            console.error('==> errorr when do broadcast ', e, data);
+          }
         }
         for (const client of wss.clients) {
           try {
@@ -137,6 +141,20 @@ class wsServer {
         }
       });
 
+      recorder.on('updateLatestWsMsg', (data) => {
+        try {
+          // console.info('==> update latestMsg ', data);
+          wss && wss.broadcast({
+            type: 'updateLatestWsMsg',
+            content: data
+          });
+        } catch (e) {
+          logUtil.error(e.message);
+          logUtil.error(e.stack);
+          console.error(e);
+        }
+      });
+
       self.wss = wss;
     });
   }

lib/wsServerMgr.js

@@ -0,0 +1,39 @@
+/**
+* manage the websocket server
+*
+*/
+const ws = require('ws');
+const logUtil = require('./log.js');
+
+const WsServer = ws.Server;
+
+/**
+* get a new websocket server based on the server
+* @param @required {object} config
+                   {string} config.server
+                   {handler} config.handler
+*/
+function getWsServer(config) {
+  const wss = new WsServer({
+    server: config.server
+  });
+
+  wss.on('connection', config.connHandler);
+
+  wss.on('headers', (headers) => {
+    headers.push('x-anyproxy-websocket:true');
+  });
+
+  wss.on('error', e => {
+    logUtil.error(`error in websocket proxy: ${e.message},\r\n ${e.stack}`);
+    console.error('error happened in proxy websocket:', e)
+  });
+
+  wss.on('close', e => {
+    console.error('==> closing the ws server');
+  });
+
+  return wss;
+}
+
+module.exports.getWsServer = getWsServer;