Merge pull request #1 from gemini-testing/dd.ws_interceptor

feat: add ability to intercept websocket messages

Author: DudaGod

.travis.yml

@@ -1,16 +1,5 @@
-sudo: false
 language: node_js
 node_js:
-  - "4"
-  - "6"
   - "8"
 script:
   - npm test
-after_success:
-  - bash <(curl -s https://codecov.io/bash)
-matrix:
-  fast_finish: true
-notifications:
-  email:
-    - [email protected]
-  irc: "irc.freenode.org#nodejitsu"

README.md

@@ -395,6 +395,26 @@ proxyServer.listen(8015);
 
     };
     ```
+*  **wsInterceptServerMsg**: Is a handler which is called when a websocket message is intercepted on its way to the server. It takes two arguments: `data` - is a websocket message and flags (fin, mask, compress, binary). If falsy value is returned then nothing will be sended to the target server.
+    ```
+      const proxy = new HttpProxy({
+        ...
+        wsInterceptServerMsg: (data, flags) {
+          return typeof data === 'string ? data.toUpperCase() : data;
+        }
+        ...
+      })
+    ```
+*  **wsInterceptClientMsg**: Is a handler which is called when a websocket message is intercepted on its way to the client. It takes two arguments: `data` - is a websocket message and flags (fin, mask, compress, binary). If falsy value is returned then nothing will be sended to the client.
+     ```
+      const proxy = new HttpProxy({
+        ...
+        wsInterceptClientMsg: (data, flags) {
+          return typeof data === 'string ? data.toUpperCase() : data;
+        }
+        ...
+      })
+    ```
 
 **NOTE:**
 `options.ws` and `options.ssl` are optional.
@@ -414,6 +434,8 @@ If you are using the `proxyServer.listen` method, the following options are also
 * `proxyReq`: This event is emitted before the data is sent. It gives you a chance to alter the proxyReq request object. Applies to "web" connections
 * `proxyReqWs`: This event is emitted before the data is sent. It gives you a chance to alter the proxyReq request object. Applies to "websocket" connections
 * `proxyRes`: This event is emitted if the request to the target got a response.
+* `wsServerMsg`: This event is emitted after websocket message is sended to the server.
+* `wsClientMsg`: This event is emitted after webscoket mesage is sended to the client.
 * `open`: This event is emitted once the proxy websocket was created and piped into the target websocket.
 * `close`: This event is emitted once the proxy websocket was closed.
 * (DEPRECATED) `proxySocket`: Deprecated in favor of `open`.

lib/http-proxy/passes/ws-incoming.js

@@ -1,6 +1,9 @@
-var http   = require('http'),
-    https  = require('https'),
-    common = require('../common');
+'use strict';
+
+const http = require('http');
+const https = require('https');
+const common = require('../common');
+const WsInterceptor = require('../ws/interceptor');
 
 /*!
  * Array of passes.
@@ -142,7 +145,12 @@ module.exports = {
       //
       socket.write(createHttpHeader('HTTP/1.1 101 Switching Protocols', proxyRes.headers));
 
-      proxySocket.pipe(socket).pipe(proxySocket);
+      if (options.wsInterceptServerMsg || options.wsInterceptClientMsg) {
+        WsInterceptor.create({socket, options, proxyReq, proxyRes, proxySocket}).intercept();
+      }
+      else {
+        proxySocket.pipe(socket).pipe(proxySocket);
+      }
 
       server.emit('open', proxySocket);
       server.emit('proxySocket', proxySocket);  //DEPRECATED.

lib/http-proxy/ws/interceptor.js

@@ -0,0 +1,102 @@
+'use strict';
+
+const PerMessageDeflate = require('ws/lib/PerMessageDeflate');
+const Extensions = require('ws/lib/Extensions');
+const Receiver = require('ws/lib/Receiver');
+const Sender = require('ws/lib/Sender');
+
+const acceptExtensions = ({extenstions, isServer}) => {
+    const {extensionName} = PerMessageDeflate;
+    const extenstion = extenstions[extensionName];
+
+    if (!extenstion) {
+        return {};
+    }
+
+    const perMessageDeflate = new PerMessageDeflate({}, isServer);
+    perMessageDeflate.accept(extenstion);
+
+    return {[extensionName]: perMessageDeflate};
+};
+
+const getMsgHandler = ({interceptor, dataSender, binary}) => {
+    return (data, flags) => {
+        if (typeof interceptor !== 'function') {
+            dataSender({data});
+        }
+
+        const modifiedData = interceptor(data, flags);
+
+        // if interceptor does not return data then nothing will be sended to the server
+        if (modifiedData) {
+            dataSender({data: modifiedData, binary});
+        }
+    }
+};
+
+module.exports = class Interceptor {
+    static create(opts = {}) {
+        return new this(opts);
+    }
+
+    constructor({socket, options, proxyReq, proxyRes, proxySocket}) {
+        this._socket = socket;
+        this._options = options;
+        this._proxyReq = proxyReq;
+        this._proxyRes = proxyRes;
+        this._proxySocket = proxySocket;
+
+        this._configure();
+    }
+
+    _configure() {
+        const secWsExtensions = this._proxyRes.headers['sec-websocket-extensions'];
+        const extenstions = Extensions.parse(secWsExtensions);
+        this._isCompressed = secWsExtensions && secWsExtensions.indexOf('permessage-deflate') != -1;
+
+        // need both versions of extensions for each side of the proxy connection
+        this._clientExtenstions = this._isCompressed ? acceptExtensions({extenstions, isServer: false}) : null;
+        this._serverExtenstions = this._isCompressed ? acceptExtensions({extenstions, isServer: true}) : null;
+    }
+
+    _getDataSender({sender, event, options}) {
+        return ({data, binary = false}) => {
+            const opts = Object.assign({fin: true, compress: this._isCompressed, binary}, options);
+            sender.send(data, opts);
+
+            this._proxyReq.emit(event, {data, binary});
+        };
+    }
+
+    _interceptServerMessages() {
+        const receiver = new Receiver(this._clientExtenstions);
+        const sender = new Sender(this._proxySocket, this._serverExtenstions);
+
+        // frame must be masked when send from client to server - https://tools.ietf.org/html/rfc6455#section-5.3
+        const options = {mask: true};
+        const dataSender = this._getDataSender({sender, event: 'wsServerMsg', options});
+
+        receiver.ontext = getMsgHandler({interceptor: this._options.wsInterceptServerMsg, dataSender, binary: false});
+        receiver.onbinary = getMsgHandler({interceptor: this._options.wsInterceptServerMsg, dataSender, binary: true});
+
+        this._socket.on('data', (data) => receiver.add(data));
+    }
+
+    _interceptClientMessages() {
+        const receiver = new Receiver(this._serverExtenstions);
+        const sender = new Sender(this._socket, this._clientExtenstions);
+
+        const options = {mask: false};
+        const dataSender = this._getDataSender({sender, event: 'wsClientMsg', options});
+
+        receiver.ontext = getMsgHandler({interceptor: this._options.wsInterceptClientMsg, dataSender, binary: false});
+        receiver.onbinary = getMsgHandler({interceptor: this._options.wsInterceptClientMsg, dataSender, binary: true});
+
+        this._proxySocket.on('data', (data) => receiver.add(data));
+    }
+
+    intercept() {
+        this._interceptServerMessages();
+        this._interceptClientMessages();
+    }
+};