From e1ddacce06d302884e674f50c12cb376e000c6ea Mon Sep 17 00:00:00 2001
From: Sergey Bakulin <sergey.bakulin.github@gmail.com>
Date: Tue, 1 Feb 2022 22:25:02 +0300
Subject: [PATCH] [feature] Introduce the `WebSocket` option (#2007)

Add the ability to use a custom class that extends the `WebSocket`
class.
---
 doc/ws.md                     |  2 ++
 lib/websocket-server.js       |  5 ++++-
 test/websocket-server.test.js | 26 ++++++++++++++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/doc/ws.md b/doc/ws.md
index 62fd36ca5..3273c1f9e 100644
--- a/doc/ws.md
+++ b/doc/ws.md
@@ -84,6 +84,8 @@ This class represents a WebSocket server. It extends the `EventEmitter`.
   - `verifyClient` {Function} A function which can be used to validate incoming
     connections. See description below. (Usage is discouraged: see
     [Issue #337](https://github.com/websockets/ws/issues/377#issuecomment-462152231))
+  - `WebSocket` {Function} Specifies the `WebSocket` class to be used. It must
+    be extended from the original `WebSocket`. Defaults to `WebSocket`.
 - `callback` {Function}
 
 Create a new server instance. One and only one of `port`, `server` or `noServer`
diff --git a/lib/websocket-server.js b/lib/websocket-server.js
index 3c7939f28..d0a29783f 100644
--- a/lib/websocket-server.js
+++ b/lib/websocket-server.js
@@ -49,6 +49,8 @@ class WebSocketServer extends EventEmitter {
    * @param {Boolean} [options.skipUTF8Validation=false] Specifies whether or
    *     not to skip UTF-8 validation for text and close messages
    * @param {Function} [options.verifyClient] A hook to reject connections
+   * @param {Function} [options.WebSocket=WebSocket] Specifies the `WebSocket`
+   *     class to use. It must be the `WebSocket` class or class that extends it
    * @param {Function} [callback] A listener for the `listening` event
    */
   constructor(options, callback) {
@@ -67,6 +69,7 @@ class WebSocketServer extends EventEmitter {
       host: null,
       path: null,
       port: null,
+      WebSocket,
       ...options
     };
 
@@ -356,7 +359,7 @@ class WebSocketServer extends EventEmitter {
       `Sec-WebSocket-Accept: ${digest}`
     ];
 
-    const ws = new WebSocket(null);
+    const ws = new this.options.WebSocket(null);
 
     if (protocols.size) {
       //
diff --git a/test/websocket-server.test.js b/test/websocket-server.test.js
index e3daf7e0b..fd494059f 100644
--- a/test/websocket-server.test.js
+++ b/test/websocket-server.test.js
@@ -89,6 +89,32 @@ describe('WebSocketServer', () => {
           wss.close(done);
         });
       });
+
+      it('honors the `WebSocket` option', (done) => {
+        class CustomWebSocket extends WebSocket.WebSocket {
+          get foo() {
+            return 'foo';
+          }
+        }
+
+        const wss = new WebSocket.Server(
+          {
+            port: 0,
+            WebSocket: CustomWebSocket
+          },
+          () => {
+            const ws = new WebSocket(`ws://localhost:${wss.address().port}`);
+
+            ws.on('open', ws.close);
+          }
+        );
+
+        wss.on('connection', (ws) => {
+          assert.ok(ws instanceof CustomWebSocket);
+          assert.strictEqual(ws.foo, 'foo');
+          wss.close(done);
+        });
+      });
     });
 
     it('emits an error if http server bind fails', (done) => {