Repository: websockets/ws
Pull request #1908 (2021-06-25 20:23:50)

Fix close edge cases
Ensure that `socket.end()` is called if an error occurs simultaneously on
both peers.

Refs: https://github.com/websockets/ws/pull/1902


Patch:
diff --git a/lib/stream.js b/lib/stream.js
index 53706cc9..c3a79369 100644
--- a/lib/stream.js
+++ b/lib/stream.js
@@ -84,13 +84,13 @@ function createWebSocketStream(ws, options) {
 
     // Prevent `ws.terminate()` from being called by `duplex._destroy()`.
     //
-    // - If the state of the `WebSocket` connection is `CONNECTING`,
-    //   `ws.terminate()` is a noop as no socket was assigned.
-    // - Otherwise, the error was re-emitted from the listener of the `'error'`
+    // - If the `'error'` event is emitted before the `'open'` event, then
+    //   `ws.terminate()` is a noop as no socket is assigned.
+    // - Otherwise, the error is re-emitted by the listener of the `'error'`
     //   event of the `Receiver` object. The listener already closes the
     //   connection by calling `ws.close()`. This allows a close frame to be
     //   sent to the other peer. If `ws.terminate()` is called right after this,
-    //   the close frame might not be sent.
+    //   then the close frame might not be sent.
     terminateOnDestroy = false;
     duplex.destroy(err);
   });
diff --git a/lib/websocket.js b/lib/websocket.js
index 40ab4e86..280364a1 100644
--- a/lib/websocket.js
+++ b/lib/websocket.js
@@ -225,7 +225,13 @@ class WebSocket extends EventEmitter {
     }
 
     if (this.readyState === WebSocket.CLOSING) {
-      if (this._closeFrameSent && this._closeFrameReceived) this._socket.end();
+      if (
+        this._closeFrameSent &&
+        (this._closeFrameReceived || this._receiver._writableState.errorEmitted)
+      ) {
+        this._socket.end();
+      }
+
       return;
     }
 
@@ -238,7 +244,13 @@ class WebSocket extends EventEmitter {
       if (err) return;
 
       this._closeFrameSent = true;
-      if (this._closeFrameReceived) this._socket.end();
+
+      if (
+        this._closeFrameReceived ||
+        this._receiver._writableState.errorEmitted
+      ) {
+        this._socket.end();
+      }
     });
 
     //


