Update requestHandler.js

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,
@@ -526,7 +529,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,6 +555,10 @@ 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);
+      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) {
@@ -797,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();
     }
@@ -825,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) => {
@@ -836,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) => {