Merge pull request #1 from evlon/feat-ws-hook

Update requestHandler.js
evlon · Jul 10, 2021 · 44a6a30 · 44a6a30
2 parents 855fb42 + 7e032e6
commit 44a6a30
Showing 1 changed file with 44 additions and 6 deletions.
diff --git 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) => {