From 720fb00f96b6b645309ef16bb01d252e85b26cd4 Mon Sep 17 00:00:00 2001
From: avwo <avwu@qq.com>
Date: Wed, 6 Jun 2018 11:16:26 +0800
Subject: [PATCH] feat: handle the tunnel request of proxifier

---
 lib/https/index.js | 16 +++++++++++++---
 lib/tunnel.js      |  3 +--
 lib/util/index.js  |  5 +++++
 3 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/lib/https/index.js b/lib/https/index.js
index 24ce2df1e..45d0cec2c 100644
--- a/lib/https/index.js
+++ b/lib/https/index.js
@@ -547,12 +547,14 @@ function setRuleHeaders(headers, rulesHeaders) {
   }
 }
 
-var HTTP_RE = /^\w+\s+[^\s]+\s+HTTP\/1.1(?:\r?\n|\r)/;
+var CONNECT_RE = /^CONNECT\s+([^\s]+)\s+HTTP\/1.\d$/mi;
+var HTTP_RE = /^\w+\s+[^\s]+\s+HTTP\/1.\d$/mi;
 var CRLF = util.toBuffer('\r\n');
 
 module.exports = function(socket, hostname, callback, next) {
   var reqSocket, destroyed, timer;
   var rulesHeaders = socket.rulesHeaders;
+  var isConnect, headersStr;
   function destroy(err) {
     if (destroyed) {
       return;
@@ -569,15 +571,23 @@ module.exports = function(socket, hostname, callback, next) {
   socket.on('data', request);
   socket.on('end', request);
   var clientIp = socket.clientIp;
-  timer = setTimeout(request, util.isLocalAddress(clientIp) ? LOCAL_TIMEOUT : TIMEOUT);
+  var timeout = util.isLocalAddress(clientIp) ? LOCAL_TIMEOUT : TIMEOUT;
+  timer = setTimeout(request, timeout);
   function request(chunk) {
     clearTimeout(timer);
+    headersStr = chunk && chunk.toString();
+    if (chunk && !isConnect && CONNECT_RE.test(headersStr)) {
+      isConnect = true;
+      hostname = RegExp.$1.split(':')[0];
+      timer = setTimeout(request, timeout);
+      util.setEstablished(socket);
+      return;
+    }
     socket.removeListener('data', request);
     socket.removeListener('end', request);
     if (!chunk) {//没有数据
       return next(chunk);
     }
-    var headersStr = chunk.toString();
     if (/upgrade\s*:\s*websocket/i.test(headersStr)) { //ws
       parseReq(socket, function(err, socket) {
         if (err) {
diff --git a/lib/tunnel.js b/lib/tunnel.js
index 5a920acc7..4ac4e1e1c 100644
--- a/lib/tunnel.js
+++ b/lib/tunnel.js
@@ -19,7 +19,6 @@ var STATUS_CODES = require('http').STATUS_CODES || {};
 var getRawHeaderNames = hparser.getRawHeaderNames;
 var formatHeaders = hparser.formatHeaders;
 var getRawHeaders = hparser.getRawHeaders;
-var ESTABLISHED_CTN = 'HTTP/1.1 200 Connection Established\r\nProxy-Agent: ' + config.name + '\r\n\r\n';
 
 function tunnelProxy(server, proxy) {
   server.on('connect', function(req, reqSocket, head) {//ws, wss, https proxy
@@ -104,7 +103,7 @@ function tunnelProxy(server, proxy) {
               }
             });
           });
-          reqSocket.write(ESTABLISHED_CTN);
+          util.setEstablished(reqSocket);
           return;
         }
         
diff --git a/lib/util/index.js b/lib/util/index.js
index e3b05d50f..1810dae57 100644
--- a/lib/util/index.js
+++ b/lib/util/index.js
@@ -46,6 +46,11 @@ function noop() {}
 
 exports.noop = noop;
 
+var ESTABLISHED_CTN = 'HTTP/1.1 200 Connection Established\r\nProxy-Agent: ' + config.name + '\r\n\r\n';
+exports.setEstablished = function(socket) {
+  socket.write(ESTABLISHED_CTN);
+};
+
 function listenerCount(emitter, eventName) {
   if (typeof emitter.listenerCount === 'function') {
     return emitter.listenerCount(eventName);