diff --git a/anyproxy/lib/requestHandler.js b/anyproxy/lib/requestHandler.js index f226266..cd93703 100644 --- a/anyproxy/lib/requestHandler.js +++ b/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, { diff --git a/anyproxy/lib/rule_default.js b/anyproxy/lib/rule_default.js index b3d29e9..56b25e8 100644 --- a/anyproxy/lib/rule_default.js +++ b/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; + }, }; diff --git a/anyproxy/lib/wsServerMgr.js b/anyproxy/lib/wsServerMgr.js index 62e5a50..a34e873 100644 --- a/anyproxy/lib/wsServerMgr.js +++ b/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'); diff --git a/extension/src/bg/background.js b/extension/src/bg/background.js index bbe96f8..075ca3d 100755 --- a/extension/src/bg/background.js +++ b/extension/src/bg/background.js @@ -19,6 +19,7 @@ var redirect_table = {}; const REQUEST_HEADER_BLACKLIST = [ 'cookie' ]; +const WEBSOCKETS = {}; const RPC_CALL_TABLE = { 'HTTP_REQUEST': perform_http_request, @@ -26,8 +27,98 @@ const RPC_CALL_TABLE = { '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(); -})(); \ No newline at end of file +})(); diff --git a/server.js b/server.js index 0e6cb41..c4d20d2 100644 --- a/server.js +++ b/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...`);