[feature] Add headers argument to verifyClient() callback (#1379)

Add ability to specify custom headers when rejecting the handshake.
websockets · May 12, 2018 · d871bdf · d871bdf
1 parent bb9c21c
commit d871bdf
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 8 deletions.
diff --git a/doc/ws.md b/doc/ws.md
@@ -53,6 +53,8 @@ if `verifyClient` is provided with two arguments then those are:
     error status code to be sent to the client.
   - `name` {String} When `result` is `false` this field determines the HTTP
     reason phrase.
+  - `headers` {Object} When `result` is `false` this field determines additional
+    HTTP headers to be sent to the client. For example, `{ 'Retry-After': 120 }`.
 
 
 `handleProtocols` takes two arguments:

diff --git a/lib/websocket-server.js b/lib/websocket-server.js
@@ -204,8 +204,10 @@ class WebSocketServer extends EventEmitter {
       };
 
       if (this.options.verifyClient.length === 2) {
-        this.options.verifyClient(info, (verified, code, message) => {
-          if (!verified) return abortHandshake(socket, code || 401, message);
+        this.options.verifyClient(info, (verified, code, message, headers) => {
+          if (!verified) {
+            return abortHandshake(socket, code || 401, message, headers);
+          }
 
           this.completeUpgrade(extensions, req, socket, head, cb);
         });
@@ -330,17 +332,22 @@ function socketOnError () {
  * @param {net.Socket} socket The socket of the upgrade request
  * @param {Number} code The HTTP response status code
  * @param {String} [message] The HTTP response body
+ * @param {Object} [headers] Additional HTTP response headers
  * @private
  */
-function abortHandshake (socket, code, message) {
+function abortHandshake (socket, code, message, headers) {
   if (socket.writable) {
     message = message || http.STATUS_CODES[code];
+    headers = Object.assign({
+      'Connection': 'close',
+      'Content-type': 'text/html',
+      'Content-Length': Buffer.byteLength(message)
+    }, headers);
+
     socket.write(
       `HTTP/1.1 ${code} ${http.STATUS_CODES[code]}\r\n` +
-      'Connection: close\r\n' +
-      'Content-type: text/html\r\n' +
-      `Content-Length: ${Buffer.byteLength(message)}\r\n` +
-      '\r\n' +
+      Object.keys(headers).map(h => `${h}: ${headers[h]}`).join('\r\n') +
+      '\r\n\r\n' +
       message
     );
   }

diff --git a/test/websocket-server.test.js b/test/websocket-server.test.js
@@ -551,7 +551,7 @@ describe('WebSocketServer', function () {
         });
       });
 
-      it('can reject client asynchronously with status code', function (done) {
+      it('can reject client asynchronously w/ status code', function (done) {
         const wss = new WebSocket.Server({
           verifyClient: (info, cb) => process.nextTick(cb, false, 404),
           port: 0
@@ -576,6 +576,35 @@ describe('WebSocketServer', function () {
           done(new Error("Unexpected 'connection' event"));
         });
       });
+
+      it('can reject client asynchronously w/ custom headers', function (done) {
+        const wss = new WebSocket.Server({
+          verifyClient: (info, cb) => {
+            process.nextTick(cb, false, 503, '', { 'Retry-After': 120 });
+          },
+          port: 0
+        }, () => {
+          const req = http.get({
+            port: wss.address().port,
+            headers: {
+              'Connection': 'Upgrade',
+              'Upgrade': 'websocket',
+              'Sec-WebSocket-Key': 'dGhlIHNhbXBsZSBub25jZQ==',
+              'Sec-WebSocket-Version': 8
+            }
+          });
+
+          req.on('response', (res) => {
+            assert.strictEqual(res.statusCode, 503);
+            assert.strictEqual(res.headers['retry-after'], '120');
+            wss.close(done);
+          });
+        });
+
+        wss.on('connection', (ws) => {
+          done(new Error("Unexpected 'connection' event"));
+        });
+      });
     });
 
     it("doesn't emit the 'connection' event if socket is closed prematurely", function (done) {