From 187cfdc328c7b59bed82cd2ee0d9b3097de19da7 Mon Sep 17 00:00:00 2001
From: Dmitriy Dudkevich <dudkevich@yandex-team.ru>
Date: Thu, 8 Jun 2023 20:41:36 +0300
Subject: [PATCH] feat: ability to pass additional headers to selenium (#23348)

---
 docs/src/selenium-grid.md                     | 20 ++++++++++
 .../src/server/chromium/chromium.ts           | 40 ++++++++++++++-----
 2 files changed, 49 insertions(+), 11 deletions(-)

diff --git a/docs/src/selenium-grid.md b/docs/src/selenium-grid.md
index 92e7d6f134..6e5f146198 100644
--- a/docs/src/selenium-grid.md
+++ b/docs/src/selenium-grid.md
@@ -64,6 +64,26 @@ SELENIUM_REMOTE_URL=http://<selenium-hub-ip>:4444 SELENIUM_REMOTE_CAPABILITIES="
 SELENIUM_REMOTE_URL=http://<selenium-hub-ip>:4444 SELENIUM_REMOTE_CAPABILITIES="{'mygrid:options':{os:'windows',username:'John',password:'secure'}}" dotnet test
 ```
 
+### Passing additional headers
+
+If your grid requires additional headers to be set (for example, you should provide authorization token to use browsers in your cloud), you can set `SELENIUM_REMOTE_HEADERS` environment variable to provide JSON-serialized headers.
+
+```bash js
+SELENIUM_REMOTE_URL=http://<selenium-hub-ip>:4444 SELENIUM_REMOTE_HEADERS="{'Authorization':'OAuth 12345'}" npx playwright test
+```
+
+```bash python
+SELENIUM_REMOTE_URL=http://<selenium-hub-ip>:4444 SELENIUM_REMOTE_HEADERS="{'Authorization':'OAuth 12345'}" pytest --browser chromium
+```
+
+```bash java
+SELENIUM_REMOTE_URL=http://<selenium-hub-ip>:4444 SELENIUM_REMOTE_HEADERS="{'Authorization':'OAuth 12345'}" mvn test
+```
+
+```bash csharp
+SELENIUM_REMOTE_URL=http://<selenium-hub-ip>:4444 SELENIUM_REMOTE_HEADERS="{'Authorization':'OAuth 12345'}" dotnet test
+```
+
 ### Detailed logs
 
 Run with `DEBUG=pw:browser*` environment variable to see how Playwright is connecting to Selenium Grid.
diff --git a/packages/playwright-core/src/server/chromium/chromium.ts b/packages/playwright-core/src/server/chromium/chromium.ts
index 9eeeea9217..80a8deaa86 100644
--- a/packages/playwright-core/src/server/chromium/chromium.ts
+++ b/packages/playwright-core/src/server/chromium/chromium.ts
@@ -36,7 +36,7 @@ import type { HTTPRequestParams } from '../../utils/network';
 import { fetchData } from '../../utils/network';
 import { getUserAgent } from '../../utils/userAgent';
 import { wrapInASCIIBox } from '../../utils/ascii';
-import { debugMode, headersArrayToObject, } from '../../utils';
+import { debugMode, headersArrayToObject, headersObjectToArray, } from '../../utils';
 import { removeFolders } from '../../utils/fileUtils';
 import { RecentLogsCollector } from '../../common/debugLogger';
 import type { Progress } from '../progress';
@@ -179,14 +179,18 @@ export class Chromium extends BrowserType {
       'browserName': isEdge ? 'MicrosoftEdge' : 'chrome',
       [isEdge ? 'ms:edgeOptions' : 'goog:chromeOptions']: { args }
     };
-    try {
-      if (process.env.SELENIUM_REMOTE_CAPABILITIES) {
-        const parsed = JSON.parse(process.env.SELENIUM_REMOTE_CAPABILITIES);
-        desiredCapabilities = { ...desiredCapabilities, ...parsed };
-        progress.log(`<selenium> using additional capabilities "${process.env.SELENIUM_REMOTE_CAPABILITIES}"`);
-      }
-    } catch (e) {
-      progress.log(`<selenium> ignoring additional capabilities "${process.env.SELENIUM_REMOTE_CAPABILITIES}": ${e}`);
+
+    if (process.env.SELENIUM_REMOTE_CAPABILITIES) {
+      const remoteCapabilities = parseSeleniumRemoteParams({ name: 'capabilities', value: process.env.SELENIUM_REMOTE_CAPABILITIES }, progress);
+      if (remoteCapabilities)
+        desiredCapabilities = { ...desiredCapabilities, ...remoteCapabilities };
+    }
+
+    let headers: { [key: string]: string } = {};
+    if (process.env.SELENIUM_REMOTE_HEADERS) {
+      const remoteHeaders = parseSeleniumRemoteParams({ name: 'headers', value: process.env.SELENIUM_REMOTE_HEADERS }, progress);
+      if (remoteHeaders)
+        headers = remoteHeaders;
     }
 
     progress.log(`<selenium> connecting to ${hubUrl}`);
@@ -194,7 +198,8 @@ export class Chromium extends BrowserType {
       url: hubUrl + 'session',
       method: 'POST',
       headers: {
-        'Content-Type': 'application/json; charset=utf-8'
+        'Content-Type': 'application/json; charset=utf-8',
+        ...headers,
       },
       data: JSON.stringify({
         desiredCapabilities,
@@ -257,7 +262,10 @@ export class Chromium extends BrowserType {
         }
       }
 
-      return await this._connectOverCDPInternal(progress, endpointURL.toString(), options, disconnectFromSelenium);
+      return await this._connectOverCDPInternal(progress, endpointURL.toString(), {
+        ...options,
+        headers: headersObjectToArray(headers),
+      }, disconnectFromSelenium);
     } catch (e) {
       await disconnectFromSelenium();
       throw e;
@@ -376,3 +384,13 @@ function streamToString(stream: stream.Readable): Promise<string> {
     stream.on('end', () => resolve(Buffer.concat(chunks).toString('utf8')));
   });
 }
+
+function parseSeleniumRemoteParams(env: {name: string, value: string}, progress: Progress) {
+  try {
+    const parsed = JSON.parse(env.value);
+    progress.log(`<selenium> using additional ${env.name} "${env.value}"`);
+    return parsed;
+  } catch (e) {
+    progress.log(`<selenium> ignoring additional ${env.name} "${env.value}": ${e}`);
+  }
+}