From 53a88881cf5da8307ecbd5020db0a8fb72cf0d20 Mon Sep 17 00:00:00 2001
From: Tim Perry <1526883+pimterry@users.noreply.github.com>
Date: Mon, 15 Apr 2024 14:20:52 +0200
Subject: [PATCH] [feature] Allow the `createConnection` option (#2219)

Allow passing in a custom `createConnection` function.
---
 doc/ws.md              |  4 ++++
 lib/websocket.js       |  7 +++++--
 test/websocket.test.js | 28 ++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/doc/ws.md b/doc/ws.md
index 017087f5f..37f4c9707 100644
--- a/doc/ws.md
+++ b/doc/ws.md
@@ -304,6 +304,10 @@ This class represents a WebSocket. It extends the `EventEmitter`.
     `'ping'`, and `'pong'` events can be emitted multiple times in the same
     tick. To improve compatibility with the WHATWG standard, the default value
     is `false`. Setting it to `true` improves performance slightly.
+  - `createConnection` {Function} An alternative function to use in place of
+    `tls.createConnection` or `net.createConnection`. This can be used to
+    manually control exactly how the connection to the server is made, or to
+    make a connection over an existing Duplex stream obtained elsewhere.
   - `finishRequest` {Function} A function which can be used to customize the
     headers of each HTTP request before it is sent. See description below.
   - `followRedirects` {Boolean} Whether or not to follow redirects. Defaults to
diff --git a/lib/websocket.js b/lib/websocket.js
index f133d08fc..a2c8edbec 100644
--- a/lib/websocket.js
+++ b/lib/websocket.js
@@ -628,6 +628,8 @@ module.exports = WebSocket;
  *     times in the same tick
  * @param {Boolean} [options.autoPong=true] Specifies whether or not to
  *     automatically send a pong in response to a ping
+ * @param {Function} [options.createConnection] An alternative function to use
+ *     in place of `tls.createConnection` or `net.createConnection`.
  * @param {Function} [options.finishRequest] A function which can be used to
  *     customize the headers of each http request before it is sent
  * @param {Boolean} [options.followRedirects=false] Whether or not to follow
@@ -660,8 +662,8 @@ function initAsClient(websocket, address, protocols, options) {
     perMessageDeflate: true,
     followRedirects: false,
     maxRedirects: 10,
-    ...options,
     createConnection: undefined,
+    ...options,
     socketPath: undefined,
     hostname: undefined,
     protocol: undefined,
@@ -732,7 +734,8 @@ function initAsClient(websocket, address, protocols, options) {
   const protocolSet = new Set();
   let perMessageDeflate;
 
-  opts.createConnection = isSecure ? tlsConnect : netConnect;
+  opts.createConnection =
+    opts.createConnection || (isSecure ? tlsConnect : netConnect);
   opts.defaultPort = opts.defaultPort || defaultPort;
   opts.port = parsedUrl.port || defaultPort;
   opts.host = parsedUrl.hostname.startsWith('[')
diff --git a/test/websocket.test.js b/test/websocket.test.js
index e1b3bd239..fc41ae755 100644
--- a/test/websocket.test.js
+++ b/test/websocket.test.js
@@ -1224,6 +1224,34 @@ describe('WebSocket', () => {
       });
     });
 
+    it('honors the `createConnection` option', (done) => {
+      const wss = new WebSocket.Server({ noServer: true, path: '/foo' });
+
+      server.once('upgrade', (req, socket, head) => {
+        assert.strictEqual(req.headers.host, 'google.com:22');
+        wss.handleUpgrade(req, socket, head, NOOP);
+      });
+
+      const ws = new WebSocket('ws://google.com:22/foo', {
+        createConnection: (options) => {
+          assert.strictEqual(options.host, 'google.com');
+          assert.strictEqual(options.port, '22');
+
+          // Ignore the invalid host address, and connect to the server manually:
+          return net.createConnection({
+            host: 'localhost',
+            port: server.address().port
+          });
+        }
+      });
+
+      ws.on('open', () => {
+        assert.strictEqual(ws.url, 'ws://google.com:22/foo');
+        ws.on('close', () => done());
+        ws.close();
+      });
+    });
+
     it('emits an error if the redirect URL is invalid (1/2)', (done) => {
       server.once('upgrade', (req, socket) => {
         socket.end('HTTP/1.1 302 Found\r\nLocation: ws://\r\n\r\n');