From cd89e077f68ba9a999d408cb4fdb3e91289096a7 Mon Sep 17 00:00:00 2001
From: Matthijs van Duin <matthijsvanduin@gmail.com>
Date: Fri, 10 Mar 2023 15:16:35 +0100
Subject: [PATCH] [feature] Add option to support late addition of headers
 (#2123)

This supports the use-case where headers need to be added that depend on
the socket connection (e.g. for TLS channel binding).
---
 doc/ws.md              | 16 +++++++++++++++-
 lib/websocket.js       |  6 +++++-
 test/websocket.test.js | 29 +++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/doc/ws.md b/doc/ws.md
index ae7993e68..0fc44d6e6 100644
--- a/doc/ws.md
+++ b/doc/ws.md
@@ -293,6 +293,8 @@ This class represents a WebSocket. It extends the `EventEmitter`.
 - `address` {String|url.URL} The URL to which to connect.
 - `protocols` {String|Array} The list of subprotocols.
 - `options` {Object}
+  - `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
     `false`.
   - `generateMask` {Function} The function used to generate the masking key. It
@@ -316,12 +318,24 @@ This class represents a WebSocket. It extends the `EventEmitter`.
     Options given do not have any effect if parsed from the URL given with the
     `address` parameter.
 
+Create a new WebSocket instance.
+
 `perMessageDeflate` default value is `true`. When using an object, parameters
 are the same of the server. The only difference is the direction of requests.
 For example, `serverNoContextTakeover` can be used to ask the server to disable
 context takeover.
 
-Create a new WebSocket instance.
+`finishRequest` is called with arguments
+
+- `request` {http.ClientRequest}
+- `websocket` {WebSocket}
+
+for each HTTP GET request (the initial one and any caused by redirects) when it
+is ready to be sent, to allow for last minute customization of the headers. If
+`finishRequest` is set then it has the responsibility to call `request.end()`
+once it is done setting request headers. This is intended for niche use-cases
+where some headers can't be provided in advance e.g. because they depend on the
+underlying socket.
 
 #### IPC connections
 
diff --git a/lib/websocket.js b/lib/websocket.js
index 35a788ac4..b2b2b0926 100644
--- a/lib/websocket.js
+++ b/lib/websocket.js
@@ -989,7 +989,11 @@ function initAsClient(websocket, address, protocols, options) {
     });
   });
 
-  req.end();
+  if (opts.finishRequest) {
+    opts.finishRequest(req, websocket);
+  } else {
+    req.end();
+  }
 }
 
 /**
diff --git a/test/websocket.test.js b/test/websocket.test.js
index 6b2f3ef5c..f80acd3d5 100644
--- a/test/websocket.test.js
+++ b/test/websocket.test.js
@@ -3857,6 +3857,35 @@ describe('WebSocket', () => {
         agent
       });
     });
+
+    it('honors the `finishRequest` option', (done) => {
+      const wss = new WebSocket.Server({ port: 0 }, () => {
+        const ws = new WebSocket(`ws://localhost:${wss.address().port}`, {
+          finishRequest(request, websocket) {
+            process.nextTick(() => {
+              assert.strictEqual(request, ws._req);
+              assert.strictEqual(websocket, ws);
+            });
+            request.on('socket', (socket) => {
+              socket.on('connect', () => {
+                request.setHeader('Cookie', 'foo=bar');
+                request.end();
+              });
+            });
+          }
+        });
+
+        ws.on('close', (code) => {
+          assert.strictEqual(code, 1005);
+          wss.close(done);
+        });
+      });
+
+      wss.on('connection', (ws, req) => {
+        assert.strictEqual(req.headers.cookie, 'foo=bar');
+        ws.close();
+      });
+    });
   });
 
   describe('permessage-deflate', () => {