Merge pull request #2 from gemini-testing/dd.correct_close_connection

fix: listen close event from client and server
gemini-testing · Jun 25, 2019 · 92ca51f · 92ca51f
2 parents cb1ce77 + 003a07e
commit 92ca51f
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 36 deletions.
diff --git a/README.md b/README.md
@@ -395,27 +395,28 @@ 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.
-    ```
+*  **wsInterceptClientMsg**: Is a handler which is called when a websocket message is intercepted on its way to the server from 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({
         ...
-        wsInterceptServerMsg: (data, flags) {
+        wsInterceptClientMsg: (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.
-     ```
+*  **wsInterceptServerMsg**: Is a handler which is called when a websocket message is intercepted on its way to the client from 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({
         ...
-        wsInterceptClientMsg: (data, flags) {
+        wsInterceptServerMsg: (data, flags) {
           return typeof data === 'string ? data.toUpperCase() : data;
         }
         ...
       })
     ```
 
+
 **NOTE:**
 `options.ws` and `options.ssl` are optional.
 `options.target` and `options.forward` cannot both be missing
@@ -434,8 +435,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.
+* `wsClientMsg`: This event is emitted after webscoket mesage is sended from the client to the server.
+* `wsServerMsg`: This event is emitted after websocket message is sended from the server 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`.

diff --git a/lib/http-proxy/passes/ws-incoming.js b/lib/http-proxy/passes/ws-incoming.js
@@ -145,12 +145,8 @@ module.exports = {
       //
       socket.write(createHttpHeader('HTTP/1.1 101 Switching Protocols', proxyRes.headers));
 
-      if (options.wsInterceptServerMsg || options.wsInterceptClientMsg) {
-        WsInterceptor.create({socket, options, proxyReq, proxyRes, proxySocket}).intercept();
-      }
-      else {
-        proxySocket.pipe(socket).pipe(proxySocket);
-      }
+      const wsInterceptor = WsInterceptor.create({socket, options, proxyReq, proxyRes, proxySocket});
+      wsInterceptor.startDataTransfer();
 
       server.emit('open', proxySocket);
       server.emit('proxySocket', proxySocket);  //DEPRECATED.

diff --git a/lib/http-proxy/ws/interceptor.js b/lib/http-proxy/ws/interceptor.js
@@ -5,16 +5,16 @@ const Extensions = require('ws/lib/Extensions');
 const Receiver = require('ws/lib/Receiver');
 const Sender = require('ws/lib/Sender');
 
-const acceptExtensions = ({extenstions, isServer}) => {
+const acceptExtensions = ({extensions, isServer}) => {
     const {extensionName} = PerMessageDeflate;
-    const extenstion = extenstions[extensionName];
+    const extension = extensions[extensionName];
 
-    if (!extenstion) {
+    if (!extension) {
         return {};
     }
 
     const perMessageDeflate = new PerMessageDeflate({}, isServer);
-    perMessageDeflate.accept(extenstion);
+    perMessageDeflate.accept(extension);
 
     return {[extensionName]: perMessageDeflate};
 };
@@ -46,18 +46,23 @@ module.exports = class Interceptor {
         this._proxyReq = proxyReq;
         this._proxyRes = proxyRes;
         this._proxySocket = proxySocket;
+        this._isSocketOpened = true;
 
         this._configure();
     }
 
     _configure() {
+        this._proxySocket.on('close', () => {
+            this._isSocketOpened = false;
+        });
+
         const secWsExtensions = this._proxyRes.headers['sec-websocket-extensions'];
-        const extenstions = Extensions.parse(secWsExtensions);
-        this._isCompressed = secWsExtensions && secWsExtensions.indexOf('permessage-deflate') != -1;
+        const extensions = Extensions.parse(secWsExtensions);
+        this._isCompressed = secWsExtensions && secWsExtensions.includes('permessage-deflate');
 
         // 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;
+        this._clientExtensions = this._isCompressed ? acceptExtensions({extensions, isServer: false}) : null;
+        this._serverExtensions = this._isCompressed ? acceptExtensions({extensions, isServer: true}) : null;
     }
 
     _getDataSender({sender, event, options}) {
@@ -69,35 +74,37 @@ module.exports = class Interceptor {
         };
     }
 
-    _interceptServerMessages() {
-        const receiver = new Receiver(this._clientExtenstions);
-        const sender = new Sender(this._proxySocket, this._serverExtenstions);
+    _interceptClientMessages() {
+        const receiver = new Receiver(this._clientExtensions);
+        const sender = new Sender(this._proxySocket, this._serverExtensions);
 
         // 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});
+        const dataSender = this._getDataSender({sender, event: 'wsClientMsg', options});
 
-        receiver.ontext = getMsgHandler({interceptor: this._options.wsInterceptServerMsg, dataSender, binary: false});
-        receiver.onbinary = getMsgHandler({interceptor: this._options.wsInterceptServerMsg, dataSender, binary: true});
+        receiver.ontext = getMsgHandler({interceptor: this._options.wsInterceptClientMsg, dataSender, binary: false});
+        receiver.onbinary = getMsgHandler({interceptor: this._options.wsInterceptClientMsg, dataSender, binary: true});
+        receiver.onclose = (code, msg, {masked: mask}) => this._isSocketOpened && sender.close(code, msg, mask);
 
         this._socket.on('data', (data) => receiver.add(data));
     }
 
-    _interceptClientMessages() {
-        const receiver = new Receiver(this._serverExtenstions);
-        const sender = new Sender(this._socket, this._clientExtenstions);
+    _interceptServerMessages() {
+        const receiver = new Receiver(this._serverExtensions);
+        const sender = new Sender(this._socket, this._clientExtensions);
 
         const options = {mask: false};
-        const dataSender = this._getDataSender({sender, event: 'wsClientMsg', options});
+        const dataSender = this._getDataSender({sender, event: 'wsServerMsg', options});
 
-        receiver.ontext = getMsgHandler({interceptor: this._options.wsInterceptClientMsg, dataSender, binary: false});
-        receiver.onbinary = getMsgHandler({interceptor: this._options.wsInterceptClientMsg, dataSender, binary: true});
+        receiver.ontext = getMsgHandler({interceptor: this._options.wsInterceptServerMsg, dataSender, binary: false});
+        receiver.onbinary = getMsgHandler({interceptor: this._options.wsInterceptServerMsg, dataSender, binary: true});
+        receiver.onclose = (code, msg, {masked: mask}) => this._isSocketOpened && sender.close(code, msg, mask);
 
         this._proxySocket.on('data', (data) => receiver.add(data));
     }
 
-    intercept() {
-        this._interceptServerMessages();
+    startDataTransfer() {
         this._interceptClientMessages();
+        this._interceptServerMessages();
     }
 };