From b32016995e9620ed6d5826464b2c3ee931e2240a Mon Sep 17 00:00:00 2001
From: Ido schachter <id0Sch@users.noreply.github.com>
Date: Thu, 27 Jul 2017 08:42:35 +0300
Subject: [PATCH] [feature] Allow client to specify handshake request timeout
 (#1177)

---
 doc/ws.md              |  1 +
 lib/WebSocket.js       | 10 ++++++++++
 test/WebSocket.test.js | 15 +++++++++++++++
 3 files changed, 26 insertions(+)

diff --git a/doc/ws.md b/doc/ws.md
index 4da6e4ea7..9e5e6f413 100644
--- a/doc/ws.md
+++ b/doc/ws.md
@@ -170,6 +170,7 @@ This class represents a WebSocket. It extends the `EventEmitter`.
 - `protocols` {String|Array} The list of subprotocols.
 - `options` {Object}
   - `protocol` {String} Value of the `Sec-WebSocket-Protocol` header.
+  - `handshakeTimeout` {Number} Timeout in milliseconds for the handshake request.
   - `perMessageDeflate` {Boolean|Object} Enable/disable permessage-deflate.
   - `localAddress` {String} Local interface to bind for network connections.
   - `protocolVersion` {Number} Value of the `Sec-WebSocket-Version` header.
diff --git a/lib/WebSocket.js b/lib/WebSocket.js
index 320c725e1..760ae34fa 100644
--- a/lib/WebSocket.js
+++ b/lib/WebSocket.js
@@ -471,6 +471,7 @@ function initAsServerClient (socket, head, options) {
  * @param {Object} options Connection options
  * @param {String} options.protocol Value of the `Sec-WebSocket-Protocol` header
  * @param {(Boolean|Object)} options.perMessageDeflate Enable/disable permessage-deflate
+ * @param {Number} options.handshakeTimeout Timeout in milliseconds for the handshake request
  * @param {String} options.localAddress Local interface to bind for network connections
  * @param {Number} options.protocolVersion Value of the `Sec-WebSocket-Version` header
  * @param {Object} options.headers An object containing request headers
@@ -493,6 +494,7 @@ function initAsClient (address, protocols, options) {
     protocolVersion: protocolVersions[1],
     protocol: protocols.join(','),
     perMessageDeflate: true,
+    handshakeTimeout: null,
     localAddress: null,
     headers: null,
     family: null,
@@ -632,6 +634,14 @@ function initAsClient (address, protocols, options) {
 
   this._req = httpObj.get(requestOptions);
 
+  if (options.handshakeTimeout) {
+    this._req.setTimeout(options.handshakeTimeout, () => {
+      this._req.abort();
+      this.emit('error', new Error('opening handshake has timed out'));
+      this.finalize(true);
+    });
+  }
+
   this._req.on('error', (error) => {
     if (this._req.aborted) return;
 
diff --git a/test/WebSocket.test.js b/test/WebSocket.test.js
index 0ebfbe52d..1aa45d12e 100644
--- a/test/WebSocket.test.js
+++ b/test/WebSocket.test.js
@@ -461,6 +461,21 @@ describe('WebSocket', function () {
         req.abort();
       });
     });
+
+    it('emits an error if the opening handshake timeout expires', function (done) {
+      server.once('upgrade', (req, socket) => socket.on('end', socket.end));
+
+      const ws = new WebSocket(`ws://localhost:${port}`, null, {
+        handshakeTimeout: 100
+      });
+
+      ws.on('open', () => assert.fail(null, null, 'connect shouldn\'t be raised here'));
+      ws.on('error', (err) => {
+        assert.ok(err instanceof Error);
+        assert.strictEqual(err.message, 'opening handshake has timed out');
+        done();
+      });
+    });
   });
 
   describe('connection with query string', function () {