Test patch:
diff --git a/test/create-websocket-stream.test.js b/test/create-websocket-stream.test.js
index bcd24097..b96ac9b1 100644
--- a/test/create-websocket-stream.test.js
+++ b/test/create-websocket-stream.test.js
@@ -204,6 +204,8 @@ describe('createWebSocketStream', () => {
 
     it('reemits errors', (done) => {
       let duplexCloseEventEmitted = false;
+      let serverClientCloseEventEmitted = false;
+
       const wss = new WebSocket.Server({ port: 0 }, () => {
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
         const duplex = createWebSocketStream(ws);
@@ -218,6 +220,7 @@ describe('createWebSocketStream', () => {
 
           duplex.on('close', () => {
             duplexCloseEventEmitted = true;
+            if (serverClientCloseEventEmitted) wss.close(done);
           });
         });
       });
@@ -225,10 +228,11 @@ describe('createWebSocketStream', () => {
       wss.on('connection', (ws) => {
         ws._socket.write(Buffer.from([0x85, 0x00]));
         ws.on('close', (code, reason) => {
-          assert.ok(duplexCloseEventEmitted);
           assert.strictEqual(code, 1002);
           assert.strictEqual(reason, '');
-          wss.close(done);
+
+          serverClientCloseEventEmitted = true;
+          if (duplexCloseEventEmitted) wss.close(done);
         });
       });
     });
diff --git a/test/websocket.test.js b/test/websocket.test.js
index 03f984fc..5642d968 100644
--- a/test/websocket.test.js
+++ b/test/websocket.test.js
@@ -430,6 +430,8 @@ describe('WebSocket', () => {
   describe('Events', () => {
     it("emits an 'error' event if an error occurs", (done) => {
       let clientCloseEventEmitted = false;
+      let serverClientCloseEventEmitted = false;
+
       const wss = new WebSocket.Server({ port: 0 }, () => {
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
 
@@ -442,19 +444,22 @@ describe('WebSocket', () => {
           );
 
           ws.on('close', (code, reason) => {
-            clientCloseEventEmitted = true;
             assert.strictEqual(code, 1006);
             assert.strictEqual(reason, '');
+
+            clientCloseEventEmitted = true;
+            if (serverClientCloseEventEmitted) wss.close(done);
           });
         });
       });
 
       wss.on('connection', (ws) => {
         ws.on('close', (code, reason) => {
-          assert.ok(clientCloseEventEmitted);
           assert.strictEqual(code, 1002);
           assert.strictEqual(reason, '');
-          wss.close(done);
+
+          serverClientCloseEventEmitted = true;
+          if (clientCloseEventEmitted) wss.close(done);
         });
 
         ws._socket.write(Buffer.from([0x85, 0x00]));
@@ -1419,16 +1424,19 @@ describe('WebSocket', () => {
     });
 
     it('honors the `mask` option', (done) => {
+      let clientCloseEventEmitted = false;
       let serverClientCloseEventEmitted = false;
+
       const wss = new WebSocket.Server({ port: 0 }, () => {
         const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
 
         ws.on('open', () => ws.send('hi', { mask: false }));
         ws.on('close', (code, reason) => {
-          assert.ok(serverClientCloseEventEmitted);
           assert.strictEqual(code, 1002);
           assert.strictEqual(reason, '');
-          wss.close(done);
+
+          clientCloseEventEmitted = true;
+          if (serverClientCloseEventEmitted) wss.close(done);
         });
       });
 
@@ -1450,9 +1458,11 @@ describe('WebSocket', () => {
           );
 
           ws.on('close', (code, reason) => {
-            serverClientCloseEventEmitted = true;
             assert.strictEqual(code, 1006);
             assert.strictEqual(reason, '');
+
+            serverClientCloseEventEmitted = true;
+            if (clientCloseEventEmitted) wss.close(done);
           });
         });
       });
@@ -2628,6 +2638,87 @@ describe('WebSocket', () => {
       });
     });
 
+    it('handles a close frame received while compressing data', (done) => {
+      const wss = new WebSocket.Server(
+        {
+          perMessageDeflate: true,
+          port: 0
+        },
+        () => {
+          const ws = new WebSocket(`ws://localhost:${wss.address().port}`, {
+            perMessageDeflate: { threshold: 0 }
+          });
+
+          ws.on('open', () => {
+            ws._receiver.on('conclude', () => {
+              assert.ok(ws._sender._deflating);
+            });
+
+            ws.send('foo');
+            ws.send('bar');
+            ws.send('baz');
+            ws.send('qux');
+          });
+        }
+      );
+
+      wss.on('connection', (ws) => {
+        const messages = [];
+
+        ws.on('message', (message) => {
+          messages.push(message);
+        });
+
+        ws.on('close', (code, reason) => {
+          assert.deepStrictEqual(messages, ['foo', 'bar', 'baz', 'qux']);
+          assert.strictEqual(code, 1000);
+          assert.strictEqual(reason, '');
+          wss.close(done);
+        });
+
+        ws.close(1000);
+      });
+    });
+
+    describe('#close', () => {
+      it('can be used while data is being decompressed', (done) => {
+        const wss = new WebSocket.Server(
+          {
+            perMessageDeflate: true,
+            port: 0
+          },
+          () => {
+            const messages = [];
+            const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
+
+            ws.on('open', () => {
+              ws._socket.on('end', () => {
+                assert.strictEqual(ws._receiver._state, 5);
+              });
+            });
+
+            ws.on('message', (message) => {
+              if (messages.push(message) > 1) return;
+
+              ws.close(1000);
+            });
+
+            ws.on('close', (code, reason) => {
+              assert.deepStrictEqual(messages, ['', '', '', '']);
+              assert.strictEqual(code, 1000);
+              assert.strictEqual(reason, '');
+              wss.close(done);
+            });
+          }
+        );
+
+        wss.on('connection', (ws) => {
+          const buf = Buffer.from('c10100c10100c10100c10100', 'hex');
+          ws._socket.write(buf);
+        });
+      });
+    });
+
     describe('#send', () => {
       it('ignores the `compress` option if the extension is disabled', (done) => {
         const wss = new WebSocket.Server({ port: 0 }, () => {
@@ -2760,4 +2851,118 @@ describe('WebSocket', () => {
       });
     });
   });
+
+  describe('Connection close edge cases', () => {
+    it('closes cleanly after simultaneous errors (1/2)', (done) => {
+      let clientCloseEventEmitted = false;
+      let serverClientCloseEventEmitted = false;
+
+      const wss = new WebSocket.Server({ port: 0 }, () => {
+        const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
+
+        ws.on('error', (err) => {
+          assert.ok(err instanceof RangeError);
+          assert.strictEqual(err.code, 'WS_ERR_INVALID_OPCODE');
+          assert.strictEqual(
+            err.message,
+            'Invalid WebSocket frame: invalid opcode 5'
+          );
+
+          ws.on('close', (code, reason) => {
+            assert.strictEqual(code, 1006);
+            assert.strictEqual(reason, '');
+
+            clientCloseEventEmitted = true;
+            if (serverClientCloseEventEmitted) wss.close(done);
+          });
+        });
+
+        ws.on('open', () => {
+          // Write an invalid frame in both directions to trigger simultaneous
+          // failure.
+          const chunk = Buffer.from([0x85, 0x00]);
+
+          wss.clients.values().next().value._socket.write(chunk);
+          ws._socket.write(chunk);
+        });
+      });
+
+      wss.on('connection', (ws) => {
+        ws.on('error', (err) => {
+          assert.ok(err instanceof RangeError);
+          assert.strictEqual(err.code, 'WS_ERR_INVALID_OPCODE');
+          assert.strictEqual(
+            err.message,
+            'Invalid WebSocket frame: invalid opcode 5'
+          );
+
+          ws.on('close', (code, reason) => {
+            assert.strictEqual(code, 1006);
+            assert.strictEqual(reason, '');
+
+            serverClientCloseEventEmitted = true;
+            if (clientCloseEventEmitted) wss.close(done);
+          });
+        });
+      });
+    });
+
+    it('closes cleanly after simultaneous errors (2/2)', (done) => {
+      let clientCloseEventEmitted = false;
+      let serverClientCloseEventEmitted = false;
+
+      const wss = new WebSocket.Server({ port: 0 }, () => {
+        const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
+
+        ws.on('error', (err) => {
+          assert.ok(err instanceof RangeError);
+          assert.strictEqual(err.code, 'WS_ERR_INVALID_OPCODE');
+          assert.strictEqual(
+            err.message,
+            'Invalid WebSocket frame: invalid opcode 5'
+          );
+
+          ws.on('close', (code, reason) => {
+            assert.strictEqual(code, 1006);
+            assert.strictEqual(reason, '');
+
+            clientCloseEventEmitted = true;
+            if (serverClientCloseEventEmitted) wss.close(done);
+          });
+        });
+
+        ws.on('open', () => {
+          // Write an invalid frame in both directions and change the
+          // `readyState` to `WebSocket.CLOSING`.
+          const chunk = Buffer.from([0x85, 0x00]);
+          const serverWs = wss.clients.values().next().value;
+
+          serverWs._socket.write(chunk);
+          serverWs.close();
+
+          ws._socket.write(chunk);
+          ws.close();
+        });
+      });
+
+      wss.on('connection', (ws) => {
+        ws.on('error', (err) => {
+          assert.ok(err instanceof RangeError);
+          assert.strictEqual(err.code, 'WS_ERR_INVALID_OPCODE');
+          assert.strictEqual(
+            err.message,
+            'Invalid WebSocket frame: invalid opcode 5'
+          );
+
+          ws.on('close', (code, reason) => {
+            assert.strictEqual(code, 1006);
+            assert.strictEqual(reason, '');
+
+            serverClientCloseEventEmitted = true;
+            if (clientCloseEventEmitted) wss.close(done);
+          });
+        });
+      });
+    });
+  });
 });