anyproxy/lib/handlers/wsHandler.js

'use strict';

const co = require('co');
const WebSocket = require('ws');
const logUtil = require('../log');

/**
 * construct the request headers based on original connection,
 * but delete the `sec-websocket-*` headers as they are already consumed by AnyProxy
 */
function getNoWsHeaders(headers) {
  const originHeaders = Object.assign({}, headers);

  Object.keys(originHeaders).forEach((key) => {
    // if the key matchs 'sec-websocket', delete it
    if (/sec-websocket/ig.test(key)) {
      delete originHeaders[key];
    }
  });

  delete originHeaders.connection;
  delete originHeaders.upgrade;
  return originHeaders;
}

/**
 * get request info from the ws client
 * @param @required wsClient the ws client of WebSocket
 */
function getWsReqInfo(wsReq) {
  const headers = wsReq.headers || {};
  const host = headers.host;
  const hostname = host.split(':')[0];
  const port = host.split(':')[1];
  // TODO 如果是windows机器，url是不是全路径？需要对其过滤，取出
  const path = wsReq.url || '/';
  const isEncript = wsReq.connection && wsReq.connection.encrypted;

  return {
    url: `${isEncript ? 'wss' : 'ws'}://${hostname}:${port}${path}`,
    headers: headers, // the full headers of origin ws connection
    noWsHeaders: getNoWsHeaders(headers),
    secure: Boolean(isEncript),
    hostname: hostname,
    port: port,
    path: path
  };
}

/**
 * When the source ws is closed, we need to close the target websocket.
 * If the source ws is normally closed, that is, the code is reserved, we need to transfrom them
 * @param {object} event CloseEvent
 */
const getCloseFromOriginEvent = (closeEvent) => {
  const code = closeEvent.code || '';
  const reason = closeEvent.reason || '';
  let targetCode = '';
  let targetReason = '';
  if (code >= 1004 && code <= 1006) {
    targetCode = 1000; // normal closure
    targetReason = `Normally closed. The origin ws is closed at code: ${code} and reason: ${reason}`;
  } else {
    targetCode = code;
    targetReason = reason;
  }

  return {
    code: targetCode,
    reason: targetReason
  };
}

/**
 * get a websocket event handler
 * @param @required {object} wsClient
 */
function handleWs(userRule, recorder, wsClient, wsReq) {
  const self = this;
  let resourceInfoId = -1;
  const resourceInfo = {
    wsMessages: [] // all ws messages go through AnyProxy
  };
  const clientMsgQueue = [];
  const serverInfo = getWsReqInfo(wsReq);
  // proxy-layer websocket client
  const proxyWs = new WebSocket(serverInfo.url, '', {
    rejectUnauthorized: !self.dangerouslyIgnoreUnauthorized,
    headers: serverInfo.noWsHeaders
  });

  if (recorder) {
    Object.assign(resourceInfo, {
      host: serverInfo.hostname,
      method: 'WebSocket',
      path: serverInfo.path,
      url: serverInfo.url,
      req: wsReq,
      startTime: new Date().getTime()
    });
    resourceInfoId = recorder.appendRecord(resourceInfo);
  }

  /**
   * store the messages before the proxy ws is ready
   */
  const sendProxyMessage = (finalMsg) => {
    const message = finalMsg.data;
    if (proxyWs.readyState === 1) {
      // if there still are msg queue consuming, keep it going
      if (clientMsgQueue.length > 0) {
        clientMsgQueue.push(message);
      } else {
        proxyWs.send(message);
      }
    } else {
      clientMsgQueue.push(message);
    }
  };

  /**
   * consume the message in queue when the proxy ws is not ready yet
   * will handle them from the first one-by-one
   */
  const consumeMsgQueue = () => {
    while (clientMsgQueue.length > 0) {
      const message = clientMsgQueue.shift();
      proxyWs.send(message);
    }
  };

  /**
   * consruct a message Record from message event
   * @param @required {object} finalMsg based on the MessageEvent from websockt.onmessage
   * @param @required {boolean} isToServer whether the message is to or from server
   */
  const recordMessage = (finalMsg, isToServer) => {
    const message = {
      time: Date.now(),
      message: finalMsg.data,
      isToServer: isToServer
    };

    // resourceInfo.wsMessages.push(message);
    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();
  };

  // this event is fired when the connection is build and headers is returned
  proxyWs.on('upgrade', (response) => {
    resourceInfo.endTime = new Date().getTime();
    const headers = response.headers;
    resourceInfo.res = { //construct a self-defined res object
      statusCode: response.statusCode,
      headers: headers,
    };

    resourceInfo.statusCode = response.statusCode;
    resourceInfo.resHeader = headers;
    resourceInfo.resBody = '';
    resourceInfo.length = resourceInfo.resBody.length;

    recorder && recorder.updateRecord(resourceInfoId, resourceInfo);
  });

  proxyWs.onerror = (e) => {
    // https://developer.mozilla.org/en-US/docs/Web/API/CloseEvent#Status_codes
    wsClient.close(1001, e.message);
    proxyWs.close(1001);
  };

  proxyWs.onmessage = (event) => {
    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) => {
    logUtil.debug(`proxy ws closed with code: ${event.code} and reason: ${event.reason}`);
    const targetCloseInfo = getCloseFromOriginEvent(event);
    wsClient.readyState !== 3 && wsClient.close(targetCloseInfo.code, targetCloseInfo.reason);
  };

  wsClient.onmessage = (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) => {
    logUtil.debug(`original ws closed with code: ${event.code} and reason: ${event.reason}`);
    const targetCloseInfo = getCloseFromOriginEvent(event);
    proxyWs.readyState !== 3 && proxyWs.close(targetCloseInfo.code, targetCloseInfo.reason);
  };
}

module.exports = function getWsHandler(userRule, recorder, wsClient, wsReq) {
  try {
    handleWs.call(this, userRule, recorder, wsClient, wsReq);
  } catch (e) {
    logUtil.debug('WebSocket Proxy Error:' + e.message);
    logUtil.debug(e.stack);
    console.error(e);
  }
}