fix: properly close the websocket connection upon handshake error

This should fix the following error:

> TypeError: Cannot read property 'writeHead' of undefined

This bug was introduced by [1], because the `if (res !== undefined) { ... }`
check was removed.

But `res` is indeed undefined when the client connects with WebSocket
directly, in that case we need to manually write the response content
(in the abortUpgrade method).

Please note that the previous behavior was invalid too, since the
WebSocket connection was left open when an error occurred during the
handshake.

[1]: 7096e98

Author: darrachequesne

lib/server.js

@@ -236,15 +236,17 @@ class Server extends EventEmitter {
           message: Server.errorMessages[errorCode],
           context: errorContext
         });
-        sendErrorMessage(req, res, errorCode, errorContext);
+        abortRequest(res, errorCode, errorContext);
         return;
       }
 
       if (req._query.sid) {
         debug("setting new request for existing client");
         this.clients[req._query.sid].transport.onRequest(req);
       } else {
-        this.handshake(req._query.transport, req);
+        const closeConnection = (errorCode, errorContext) =>
+          abortRequest(res, errorCode, errorContext);
+        this.handshake(req._query.transport, req, closeConnection);
       }
     };
 
@@ -273,9 +275,11 @@ class Server extends EventEmitter {
    *
    * @param {String} transport name
    * @param {Object} request object
+   * @param {Function} closeConnection
+   *
    * @api private
    */
-  async handshake(transportName, req) {
+  async handshake(transportName, req, closeConnection) {
     const protocol = req._query.EIO === "4" ? 4 : 3; // 3rd revision by default
     if (protocol === 3 && !this.opts.allowEIO3) {
       debug("unsupported protocol version");
@@ -288,11 +292,7 @@ class Server extends EventEmitter {
           protocol
         }
       });
-      sendErrorMessage(
-        req,
-        req.res,
-        Server.errors.UNSUPPORTED_PROTOCOL_VERSION
-      );
+      closeConnection(Server.errors.UNSUPPORTED_PROTOCOL_VERSION);
       return;
     }
 
@@ -310,7 +310,7 @@ class Server extends EventEmitter {
           error: e
         }
       });
-      sendErrorMessage(req, req.res, Server.errors.BAD_REQUEST);
+      closeConnection(Server.errors.BAD_REQUEST);
       return;
     }
 
@@ -341,7 +341,7 @@ class Server extends EventEmitter {
           error: e
         }
       });
-      sendErrorMessage(req, req.res, Server.errors.BAD_REQUEST);
+      closeConnection(Server.errors.BAD_REQUEST);
       return;
     }
     const socket = new Socket(id, this, transport, req, protocol);
@@ -389,16 +389,16 @@ class Server extends EventEmitter {
           message: Server.errorMessages[errorCode],
           context: errorContext
         });
-        abortConnection(socket, errorCode, errorContext);
+        abortUpgrade(socket, errorCode, errorContext);
         return;
       }
 
       const head = Buffer.from(upgradeHead); // eslint-disable-line node/no-deprecated-api
       upgradeHead = null;
 
       // delegate to ws
-      this.ws.handleUpgrade(req, socket, head, conn => {
-        this.onWebSocket(req, conn);
+      this.ws.handleUpgrade(req, socket, head, websocket => {
+        this.onWebSocket(req, socket, websocket);
       });
     });
   }
