From 855fb42e3e352de1b369996cfef0a68af1cd9b1b Mon Sep 17 00:00:00 2001
From: evlon <liamgnc@gmail.com>
Date: Fri, 9 Jul 2021 13:36:39 +0800
Subject: [PATCH 1/2] Update requestHandler.js

enable rule.beforeDealHttpsRequest to overwrite host and port
---
 lib/requestHandler.js | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/lib/requestHandler.js b/lib/requestHandler.js
index 6a6c6c034..ce5426f38 100644
--- a/lib/requestHandler.js
+++ b/lib/requestHandler.js
@@ -526,7 +526,7 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
   const reqHandlerCtx = this; reqHandlerCtx.conns = new Map(); reqHandlerCtx.cltSockets = new Map()
 
   return function (req, cltSocket, head) {
-    const host = req.url.split(':')[0],
+    let host = req.url.split(':')[0],
       targetPort = req.url.split(':')[1];
     let shouldIntercept;
     let interceptWsRequest = false;
@@ -552,7 +552,12 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
       // the return value in default rule is null
       // so if the value is null, will take it as final value
       shouldIntercept = yield userRule.beforeDealHttpsRequest(requestDetail);
-
+      
+      //enable rule.beforeDealHttpsRequest to overwrite host and port  @evlon 2021-7-9
+      if(requestDetail.host != req.url){
+        host = requestDetail.host.split(':')[0],
+        targetPort = requestDetail.host.split(':')[1];
+      }
       // otherwise, will take the passed in option
       if (shouldIntercept === null) {
         shouldIntercept = reqHandlerCtx.forceProxyHttps;

From 7e032e6ed1b4c701038d743a752a1d34311f6773 Mon Sep 17 00:00:00 2001
From: evlon <liamgnc@gmail.com>
Date: Sat, 10 Jul 2021 08:30:42 +0800
Subject: [PATCH 2/2] Update requestHandler.js
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

添加websocket hook
---
 lib/requestHandler.js | 50 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 6 deletions(-)

diff --git a/lib/requestHandler.js b/lib/requestHandler.js
index ce5426f38..271c171d4 100644
--- a/lib/requestHandler.js
+++ b/lib/requestHandler.js
@@ -218,6 +218,8 @@ function getWsReqInfo(wsReq) {
   const hostName = host.split(':')[0];
   const port = host.split(':')[1];
 
+   // https://github.com/alibaba/anyproxy/pull/450
+   const portPart = port ? `:${port}` : ''; 
   // TODO 如果是windows机器，url是不是全路径？需要对其过滤，取出
   const path = wsReq.url || '/';
 
@@ -243,6 +245,7 @@ function getWsReqInfo(wsReq) {
 
 
   return {
+    url: `${isEncript ? 'wss' : 'ws'}://${hostname}${portPart}${path}`,
     headers: headers, // the full headers of origin ws connection
     noWsHeaders: getNoWsHeaders(),
     hostName: hostName,
@@ -552,12 +555,11 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
       // the return value in default rule is null
       // so if the value is null, will take it as final value
       shouldIntercept = yield userRule.beforeDealHttpsRequest(requestDetail);
-      
-      //enable rule.beforeDealHttpsRequest to overwrite host and port  @evlon 2021-7-9
       if(requestDetail.host != req.url){
         host = requestDetail.host.split(':')[0],
         targetPort = requestDetail.host.split(':')[1];
       }
+
       // otherwise, will take the passed in option
       if (shouldIntercept === null) {
         shouldIntercept = reqHandlerCtx.forceProxyHttps;
@@ -802,6 +804,24 @@ function getWsHandler(userRule, recorder, wsClient, wsReq) {
       recorder && recorder.updateRecordWsMessage(resourceInfoId, message);
     };
 
+      /**
+     * prepare messageDetail object for intercept hooks
+     * @param {object} messageEvent
+     * @returns {object}
+     */
+    const prepareMessageDetail = (messageEvent) => {
+      return {
+        requestOptions: {
+          port: serverInfo.port,
+          hostname: serverInfo.hostname,
+          path: serverInfo.path,
+          secure: serverInfo.secure,
+        },
+        url: serverInfo.url,
+        data: messageEvent.data,
+      };
+    };
+
     proxyWs.onopen = () => {
       consumeMsgQueue();
     }
@@ -830,8 +850,18 @@ function getWsHandler(userRule, recorder, wsClient, wsReq) {
     }
 
     proxyWs.onmessage = (event) => {
-      recordMessage(event, false);
-      wsClient.readyState === 1 && wsClient.send(event.data);
+
+      // recordMessage(event, false);
+      // wsClient.readyState === 1 && wsClient.send(event.data);
+      co(function *() {
+        const modifiedMsg = (yield userRule.beforeSendWsMessageToClient(prepareMessageDetail(event))) || {};
+        const finalMsg = {
+          data: modifiedMsg.data || event.data,
+        };
+        recordMessage(finalMsg, false);
+        wsClient.readyState === 1 && wsClient.send(finalMsg.data);
+      });
+      
     }
 
     proxyWs.onclose = (event) => {
@@ -841,8 +871,16 @@ function getWsHandler(userRule, recorder, wsClient, wsReq) {
     }
 
     wsClient.onmessage = (event) => {
-      recordMessage(event, true);
-      sendProxyMessage(event);
+      // recordMessage(event, true);
+      // sendProxyMessage(event);
+      co(function *() {
+        const modifiedMsg = (yield userRule.beforeSendWsMessageToServer(prepareMessageDetail(event))) || {};
+        const finalMsg = {
+          data: modifiedMsg.data || event.data,
+        };
+        recordMessage(finalMsg, true);
+        sendProxyMessage(finalMsg);
+      });
     }
 
     wsClient.onclose = (event) => {