Merge pull request #344 from nkzawa/patch-6

Fix sockets can stay open when upgrade failed

Author: rauchg

lib/server.js

@@ -334,10 +334,14 @@ Server.prototype.onWebSocket = function(req, socket){
   req.websocket = socket;
 
   if (id) {
-    if (!this.clients[id]) {
+    var client = this.clients[id];
+    if (!client) {
       debug('upgrade attempt for closed client');
       socket.close();
-    } else if (this.clients[id].upgraded) {
+    } else if (client.upgrading) {
+      debug('transport has already been trying to upgrade');
+      socket.close();
+    } else if (client.upgraded) {
       debug('transport had already been upgraded');
       socket.close();
     } else {
@@ -348,7 +352,7 @@ Server.prototype.onWebSocket = function(req, socket){
       } else {
         transport.supportsBinary = true;
       }
-      this.clients[id].maybeUpgrade(transport);
+      client.maybeUpgrade(transport);
     }
   } else {
     this.handshake(req._query.transport, req);

lib/socket.js

@@ -20,6 +20,7 @@ module.exports = Socket;
 function Socket (id, server, transport, req) {
   this.id = id;
   this.server = server;
+  this.upgrading = false;
   this.upgraded = false;
   this.readyState = 'opening';
   this.writeBuffer = [];
@@ -161,13 +162,14 @@ Socket.prototype.maybeUpgrade = function (transport) {
   debug('might upgrade socket transport from "%s" to "%s"'
     , this.transport.name, transport.name);
 
+  this.upgrading = true;
+
   var self = this;
 
   // set transport upgrade timer
   self.upgradeTimeoutTimer = setTimeout(function () {
     debug('client did not complete upgrade - closing transport');
-    clearInterval(self.checkIntervalTimer);
-    self.checkIntervalTimer = null;
+    cleanup();
     if ('open' == transport.readyState) {
       transport.close();
     }
@@ -180,22 +182,20 @@ Socket.prototype.maybeUpgrade = function (transport) {
       self.checkIntervalTimer = setInterval(check, 100);
     } else if ('upgrade' == packet.type && self.readyState != 'closed') {
       debug('got upgrade packet - upgrading');
+      cleanup();
       self.upgraded = true;
       self.clearTransport();
       self.setTransport(transport);
       self.emit('upgrade', transport);
       self.setPingTimeout();
       self.flush();
-      clearInterval(self.checkIntervalTimer);
-      self.checkIntervalTimer = null;
-      clearTimeout(self.upgradeTimeoutTimer);
-      transport.removeListener('packet', onPacket);
       if (self.readyState == 'closing') {
         transport.close(function () {
           self.onClose('forced close');
         });
       }
     } else {
+      cleanup();
       transport.close();
     }
   }
@@ -208,7 +208,41 @@ Socket.prototype.maybeUpgrade = function (transport) {
     }
   }
 
+  function cleanup() {
+    self.upgrading = false;
+
+    clearInterval(self.checkIntervalTimer);
+    self.checkIntervalTimer = null;
+
+    clearTimeout(self.upgradeTimeoutTimer);
+    self.upgradeTimeoutTimer = null;
+
+    transport.removeListener('packet', onPacket);
+    transport.removeListener('close', onTransportClose);
+    transport.removeListener('error', onError);
+    self.removeListener('close', onClose);
+  }
+
+  function onError(err) {
+    debug('client did not complete upgrade - %s', err);
+    cleanup();
+    transport.close();
+    transport = null;
+  }
+
+  function onTransportClose(){
+    onError("transport closed");
+  }
+
+  function onClose() {
+    onError("socket closed");
+  }
+
   transport.on('packet', onPacket);
+  transport.once('close', onTransportClose);
+  transport.once('error', onError);
+
+  self.once('close', onClose);
 };
 
 /**