Fixing request resolution order of operations and adding test case to…

… verify. Resolves #171

lib/WebSocketRequest.js

@@ -248,7 +248,7 @@ WebSocketRequest.prototype.parseCookies = function(str) {
 
 WebSocketRequest.prototype.accept = function(acceptedProtocol, allowedOrigin, cookies) {
     this._verifyResolution();
-
+    
     // TODO: Handle extensions
 
     var protocolFullCase;
@@ -427,6 +427,11 @@ WebSocketRequest.prototype.accept = function(acceptedProtocol, allowedOrigin, co
     //     response += 'Sec-WebSocket-Extensions: ' + negotiatedExtensions.join(', ') + '\r\n';
     // }
 
+    // Mark the request resolved now so that the user can't call accept or
+    // reject a second time.
+    this._resolved = true;
+    this.emit('requestResolved', this);
+
     response += '\r\n';
 
     var connection = new WebSocketConnection(this.socket, [], acceptedProtocol, false, this.serverConfig);
@@ -453,13 +458,18 @@ WebSocketRequest.prototype.accept = function(acceptedProtocol, allowedOrigin, co
         });
     }
 
-    this.emit('requestAccepted', connection);    
+    this.emit('requestAccepted', connection);
     return connection;
 };
 
 WebSocketRequest.prototype.reject = function(status, reason, extraHeaders) {
     this._verifyResolution();
 
+    // Mark the request resolved now so that the user can't call accept or
+    // reject a second time.
+    this._resolved = true;
+    this.emit('requestResolved', this);
+
     if (typeof(status) !== 'number') {
         status = 403;
     }
@@ -498,7 +508,6 @@ WebSocketRequest.prototype._verifyResolution = function() {
     if (this._resolved) {
         throw new Error('WebSocketRequest may only be accepted or rejected one time.');
     }
-    this._resolved = true;
 };
 
 function cleanupFailedConnection(connection) {

lib/WebSocketServer.js

@@ -27,9 +27,11 @@ var WebSocketServer = function WebSocketServer(config) {
 
     this._handlers = {
         upgrade: this.handleUpgrade.bind(this),
-        requestAccepted: this.handleRequestAccepted.bind(this)
+        requestAccepted: this.handleRequestAccepted.bind(this),
+        requestResolved: this.handleRequestResolved.bind(this)
     };
     this.connections = [];
+    this.pendingRequests = [];
     if (config) {
         this.mount(config);
     }
@@ -154,6 +156,11 @@ WebSocketServer.prototype.closeAllConnections = function() {
     this.connections.forEach(function(connection) {
         connection.close();
     });
+    this.pendingRequests.forEach(function(request) {
+        process.nextTick(function() {
+          request.reject(503); // HTTP 503 Service Unavailable
+        });
+    });
 };
 
 WebSocketServer.prototype.broadcast = function(data) {
@@ -196,8 +203,11 @@ WebSocketServer.prototype.handleUpgrade = function(request, socket) {
         debug('Invalid handshake: %s', e.message);
         return;
     }
+
+    this.pendingRequests.push(wsRequest);
 
     wsRequest.once('requestAccepted', this._handlers.requestAccepted);
+    wsRequest.once('requestResolved', this._handlers.requestResolved);
 
     if (!this.config.autoAcceptConnections && utils.eventEmitterListenerCount(this, 'request') > 0) {
         this.emit('request', wsRequest);
@@ -227,4 +237,9 @@ WebSocketServer.prototype.handleConnectionClose = function(connection, closeReas
     this.emit('close', connection, closeReason, description);
 };
 
+WebSocketServer.prototype.handleRequestResolved = function(request) {
+    var index = this.pendingRequests.indexOf(request);
+    if (index !== -1) { this.pendingRequests.splice(index, 1); }
+};
+
 module.exports = WebSocketServer;

test/shared/test-server.js

@@ -34,7 +34,9 @@ function stopServer() {
     wsServer.shutDown();
     server.close();
   }
-  catch(e) { /* do nothing */ }
+  catch(e) {
+    console.warn("stopServer threw", e);
+  }
 }
 
 module.exports = {

test/unit/request.js

@@ -4,9 +4,13 @@ var WebSocketClient = require('../../lib/WebSocketClient');
 var server = require('../shared/test-server');
 var stopServer = server.stopServer;
 
-var testCase = test('Request can only be rejected or accepted once.', function(t) {
+test('Request can only be rejected or accepted once.', function(t) {
   t.plan(6);
 
+  t.on('end', function() {
+    stopServer();
+  });
+
   server.prepare(function(err, wsServer) {
     if (err) {
       t.fail('Unable to start test server');
@@ -49,6 +53,53 @@ var testCase = test('Request can only be rejected or accepted once.', function(t
   });
 });
 
-testCase.on('end', function() {
-  stopServer();
+
+test('Protocol mismatch should be handled gracefully', function(t) {
+  var wsServer;
+
+  t.test('setup', function(t) {
+    server.prepare(function(err, result) {
+      if (err) {
+        t.fail('Unable to start test server');
+        return t.end();
+      }
+
+      wsServer = result;
+      t.end();
+    });
+  });
+
+  t.test('mismatched protocol connection', function(t) {
+    t.plan(2);
+    wsServer.on('request', handleRequest);
+
+    var client = new WebSocketClient();
+
+    var timer = setTimeout(function() {
+      t.fail('Timeout waiting for client event');
+    }, 2000);
+
+    client.connect('ws://localhost:64321/', 'some_protocol_here');
+    client.on('connect', function(connection) {
+      clearTimeout(timer);
+      connection.close();
+      t.fail('connect event should not be emitted on client');
+    });
+    client.on('connectFailed', function() {
+      clearTimeout(timer);
+      t.pass('connectFailed event should be emitted on client');
+    });
+
+
+
+    function handleRequest(request) {
+      var accept = request.accept.bind(request, 'this_is_the_wrong_protocol', request.origin);
+      t.throws(accept, 'request.accept() should throw');
+    }
+  });
+
+  t.test('teardown', function(t) {
+    stopServer();
+    t.end();
+  });
 });