Fix subprotocol handling

doc/ws.md

@@ -50,16 +50,19 @@ if `verifyClient` is provided with two arguments then those are:
     reason phrase.
 
 
-If `handleProtocols` is not set then the handshake is automatically accepted,
-otherwise the function takes two arguments:
+`handleProtocols` takes two arguments:
 
 - `protocols` {Array} The list of WebSocket subprotocols indicated by the
   client in the `Sec-WebSocket-Protocol` header.
 - `request` {http.IncomingMessage} The client HTTP GET request.
 
-If returned value is `false` then the handshake is rejected with the HTTP 401
-status code, otherwise the returned value sets the value of the
-`Sec-WebSocket-Protocol` header in the HTTP 101 response.
+The returned value sets the value of the `Sec-WebSocket-Protocol` header in the
+HTTP 101 response. If returned value is `false` the header is not added in the
+response.
+
+If `handleProtocols` is not set then the first of the client's requested
+subprotocols is used.
+
 
 `perMessageDeflate` can be used to control the behavior of
 [permessage-deflate extension][permessage-deflate].

lib/websocket-server.js

@@ -196,18 +196,6 @@ class WebSocketServer extends EventEmitter {
       }
     }
 
-    var protocol = (req.headers['sec-websocket-protocol'] || '').split(/, */);
-
-    //
-    // Optionally call external protocol selection handler.
-    //
-    if (this.options.handleProtocols) {
-      protocol = this.options.handleProtocols(protocol, req);
-      if (protocol === false) return abortConnection(socket, 401);
-    } else {
-      protocol = protocol[0];
-    }
-
     //
     // Optionally call external client verification handler.
     //
@@ -222,29 +210,28 @@ class WebSocketServer extends EventEmitter {
         this.options.verifyClient(info, (verified, code, message) => {
           if (!verified) return abortConnection(socket, code || 401, message);
 
-          this.completeUpgrade(protocol, extensions, req, socket, head, cb);
+          this.completeUpgrade(extensions, req, socket, head, cb);
         });
         return;
       }
 
       if (!this.options.verifyClient(info)) return abortConnection(socket, 401);
     }
 
-    this.completeUpgrade(protocol, extensions, req, socket, head, cb);
+    this.completeUpgrade(extensions, req, socket, head, cb);
   }
 
   /**
    * Upgrade the connection to WebSocket.
    *
-   * @param {String} protocol The chosen subprotocol
    * @param {Object} extensions The accepted extensions
    * @param {http.IncomingMessage} req The request object
    * @param {net.Socket} socket The network socket between the server and client
    * @param {Buffer} head The first packet of the upgraded stream
    * @param {Function} cb Callback
    * @private
    */
-  completeUpgrade (protocol, extensions, req, socket, head, cb) {
+  completeUpgrade (extensions, req, socket, head, cb) {
     //
     // Destroy the socket if the client has already sent a FIN packet.
     //
@@ -262,11 +249,26 @@ class WebSocketServer extends EventEmitter {
     ];
 
     const ws = new WebSocket(null);
+    var protocol = req.headers['sec-websocket-protocol'];
 
     if (protocol) {
-      headers.push(`Sec-WebSocket-Protocol: ${protocol}`);
-      ws.protocol = protocol;
+      protocol = protocol.trim().split(/ *, */);
+
+      //
+      // Optionally call external protocol selection handler.
+      //
+      if (this.options.handleProtocols) {
+        protocol = this.options.handleProtocols(protocol, req);
+      } else {
+        protocol = protocol[0];
+      }
+
+      if (protocol) {
+        headers.push(`Sec-WebSocket-Protocol: ${protocol}`);
+        ws.protocol = protocol;
+      }
     }
+
     if (extensions[PerMessageDeflate.extensionName]) {
       const params = extensions[PerMessageDeflate.extensionName].params;
       const value = extension.format({

test/websocket-server.test.js

@@ -649,7 +649,7 @@ describe('WebSocketServer', function () {
     });
 
     describe('`handleProtocols`', function () {
-      it('can select the last protocol', function (done) {
+      it('allows to select a subprotocol', function (done) {
         const handleProtocols = (protocols, request) => {
           assert.ok(request instanceof http.IncomingMessage);
           assert.strictEqual(request.url, '/');
@@ -667,28 +667,6 @@ describe('WebSocketServer', function () {
           });
         });
       });
-
-      it('closes the connection if return value is `false`', function (done) {
-        const wss = new WebSocket.Server({
-          handleProtocols: (protocols) => false,
-          port: 0
-        }, () => {
-          const req = http.get({
-            port: wss.address().port,
-            headers: {
-              'Connection': 'Upgrade',
-              'Upgrade': 'websocket',
-              'Sec-WebSocket-Key': 'dGhlIHNhbXBsZSBub25jZQ==',
-              'Sec-WebSocket-Version': 13
-            }
-          });
-
-          req.on('response', (res) => {
-            assert.strictEqual(res.statusCode, 401);
-            wss.close(done);
-          });
-        });
-      });
     });
 
     it('emits the `headers` event', function (done) {

test/websocket.test.js

@@ -614,9 +614,10 @@ describe('WebSocket', function () {
     });
 
     it('fails if server sends a subprotocol when none was requested', function (done) {
-      const wss = new WebSocket.Server({
-        handleProtocols: () => 'foo',
-        server
+      const wss = new WebSocket.Server({ server });
+
+      wss.on('headers', (headers) => {
+        headers.push('Sec-WebSocket-Protocol: foo');
       });
 
       const ws = new WebSocket(`ws://localhost:${server.address().port}`);