From 1aba8152c02d9cbab03e9b1d02610c60d4ddec2c Mon Sep 17 00:00:00 2001
From: Luigi Pinca <luigipinca@gmail.com>
Date: Sat, 14 Jan 2017 14:16:55 +0100
Subject: [PATCH] [fix] 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.
---
 lib/WebSocket.js       |  22 +++--
 test/WebSocket.test.js | 182 +++++++++++++++++++++--------------------
 2 files changed, 107 insertions(+), 97 deletions(-)

diff --git a/lib/WebSocket.js b/lib/WebSocket.js
index 741a05066..02b9fb951 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) {
@@ -375,9 +378,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 +635,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 +649,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 +695,6 @@ function initAsClient (address, protocols, options) {
       this.extensions[PerMessageDeflate.extensionName] = perMessageDeflate;
     }
 
-    this._req = null;
     this.setSocket(socket, head);
   });
 }
diff --git a/test/WebSocket.test.js b/test/WebSocket.test.js
index cf44a69fb..7ac77dd96 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 = () => {};