From 184b91f62ce73c61aed706fbd258552ac4fcdce6 Mon Sep 17 00:00:00 2001
From: avwo <avwu@qq.com>
Date: Wed, 2 Dec 2020 14:18:42 +0800
Subject: [PATCH] feat: add tunnel ack

---
 lib/config.js     |  5 +++++
 lib/tunnel.js     | 22 ++++++++++++++++++----
 lib/util/patch.js |  5 +++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/lib/config.js b/lib/config.js
index eea367766..3c91162c7 100644
--- a/lib/config.js
+++ b/lib/config.js
@@ -44,6 +44,8 @@ config.HTTPS_FIELD = 'x-whistle-https-request';
 config.HTTPS_PROTO_HEADER = 'x-forwarded-proto';
 config.INTERNAL_ID = Date.now() + '/' + process.pid + '/' + Math.floor(Math.random() * 100000);
 config.INTERNAL_ID_HEADER = '_x-whistle-internal-id';
+config.REQUEST_TUNNEL_ACK = 'x-whistle-request-tunnel-ack';
+config.ALLOW_TUNNEL_ACK = 'x-whistle-allow-tunnel-ack';
 config.DATA_ID = 'x-whistle-data-id' + '-' + uid;
 config.PROXY_ID_HEADER = 'x-whistle-proxy-id-' + uid;
 config.CLIENT_PORT_HEAD = 'x-whistle-client-port';
@@ -412,6 +414,9 @@ function connect(options, cb) {
       });
       return;
     }
+    if (headers['x-whistle-request-tunnel-ack'] && res.headers['x-whistle-allow-tunnel-ack']) {
+      socket.write('1');
+    }
     cb(socket, res);
   }).end();
   return req;
diff --git a/lib/tunnel.js b/lib/tunnel.js
index b42676b24..8e66c6f5a 100644
--- a/lib/tunnel.js
+++ b/lib/tunnel.js
@@ -62,6 +62,7 @@ function tunnelProxy(server, proxy) {
     var isIPHost = !isICloundCKDB && net.isIP(hostname);
     var policy = headers[config.WHISTLE_POLICY_HEADER];
     var useTunnelPolicy = policy == 'tunnel';
+    var enableTunnelAck = useTunnelPolicy && req.headers[config.REQUEST_TUNNEL_ACK];
     var isLocalUIUrl = !useTunnelPolicy && config.isLocalUIUrl(hostname);
     if (isLocalUIUrl ? isIPHost : util.isLocalHost(hostname)) {
       isLocalUIUrl = options.port == config.port || options.port == config.uiport;
@@ -302,6 +303,7 @@ function tunnelProxy(server, proxy) {
                       if (isProxyPort) {
                         _headers[config.WEBUI_HEAD] = 1;
                       }
+                      _headers[config.REQUEST_TUNNEL_ACK] = 1;
                     }
                     var netMgr = isSocks ? socks : config;
                     var reqDelay = util.getMatcherValue(_rules.reqDelay);
@@ -443,9 +445,7 @@ function tunnelProxy(server, proxy) {
             };
             var handleEstablished = function() {
               if (useTunnelPolicy) {
-                sendEstablished(200, function() {
-                  setTimeout(connHandler, 16);
-                });
+                sendEstablished(200, connHandler);
               } else {
                 connHandler();
                 sendEstablished();
@@ -500,6 +500,9 @@ function tunnelProxy(server, proxy) {
               rawHeaders: ['proxy-agent', 'Proxy-Agent']
             };
           }
+          if (enableTunnelAck && cb && code == 200) {
+            res.headers[config.ALLOW_TUNNEL_ACK] = 1;
+          }
           var resHeaders = res.headers;
           pluginMgr.getResRules(req, res, function() {
             var reqRules = req.rules;
@@ -528,7 +531,18 @@ function tunnelProxy(server, proxy) {
                 if (code && code != 200) {
                   reqSocket.end(rawData, cb);
                 } else {
-                  reqSocket.write(rawData, cb);
+                  if (enableTunnelAck && cb) {
+                    reqSocket.write(rawData);
+                    reqSocket.once('data', function(chunk) {
+                      buf = chunk.length > 1 ? chunk.slice(1) : null;
+                      reqSocket.pause();
+                      cb();
+                    });
+                  } else {
+                    reqSocket.write(rawData, cb && function() {
+                      setTimeout(cb, 16);
+                    });
+                  }
                 }
               }
               if (reqEmitter) {
diff --git a/lib/util/patch.js b/lib/util/patch.js
index 0e1774078..921707e1f 100644
--- a/lib/util/patch.js
+++ b/lib/util/patch.js
@@ -85,6 +85,11 @@ http.request = function(options) {
         }
         res.statusCode = parseInt(_res.statusCode, 10);
         res.headers = _res.headers;
+        if (res.statusCode === 200 && options.headers &&
+          options.headers['x-whistle-request-tunnel-ack'] &&
+          res.headers['x-whistle-allow-tunnel-ack']) {
+          socket.write('1');
+        }
         !options.keepStreamResume && socket.pause();
         listener.call(this, res, socket, head);
       }, true);