From afb132cd11a627beb9a8634fec4c5499fd6f41e0 Mon Sep 17 00:00:00 2001
From: Peter Hedenskog <peter@soulgalore.com>
Date: Mon, 15 Jul 2024 05:58:07 +0200
Subject: [PATCH] Set headers using Bidi (Firefox 128) (#2108)

---
 lib/firefox/firefoxBidi.js       | 49 ++++++++++++++++++++++++++++++++
 lib/firefox/webdriver/firefox.js | 33 +++++++++++----------
 2 files changed, 67 insertions(+), 15 deletions(-)

diff --git a/lib/firefox/firefoxBidi.js b/lib/firefox/firefoxBidi.js
index 5800e8d61..89c3c4fcd 100644
--- a/lib/firefox/firefoxBidi.js
+++ b/lib/firefox/firefoxBidi.js
@@ -78,6 +78,7 @@ export class FirefoxBidi {
       }
     });
   }
+
   async setBasicAuth(basicAuth) {
     const parts = basicAuth.split('@');
     const bidi = new Bidi(this.driver, this.options.browser);
@@ -146,4 +147,52 @@ export class FirefoxBidi {
       }
     }
   }
+
+  async setRequestHeaders(requestHeaders) {
+    const headersArray = toArray(requestHeaders);
+    const headers = [];
+    for (let header of headersArray) {
+      if (header.indexOf && header.includes(':')) {
+        const parts = header.split(':');
+        headers.push({
+          name: parts[0],
+          value: {
+            type: 'string',
+            value: parts[1]
+          }
+        });
+      } else {
+        log.error(
+          'Request headers need to be of the format key:value not ' + header
+        );
+      }
+    }
+
+    const command = {
+      method: 'network.addIntercept',
+      params: {
+        phases: ['beforeRequestSent']
+      }
+    };
+    const bidi = new Bidi(this.driver, this.options.browser);
+
+    await bidi.send(command);
+    await bidi.subscribe('network.beforeRequestSent');
+    await bidi.onMessage(async function (event) {
+      const parsedEvent = JSON.parse(Buffer.from(event.toString()));
+      if (parsedEvent.method === 'network.beforeRequestSent') {
+        const continueRequest = {
+          method: 'network.continueRequest',
+          params: {
+            request: parsedEvent.params.request.request,
+            headers: [...parsedEvent.params.request.headers, ...headers]
+          }
+        };
+        const result = await bidi.send(continueRequest);
+        if (result.type != 'success') {
+          log.error(result);
+        }
+      }
+    });
+  }
 }
diff --git a/lib/firefox/webdriver/firefox.js b/lib/firefox/webdriver/firefox.js
index 03682ae7d..bb58bb83c 100644
--- a/lib/firefox/webdriver/firefox.js
+++ b/lib/firefox/webdriver/firefox.js
@@ -102,22 +102,25 @@ export class Firefox {
       await this.browsertimeBidi.setBasicAuth(this.options.basicAuth);
     }
 
-    if (this.options.block) {
-      await this.browsertimeBidi.blockUrls(this.options.block);
-    }
+    if (this.options.requestheader) {
+      await this.browsertimeBidi.setRequestHeaders(this.options.requestheader);
+      if (this.options.block) {
+        await this.browsertimeBidi.blockUrls(this.options.block);
+      }
 
-    if (
-      this.firefoxConfig.appendToUserAgent ||
-      this.options.appendToUserAgent
-    ) {
-      const currentUserAgent = await runner.runScript(
-        'return navigator.userAgent;',
-        'GET_USER_AGENT'
-      );
-      let script = `Services.prefs.setStringPref('general.useragent.override', '${currentUserAgent} ${
-        this.firefoxConfig.appendToUserAgent || this.options.appendToUserAgent
-      }');`;
-      return runner.runPrivilegedScript(script, 'SET_USER_AGENT');
+      if (
+        this.firefoxConfig.appendToUserAgent ||
+        this.options.appendToUserAgent
+      ) {
+        const currentUserAgent = await runner.runScript(
+          'return navigator.userAgent;',
+          'GET_USER_AGENT'
+        );
+        let script = `Services.prefs.setStringPref('general.useragent.override', '${currentUserAgent} ${
+          this.firefoxConfig.appendToUserAgent || this.options.appendToUserAgent
+        }');`;
+        return runner.runPrivilegedScript(script, 'SET_USER_AGENT');
+      }
     }
   }