@@ -409,40 +409,40 @@ class Server extends EventEmitter {
    * @param {ws.Socket} websocket
    * @api private
    */
-  onWebSocket(req, socket) {
-    socket.on("error", onUpgradeError);
+  onWebSocket(req, socket, websocket) {
+    websocket.on("error", onUpgradeError);
 
     if (
       transports[req._query.transport] !== undefined &&
       !transports[req._query.transport].prototype.handlesUpgrades
     ) {
       debug("transport doesnt handle upgraded requests");
-      socket.close();
+      websocket.close();
       return;
     }
 
     // get client id
     const id = req._query.sid;
 
     // keep a reference to the ws.Socket
-    req.websocket = socket;
+    req.websocket = websocket;
 
     if (id) {
       const client = this.clients[id];
       if (!client) {
         debug("upgrade attempt for closed client");
-        socket.close();
+        websocket.close();
       } else if (client.upgrading) {
         debug("transport has already been trying to upgrade");
-        socket.close();
+        websocket.close();
       } else if (client.upgraded) {
         debug("transport had already been upgraded");
-        socket.close();
+        websocket.close();
       } else {
         debug("upgrading existing transport");
 
         // transport error handling takes over
-        socket.removeListener("error", onUpgradeError);
+        websocket.removeListener("error", onUpgradeError);
 
         const transport = new transports[req._query.transport](req);
         if (req._query && req._query.b64) {
@@ -455,14 +455,16 @@ class Server extends EventEmitter {
       }
     } else {
       // transport error handling takes over
-      socket.removeListener("error", onUpgradeError);
+      websocket.removeListener("error", onUpgradeError);
 
-      this.handshake(req._query.transport, req);
+      const closeConnection = (errorCode, errorContext) =>
+        abortUpgrade(socket, errorCode, errorContext);
+      this.handshake(req._query.transport, req, closeConnection);
     }
 
     function onUpgradeError() {
       debug("websocket error before upgrade");
-      // socket.close() not needed
+      // websocket.close() not needed
     }
   }
 
@@ -548,17 +550,16 @@ Server.errorMessages = {
 };
 
 /**
- * Sends an Engine.IO Error Message
+ * Close the HTTP long-polling request
  *
- * @param req - the request object
  * @param res - the response object
  * @param errorCode - the error code
  * @param errorContext - additional error context
  *
  * @api private
  */
 
-function sendErrorMessage(req, res, errorCode, errorContext) {
+function abortRequest(res, errorCode, errorContext) {
   const statusCode = errorCode === Server.errors.FORBIDDEN ? 403 : 400;
   const message =
     errorContext && errorContext.message
@@ -575,7 +576,7 @@ function sendErrorMessage(req, res, errorCode, errorContext) {
 }
 
 /**
- * Closes the connection
+ * Close the WebSocket connection
  *
  * @param {net.Socket} socket
  * @param {string} errorCode - the error code
@@ -584,7 +585,7 @@ function sendErrorMessage(req, res, errorCode, errorContext) {
  * @api private
  */
 
-function abortConnection(socket, errorCode, errorContext) {
+function abortUpgrade(socket, errorCode, errorContext = {}) {
   socket.on("error", () => {
     debug("ignoring error from closed connection");
   });

test/server.js

@@ -374,6 +374,34 @@ describe("server", () => {
       });
     });
 
+    it("should disallow connection that are rejected by `generateId` (websocket only)", function(done) {
+      if (process.env.EIO_WS_ENGINE === "eiows") {
+        this.skip();
+      }
+      const partialDone = createPartialDone(done, 2);
+
+      engine = listen({ allowUpgrades: false }, port => {
+        engine.generateId = () => {
+          return Promise.reject(new Error("nope"));
+        };
+
+        engine.on("connection_error", err => {
+          expect(err.req).to.be.an(http.IncomingMessage);
+          expect(err.code).to.be(3);
+          expect(err.message).to.be("Bad request");
+          expect(err.context.name).to.be("ID_GENERATION_ERROR");
+          partialDone();
+        });
+
+        const socket = new eioc.Socket("ws://localhost:%d".s(port), {
+          transports: ["websocket"]
+        });
+        socket.on("error", () => {
+          partialDone();
+        });
+      });
+    });
+
     it("should exchange handshake data", done => {
       listen({ allowUpgrades: false }, port => {
         const socket = new eioc.Socket("ws://localhost:%d".s(port));