refactor: refine tunnel proxy

avwo · Dec 1, 2020 · 1ddbc07 · 1ddbc07
1 parent 18c4f88
commit 1ddbc07
Showing 1 changed file with 18 additions and 6 deletions.
diff --git a/lib/tunnel.js b/lib/tunnel.js
@@ -438,12 +438,23 @@ function tunnelProxy(server, proxy) {
               reqSocket.inspectFrames = data.inspect = true;
               reqSocket.customParser = false;
             }
-            socketMgr.handleConnect(reqSocket, resSocket);
+            var handleEstablished = function() {
+              if (useTunnelPolicy) {
+                sendEstablished(200, function() {
+                  setTimeout(function() {
+                    socketMgr.handleConnect(reqSocket, resSocket);
+                  }, 16);
+                });
+              } else {
+                socketMgr.handleConnect(reqSocket, resSocket);
+                sendEstablished();
+              }
+            };
             var resDelay = util.getMatcherValue(_rules.resDelay);
             if (resDelay > 0) {
-              setTimeout(sendEstablished, resDelay);
+              setTimeout(handleEstablished, resDelay);
             } else {
-              sendEstablished();
+              handleEstablished();
             }
           });
         }
@@ -470,7 +481,7 @@ function tunnelProxy(server, proxy) {
           util.onSocketEnd(socket, emitError);
         }
 
-        function sendEstablished(code) {
+        function sendEstablished(code, cb) {
           if (res) {
             code = res.statusCode;
             if (!res.headers['proxy-agent']) {
@@ -494,6 +505,7 @@ function tunnelProxy(server, proxy) {
             util.parseRuleJson(rollBackTunnel ? null : reqRules.resHeaders, function(newResHeaders) {
               if (rollBackTunnel) {
                 reqSocket.resume();
+                cb && cb();
               } else {
                 var rawHeaderNames = getRawHeaderNames(res.rawHeaders) || {};
                 if (newResHeaders) {
@@ -513,9 +525,9 @@ function tunnelProxy(server, proxy) {
                 var statusLine = ['HTTP/1.1', code, message].join(' ');
                 var rawData = [statusLine, getRawHeaders(formatHeaders(resHeaders, rawHeaderNames))].join('\r\n') + '\r\n\r\n';
                 if (code && code != 200) {
-                  reqSocket.end(rawData);
+                  reqSocket.end(rawData, cb);
                 } else {
-                  reqSocket.write(rawData);
+                  reqSocket.write(rawData, cb);
                 }
               }
               if (reqEmitter) {