refactor: refine code

avwo · Dec 2, 2020 · 6a98a32 · 6a98a32
1 parent 184b91f
commit 6a98a32
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 15 deletions.
diff --git a/lib/config.js b/lib/config.js
@@ -44,8 +44,6 @@ 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';

diff --git a/lib/tunnel.js b/lib/tunnel.js
@@ -62,7 +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 enableTunnelAck = useTunnelPolicy && req.headers['x-whistle-request-tunnel-ack'];
     var isLocalUIUrl = !useTunnelPolicy && config.isLocalUIUrl(hostname);
     if (isLocalUIUrl ? isIPHost : util.isLocalHost(hostname)) {
       isLocalUIUrl = options.port == config.port || options.port == config.uiport;
@@ -303,7 +303,7 @@ function tunnelProxy(server, proxy) {
                       if (isProxyPort) {
                         _headers[config.WEBUI_HEAD] = 1;
                       }
-                      _headers[config.REQUEST_TUNNEL_ACK] = 1;
+                      _headers['x-whistle-request-tunnel-ack'] = 1;
                     }
                     var netMgr = isSocks ? socks : config;
                     var reqDelay = util.getMatcherValue(_rules.reqDelay);
@@ -484,7 +484,7 @@ function tunnelProxy(server, proxy) {
 
         function sendEstablished(code, cb) {
           if (res) {
-            code = res.statusCode;
+            code = res.statusCode || 200;
             if (!res.headers['proxy-agent']) {
               res.headers['proxy-agent'] = config.name;
               res.rawHeaders = res.rawHeaders || [];
@@ -500,8 +500,9 @@ function tunnelProxy(server, proxy) {
               rawHeaders: ['proxy-agent', 'Proxy-Agent']
             };
           }
-          if (enableTunnelAck && cb && code == 200) {
-            res.headers[config.ALLOW_TUNNEL_ACK] = 1;
+          var tunnelAck = enableTunnelAck && cb && code == 200;
+          if (tunnelAck) {
+            res.headers['x-whistle-allow-tunnel-ack'] = 1;
           }
           var resHeaders = res.headers;
           pluginMgr.getResRules(req, res, function() {
@@ -528,10 +529,10 @@ function tunnelProxy(server, proxy) {
                 var message = code == 200 ? 'Connection Established' : (STATUS_CODES[code] || 'unknown');
                 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) {
+                if (code != 200) {
                   reqSocket.end(rawData, cb);
                 } else {
-                  if (enableTunnelAck && cb) {
+                  if (tunnelAck) {
                     reqSocket.write(rawData);
                     reqSocket.once('data', function(chunk) {
                       buf = chunk.length > 1 ? chunk.slice(1) : null;

diff --git a/lib/util/patch.js b/lib/util/patch.js
@@ -78,18 +78,19 @@ http.request = function(options) {
         }
         headers = hparser.getRawHeaders(headers);
       }
-      socket.write(['CONNECT ' + tunnelPath + ' HTTP/1.1', headers, '\r\n'].join('\r\n'));
+      var rawData = ['CONNECT ' + tunnelPath + ' HTTP/1.1', headers, '\r\n'].join('\r\n');
+      if (res.statusCode === 200 && options.headers &&
+        options.headers['x-whistle-request-tunnel-ack'] &&
+        res.headers['x-whistle-allow-tunnel-ack']) {
+          rawData = '1' + rawData;
+      }
+      socket.write(rawData);
       hparser.parse(socket, function(err, _res) {
         if (err) {
           return client.emit('error', err);
         }
         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);