From 02824de4e3393948e709c6b34132db579efbf23b Mon Sep 17 00:00:00 2001
From: avwo <avwu@qq.com>
Date: Mon, 29 Oct 2018 15:31:20 +0800
Subject: [PATCH] feat: add enable://customInspect

---
 lib/https/index.js |  2 +-
 lib/socket-mgr.js  | 52 ++++++++++++++++++++++++++--------------------
 lib/tunnel.js      |  9 ++++----
 lib/util/index.js  | 12 ++++++++---
 4 files changed, 43 insertions(+), 32 deletions(-)

diff --git a/lib/https/index.js b/lib/https/index.js
index 6ec0cf522..255a382d2 100644
--- a/lib/https/index.js
+++ b/lib/https/index.js
@@ -455,7 +455,7 @@ function resolveWebsocket(socket, callback, wss) {
               util.disableResProps(socket, headers);
               socket.write(res.getBuffer(headers));
               reqSocket.reqId = data.id;
-              socket.customInspect = util.getInspectStatus(socket.enable);
+              socket.customInspect = util.getInspectStatus(socket);
               socketMgr.handleWebsocket(socket, reqSocket);
               resData.headers = headers;
               if (res.statusCode != 101) {
diff --git a/lib/socket-mgr.js b/lib/socket-mgr.js
index f26c7a51c..69c3a2836 100644
--- a/lib/socket-mgr.js
+++ b/lib/socket-mgr.js
@@ -102,28 +102,36 @@ function initStatus(ctx, enable) {
 exports.handleConnect = function(req, res) {
   var hasEvent = util.listenerCount(proxy, 'tunnelRequest');
   var isConn = req.inspectFrames;
-  var customInspect = !req.pluginServerPort && req.customInspect;
-  if (hasEvent || customInspect || isConn) {
+  
+  if (hasEvent || isConn) {
     var url = req.fullUrl;
     var reqId = req.reqId;
+    var customInspect = req.customInspect;
     var charset = util.getCharset(req.headers && req.headers['content-type']) || '';
     var sendStatus = {};
     var receiveStatus = {};
-    
-    var emitData = function(chunk, fromClient, ignore) {
-      if (req.filter.hide) {
-        return;
-      }
-      proxy.emit('frame', {
-        reqId: reqId,
-        frameId: getFrameId(),
-        isClient: fromClient,
-        charset: charset,
-        length: chunk.length,
-        ignore: ignore,
-        bin: chunk
-      });
-    };
+    var emitData, handleInspect;
+
+    if (customInspect || !isConn) {
+      emitData = function(chunk, fromClient, ignore) {
+        if (req.filter.hide) {
+          return;
+        }
+        proxy.emit('frame', {
+          reqId: reqId,
+          frameId: getFrameId(),
+          isClient: fromClient,
+          charset: charset,
+          length: chunk.length,
+          ignore: ignore,
+          bin: chunk
+        });
+      };
+      handleInspect = (customInspect || !isConn) && function(chunk, encoding, cb) {
+        hasEvent && proxy.emit('tunnelRequest', url);
+        cb(null, chunk);
+      };
+    }
     if (isConn) {
       var ctx = conns[reqId] = {
         customInspect: customInspect,
@@ -154,10 +162,6 @@ exports.handleConnect = function(req, res) {
       };
       initStatus(ctx, req.enable);
     }
-    var handleInspect = (customInspect || !isConn) && function(chunk, encoding, cb) {
-      hasEvent && proxy.emit('tunnelRequest', url);
-      cb(null, chunk);
-    };
     var reqTrans = new PassThrough();
     reqTrans.on('error', function(err) {
       req.emit('error', err);
@@ -492,8 +496,10 @@ exports.sendData = function(data) {
     if (dataList.length > 3) {
       return false;
     }
-    data.data = data.data.toString('base64');
-    return dataList.push(data);
+    data.base64 = data.data.toString('base64');
+    delete data.data;
+    dataList.push(data);
+    return;
   }
   if (isServer) {
     data.mask = true;
diff --git a/lib/tunnel.js b/lib/tunnel.js
index bf2e7a983..bdaaa7db4 100644
--- a/lib/tunnel.js
+++ b/lib/tunnel.js
@@ -38,7 +38,7 @@ function tunnelProxy(server, proxy) {
     tunnelUrl = req.fullUrl = 'tunnel://' + options.host;
     proxy.emit('_request', tunnelUrl);
     var resSocket, responsed, reqEmitter, data, rulesMgr, originPort;
-    var headers, inspect, customInspect, reqData, resData, res, rollBackTunnel, buf;
+    var headers, inspect, reqData, resData, res, rollBackTunnel, buf;
     req.isTunnel = true;
     req.method = util.toUpperCase(req.method) || 'CONNECT';
     reqSocket.clientIp = req.clientIp = util.getClientIp(req);
@@ -123,8 +123,7 @@ function tunnelProxy(server, proxy) {
         function handleTunnel() {
           var reqRawHeaderNames = getRawHeaderNames(req.rawHeaders) || {};
           var enable = req.enable;
-          inspect = enable.inspect;
-          customInspect = util.getInspectStatus(req.enable);
+          inspect = enable.inspect || util.isCustomInspect(req);
           headers = req.headers;
           req.clientPort = reqSocket.remotePort;
           reqData = {
@@ -148,7 +147,7 @@ function tunnelProxy(server, proxy) {
             req: reqData,
             res: resData,
             isHttps: true,
-            inspect: inspect || !!customInspect,
+            inspect: inspect,
             rulesHeaders: req.rulesHeaders
           };
           if (!filter.hide) {
@@ -332,7 +331,7 @@ function tunnelProxy(server, proxy) {
           reqSocket.headers = headers;
           reqSocket.fullUrl = tunnelUrl;
           reqSocket.inspectFrames = useTunnelPolicy || inspect;
-          reqSocket.customInspect = customInspect;
+          reqSocket.customInspect = util.getInspectStatus(req);
           if (buf) {
             resSocket.write(buf);
             buf = null;
diff --git a/lib/util/index.js b/lib/util/index.js
index 98a5c5259..81a179adf 100644
--- a/lib/util/index.js
+++ b/lib/util/index.js
@@ -2130,11 +2130,17 @@ exports.getMatcherUrl = function(rule) {
   }
 };
 
-exports.getInspectStatus = function(enable) {
-  if (!enable.customInspect && !enable.customParser 
-      && !enable.customFrame) {
+function isCustomInspect(req) {
+  var enable = req.enable || '';
+  return enable.customInspect || enable.customParser || enable.customFrame;
+}
+exports.isCustomInspect = isCustomInspect;
+
+exports.getInspectStatus = function(req) {
+  if (!req.pluginServerPort || !isCustomInspect(req)) {
     return;
   }
+  var enable = req.enable;
   var customInspect = ['inspect'];
   if (enable.pauseSend) {
     customInspect.push('pauseSend');