From 5a5873048005cf5d25a2186fb9dc6db2a85096b0 Mon Sep 17 00:00:00 2001
From: Luigi Pinca <luigipinca@gmail.com>
Date: Fri, 9 Jul 2021 11:06:26 +0200
Subject: [PATCH] [fix] Emit the `'close'` event after the server is closed

Ensure that `WebSocketServer.prototype.close()` does not emit a
`'close'` event prematurely if called while the internal HTTP/S server
is closing.
---
 lib/websocket-server.js       | 16 ++++++++++++-
 test/websocket-server.test.js | 42 +++++++++++++++++++++++++++++++----
 2 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/lib/websocket-server.js b/lib/websocket-server.js
index 20276edb3..e9ec2cbd7 100644
--- a/lib/websocket-server.js
+++ b/lib/websocket-server.js
@@ -16,6 +16,10 @@ const { GUID, kWebSocket } = require('./constants');
 
 const keyRegex = /^[+/0-9A-Za-z]{22}==$/;
 
+const RUNNING = 0;
+const CLOSING = 1;
+const CLOSED = 2;
+
 /**
  * Class representing a WebSocket server.
  *
@@ -108,6 +112,7 @@ class WebSocketServer extends EventEmitter {
     if (options.perMessageDeflate === true) options.perMessageDeflate = {};
     if (options.clientTracking) this.clients = new Set();
     this.options = options;
+    this._state = RUNNING;
   }
 
   /**
@@ -137,6 +142,14 @@ class WebSocketServer extends EventEmitter {
   close(cb) {
     if (cb) this.once('close', cb);
 
+    if (this._state === CLOSED) {
+      process.nextTick(emitClose, this);
+      return;
+    }
+
+    if (this._state === CLOSING) return;
+    this._state = CLOSING;
+
     //
     // Terminate all associated clients.
     //
@@ -154,7 +167,7 @@ class WebSocketServer extends EventEmitter {
       // Close the http server if it was internally created.
       //
       if (this.options.port != null) {
-        server.close(() => this.emit('close'));
+        server.close(emitClose.bind(undefined, this));
         return;
       }
     }
@@ -373,6 +386,7 @@ function addListeners(server, map) {
  * @private
  */
 function emitClose(server) {
+  server._state = CLOSED;
   server.emit('close');
 }
 
diff --git a/test/websocket-server.test.js b/test/websocket-server.test.js
index 9b43284a1..310c87c2e 100644
--- a/test/websocket-server.test.js
+++ b/test/websocket-server.test.js
@@ -278,11 +278,45 @@ describe('WebSocketServer', () => {
       });
     });
 
-    it("emits the 'close' event", (done) => {
-      const wss = new WebSocket.Server({ noServer: true });
+    it("emits the 'close' event after the server closes", (done) => {
+      let serverCloseEventEmitted = false;
+
+      const wss = new WebSocket.Server({ port: 0 }, () => {
+        net.createConnection({ port: wss.address().port });
+      });
+
+      wss._server.on('connection', (socket) => {
+        wss.close();
+
+        //
+        // The server is closing. Ensure this does not emit a `'close'`
+        // event before the server is actually closed.
+        //
+        wss.close();
+
+        process.nextTick(() => {
+          socket.end();
+        });
+      });
 
-      wss.on('close', done);
-      wss.close();
+      wss._server.on('close', () => {
+        serverCloseEventEmitted = true;
+      });
+
+      wss.on('close', () => {
+        assert.ok(serverCloseEventEmitted);
+        done();
+      });
+    });
+
+    it("emits the 'close' event if the server is already closed", (done) => {
+      const wss = new WebSocket.Server({ port: 0 }, () => {
+        wss.close(() => {
+          assert.strictEqual(wss._state, 2);
+          wss.on('close', done);
+          wss.close();
+        });
+      });
     });
   });