Repository: websockets/ws
Pull request #969 (2017-01-14 13:18:23)

Prevent `WebSocket#close()` from triggering an infinite loop
This prevents `WebSocket.prototype.close()` from triggering an infinite loop if called from an error listener while connecting.

Refs: https://github.com/websockets/ws/pull/956#issuecomment-272585388.


Patch:
diff --git a/lib/WebSocket.js b/lib/WebSocket.js
index 741a0506..ea351355 100644
--- a/lib/WebSocket.js
+++ b/lib/WebSocket.js
@@ -263,9 +263,12 @@ class WebSocket extends EventEmitter {
   close (code, data) {
     if (this.readyState === WebSocket.CLOSED) return;
     if (this.readyState === WebSocket.CONNECTING) {
-      this._req.abort();
-      this.emit('error', new Error('closed before the connection is established'));
-      return this.finalize(true);
+      if (this._req && !this._req.aborted) {
+        this._req.abort();
+        this.emit('error', new Error('closed before the connection is established'));
+        this.finalize(true);
+      }
+      return;
     }
 
     if (this.readyState === WebSocket.CLOSING) {
@@ -355,16 +358,18 @@ class WebSocket extends EventEmitter {
     if (typeof data === 'number') data = data.toString();
     else if (!data) data = '';
 
-    options = options || {};
-    if (options.fin !== false) options.fin = true;
-    if (options.binary === undefined) options.binary = typeof data !== 'string';
-    if (options.mask === undefined) options.mask = !this._isServer;
-    if (options.compress === undefined) options.compress = true;
+    const opts = Object.assign({
+      fin: true,
+      binary: typeof data !== 'string',
+      mask: !this._isServer,
+      compress: true
+    }, options);
+
     if (!this.extensions[PerMessageDeflate.extensionName]) {
-      options.compress = false;
+      opts.compress = false;
     }
 
-    this._sender.send(data, options, cb);
+    this._sender.send(data, opts, cb);
   }
 
   /**
@@ -375,9 +380,12 @@ class WebSocket extends EventEmitter {
   terminate () {
     if (this.readyState === WebSocket.CLOSED) return;
     if (this.readyState === WebSocket.CONNECTING) {
-      this._req.abort();
-      this.emit('error', new Error('closed before the connection is established'));
-      return this.finalize(true);
+      if (this._req && !this._req.aborted) {
+        this._req.abort();
+        this.emit('error', new Error('closed before the connection is established'));
+        this.finalize(true);
+      }
+      return;
     }
 
     if (this._socket) {
@@ -629,6 +637,7 @@ function initAsClient (address, protocols, options) {
   this._req.on('error', (error) => {
     if (this._req.aborted) return;
 
+    this._req = null;
     this.emit('error', error);
     this.finalize(true);
   });
@@ -642,6 +651,8 @@ function initAsClient (address, protocols, options) {
   });
 
   this._req.on('upgrade', (res, socket, head) => {
+    this._req = null;
+
     const digest = crypto.createHash('sha1')
       .update(key + GUID, 'binary')
       .digest('base64');
@@ -686,7 +697,6 @@ function initAsClient (address, protocols, options) {
       this.extensions[PerMessageDeflate.extensionName] = perMessageDeflate;
     }
 
-    this._req = null;
     this.setSocket(socket, head);
   });
 }


Test patch:
diff --git a/test/WebSocket.test.js b/test/WebSocket.test.js
index cf44a69f..7ac77dd9 100644
--- a/test/WebSocket.test.js
+++ b/test/WebSocket.test.js
@@ -305,77 +305,6 @@ describe('WebSocket', function () {
   });
 
   describe('connection establishing', function () {
-    it('can terminate before connection is established (1/2)', function (done) {
-      const wss = new WebSocketServer({ port: ++port }, () => {
-        const ws = new WebSocket(`ws://localhost:${port}`);
-
-        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
-        ws.on('error', (err) => {
-          assert.ok(err instanceof Error);
-          assert.strictEqual(err.message, 'closed before the connection is established');
-          ws.on('close', () => wss.close(done));
-        });
-        ws.terminate();
-      });
-    });
-
-    it('can terminate before connection is established (2/2)', function (done) {
-      const wss = new WebSocketServer({
-        verifyClient: (info, cb) => setTimeout(cb, 300, true),
-        port: ++port
-      }, () => {
-        const ws = new WebSocket(`ws://localhost:${port}`);
-
-        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
-        ws.on('error', (err) => {
-          assert.ok(err instanceof Error);
-          assert.strictEqual(err.message, 'closed before the connection is established');
-          ws.on('close', () => wss.close(done));
-        });
-        setTimeout(() => ws.terminate(), 150);
-      });
-    });
-
-    it('can close before connection is established (1/2)', function (done) {
-      const wss = new WebSocketServer({ port: ++port }, () => {
-        const ws = new WebSocket(`ws://localhost:${port}`);
-
-        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
-        ws.on('error', (err) => {
-          assert.ok(err instanceof Error);
-          assert.strictEqual(err.message, 'closed before the connection is established');
-          ws.on('close', () => wss.close(done));
-        });
-        ws.close(1001);
-      });
-    });
-
-    it('can close before connection is established (2/2)', function (done) {
-      const wss = new WebSocketServer({
-        verifyClient: (info, cb) => setTimeout(cb, 300, true),
-        port: ++port
-      }, () => {
-        const ws = new WebSocket(`ws://localhost:${port}`);
-
-        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
-        ws.on('error', (err) => {
-          assert.ok(err instanceof Error);
-          assert.strictEqual(err.message, 'closed before the connection is established');
-          ws.on('close', () => wss.close(done));
-        });
-        setTimeout(() => ws.close(1001), 150);
-      });
-    });
-
-    it('can handle error before request is upgraded', function (done) {
-      // Here, we don't create a server, to guarantee that the connection will
-      // fail before the request is upgraded
-      const ws = new WebSocket(`ws://localhost:${++port}`);
-
-      ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
-      ws.on('error', () => done());
-    });
-
     it('invalid server key is denied', function (done) {
       server.createServer(++port, server.handlers.invalidKey, (srv) => {
         const ws = new WebSocket(`ws://localhost:${port}`);
@@ -956,7 +885,50 @@ describe('WebSocket', function () {
   });
 
   describe('#close', function () {
-    it('without invalid first argument throws exception', function (done) {
+    it('closes the connection if called while connecting (1/2)', function (done) {
+      const wss = new WebSocketServer({ port: ++port }, () => {
+        const ws = new WebSocket(`ws://localhost:${port}`);
+
+        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
+        ws.on('error', (err) => {
+          assert.ok(err instanceof Error);
+          assert.strictEqual(err.message, 'closed before the connection is established');
+          ws.on('close', () => wss.close(done));
+        });
+        ws.close(1001);
+      });
+    });
+
+    it('closes the connection if called while connecting (2/2)', function (done) {
+      const wss = new WebSocketServer({
+        verifyClient: (info, cb) => setTimeout(cb, 300, true),
+        port: ++port
+      }, () => {
+        const ws = new WebSocket(`ws://localhost:${port}`);
+
+        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
+        ws.on('error', (err) => {
+          assert.ok(err instanceof Error);
+          assert.strictEqual(err.message, 'closed before the connection is established');
+          ws.on('close', () => wss.close(done));
+        });
+        setTimeout(() => ws.close(1001), 150);
+      });
+    });
+
+    it('can be called from an error listener while connecting', function (done) {
+      const ws = new WebSocket(`ws://localhost:${++port}`);
+
+      ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
+      ws.on('error', (err) => {
+        assert.ok(err instanceof Error);
+        assert.strictEqual(err.code, 'ECONNREFUSED');
+        ws.close();
+        ws.on('close', () => done());
+      });
+    });
+
+    it('throws an error if the first argument is invalid (1/2)', function (done) {
       server.createServer(++port, (srv) => {
         const ws = new WebSocket(`ws://localhost:${port}`);
 
@@ -971,7 +943,7 @@ describe('WebSocket', function () {
       });
     });
 
-    it('without reserved error code 1004 throws exception', function (done) {
+    it('throws an error if the first argument is invalid (2/2)', function (done) {
       server.createServer(++port, (srv) => {
         const ws = new WebSocket(`ws://localhost:${port}`);
 
@@ -986,7 +958,7 @@ describe('WebSocket', function () {
       });
     });
 
-    it('without message is successfully transmitted to the server', function (done) {
+    it('works when close reason is not specified', function (done) {
       server.createServer(++port, (srv) => {
         const ws = new WebSocket(`ws://localhost:${port}`);
 
@@ -1000,22 +972,7 @@ describe('WebSocket', function () {
       });
     });
 
-    it('with message is successfully transmitted to the server', function (done) {
-      server.createServer(++port, (srv) => {
-        const ws = new WebSocket(`ws://localhost:${port}`);
-
-        ws.on('open', () => ws.close(1000, 'some reason'));
-
-        srv.on('close', (code, message, flags) => {
-          assert.ok(flags.masked);
-          assert.strictEqual(message, 'some reason');
-          srv.close(done);
-          ws.terminate();
-        });
-      });
-    });
-
-    it('with encoded message is successfully transmitted to the server', function (done) {
+    it('works when close reason is specified', function (done) {
       server.createServer(++port, (srv) => {
         const ws = new WebSocket(`ws://localhost:${port}`);
 
@@ -1119,6 +1076,51 @@ describe('WebSocket', function () {
     });
   });
 
+  describe('#terminate', function () {
+    it('closes the connection if called while connecting (1/2)', function (done) {
+      const wss = new WebSocketServer({ port: ++port }, () => {
+        const ws = new WebSocket(`ws://localhost:${port}`);
+
+        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
+        ws.on('error', (err) => {
+          assert.ok(err instanceof Error);
+          assert.strictEqual(err.message, 'closed before the connection is established');
+          ws.on('close', () => wss.close(done));
+        });
+        ws.terminate();
+      });
+    });
+
+    it('closes the connection if called while connecting (2/2)', function (done) {
+      const wss = new WebSocketServer({
+        verifyClient: (info, cb) => setTimeout(cb, 300, true),
+        port: ++port
+      }, () => {
+        const ws = new WebSocket(`ws://localhost:${port}`);
+
+        ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
+        ws.on('error', (err) => {
+          assert.ok(err instanceof Error);
+          assert.strictEqual(err.message, 'closed before the connection is established');
+          ws.on('close', () => wss.close(done));
+        });
+        setTimeout(() => ws.terminate(), 150);
+      });
+    });
+
+    it('can be called from an error listener while connecting', function (done) {
+      const ws = new WebSocket(`ws://localhost:${++port}`);
+
+      ws.on('open', () => assert.fail(null, null, 'connect shouldnt be raised here'));
+      ws.on('error', (err) => {
+        assert.ok(err instanceof Error);
+        assert.strictEqual(err.code, 'ECONNREFUSED');
+        ws.terminate();
+        ws.on('close', () => done());
+      });
+    });
+  });
+
   describe('WHATWG API emulation', function () {
     it('should not throw errors when getting and setting', function () {
       const listener = () => {};