feat(webdriver): support page.setUserAgent for WebDriver BiDi (#12330)

puppeteer · Apr 25, 2024 · 1f99e66 · 1f99e66
1 parent 88b46ee
commit 1f99e66
Show file tree

Hide file tree

Showing 4 changed files with 106 additions and 23 deletions.
diff --git a/packages/puppeteer-core/src/bidi/HTTPRequest.ts b/packages/puppeteer-core/src/bidi/HTTPRequest.ts
@@ -76,7 +76,7 @@ export class BidiHTTPRequest extends HTTPRequest {
 
     this.#frame.page().trustedEmitter.emit(PageEvent.Request, this);
 
-    if (Object.keys(this.#extraHTTPHeaders).length) {
+    if (this.#hasInternalHeaderOverwrite) {
       this.interception.handlers.push(async () => {
         await this.continue(
           {
@@ -112,10 +112,21 @@ export class BidiHTTPRequest extends HTTPRequest {
     throw new UnsupportedOperation();
   }
 
+  get #hasInternalHeaderOverwrite(): boolean {
+    return Boolean(
+      Object.keys(this.#extraHTTPHeaders).length ||
+        Object.keys(this.#userAgentHeaders).length
+    );
+  }
+
   get #extraHTTPHeaders(): Record<string, string> {
     return this.#frame?.page()._extraHTTPHeaders ?? {};
   }
 
+  get #userAgentHeaders(): Record<string, string> {
+    return this.#frame?.page()._userAgentHeaders ?? {};
+  }
+
   override headers(): Record<string, string> {
     const headers: Record<string, string> = {};
     for (const header of this.#request.headers) {
@@ -124,6 +135,7 @@ export class BidiHTTPRequest extends HTTPRequest {
     return {
       ...headers,
       ...this.#extraHTTPHeaders,
+      ...this.#userAgentHeaders,
     };
   }
 
@@ -169,9 +181,7 @@ export class BidiHTTPRequest extends HTTPRequest {
   ): Promise<void> {
     return await super.continue(
       {
-        headers: Object.keys(this.#extraHTTPHeaders).length
-          ? this.headers()
-          : undefined,
+        headers: this.#hasInternalHeaderOverwrite ? this.headers() : undefined,
         ...overrides,
       },
       priority

diff --git a/packages/puppeteer-core/src/bidi/Page.ts b/packages/puppeteer-core/src/bidi/Page.ts
@@ -125,16 +125,71 @@ export class BidiPage extends Page {
       this.#workers.delete(worker as BidiWebWorker);
     });
   }
-
+  /**
+   * @internal
+   */
+  _userAgentHeaders: Record<string, string> = {};
+  #userAgentInterception?: string;
+  #userAgentPreloadScript?: string;
   override async setUserAgent(
     userAgent: string,
-    userAgentMetadata?: Protocol.Emulation.UserAgentMetadata | undefined
+    userAgentMetadata?: Protocol.Emulation.UserAgentMetadata
   ): Promise<void> {
-    // TODO: handle CDP-specific cases such as mprach.
-    await this._client().send('Network.setUserAgentOverride', {
-      userAgent: userAgent,
-      userAgentMetadata: userAgentMetadata,
-    });
+    if (!this.#browserContext.browser().cdpSupported && userAgentMetadata) {
+      throw new UnsupportedOperation(
+        'Current Browser does not support `userAgentMetadata`'
+      );
+    } else if (
+      this.#browserContext.browser().cdpSupported &&
+      userAgentMetadata
+    ) {
+      return await this._client().send('Network.setUserAgentOverride', {
+        userAgent: userAgent,
+        userAgentMetadata: userAgentMetadata,
+      });
+    }
+    const enable = userAgent !== '';
+    userAgent = userAgent ?? (await this.#browserContext.browser().userAgent());
+
+    this._userAgentHeaders = enable
+      ? {
+          'User-Agent': userAgent,
+        }
+      : {};
+
+    this.#userAgentInterception = await this.#toggleInterception(
+      [Bidi.Network.InterceptPhase.BeforeRequestSent],
+      this.#userAgentInterception,
+      enable
+    );
+
+    const changeUserAgent = (userAgent: string) => {
+      Object.defineProperty(navigator, 'userAgent', {
+        value: userAgent,
+      });
+    };
+
+    const frames = [this.#frame];
+    for (const frame of frames) {
+      frames.push(...frame.childFrames());
+    }
+
+    if (this.#userAgentPreloadScript) {
+      await this.removeScriptToEvaluateOnNewDocument(
+        this.#userAgentPreloadScript
+      );
+    }
+    const [evaluateToken] = await Promise.all([
+      enable
+        ? this.evaluateOnNewDocument(changeUserAgent, userAgent)
+        : undefined,
+      // When we disable the UserAgent we want to
+      // evaluate the original value in all Browsing Contexts
+      frames.map(frame => {
+        return frame.evaluate(changeUserAgent, userAgent);
+      }),
+    ]);
+    this.#userAgentPreloadScript = evaluateToken?.identifier;
   }
 
   override async setBypassCSP(enabled: boolean): Promise<void> {

diff --git a/test/TestExpectations.json b/test/TestExpectations.json
@@ -676,8 +676,8 @@
     "testIdPattern": "[page.spec] Page Page.setUserAgent *",
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["firefox", "webDriverBiDi"],
-    "expectations": ["FAIL"],
-    "comment": "TODO: add a comment explaining why this expectation is required (include links to issues)"
+    "expectations": ["SKIP"],
+    "comment": "Firefox does not support headers override"
   },
   {
     "testIdPattern": "[pdf.spec] Page.pdf *",
@@ -1424,15 +1424,8 @@
     "testIdPattern": "[emulation.spec] Emulation Page.emulate should support clicking",
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["firefox", "webDriverBiDi"],
-    "expectations": ["FAIL"],
-    "comment": "TODO: add a comment explaining why this expectation is required (include links to issues)"
-  },
-  {
-    "testIdPattern": "[emulation.spec] Emulation Page.emulate should work",
-    "platforms": ["darwin", "linux", "win32"],
-    "parameters": ["firefox", "webDriverBiDi"],
-    "expectations": ["FAIL"],
-    "comment": "TODO: add a comment explaining why this expectation is required (include links to issues)"
+    "expectations": ["SKIP"],
+    "comment": "Firefox does not support headers override"
   },
   {
     "testIdPattern": "[emulation.spec] Emulation Page.emulateCPUThrottling should change the CPU throttling rate successfully",
@@ -3074,7 +3067,7 @@
     "platforms": ["darwin", "linux", "win32"],
     "parameters": ["firefox", "webDriverBiDi"],
     "expectations": ["FAIL"],
-    "comment": "TODO: add a comment explaining why this expectation is required (include links to issues)"
+    "comment": "https://bugzilla.mozilla.org/show_bug.cgi?id=1866749"
   },
   {
     "testIdPattern": "[page.spec] Page Page.Events.error should throw when page crashes",

diff --git a/test/src/page.spec.ts b/test/src/page.spec.ts
@@ -1327,6 +1327,31 @@ describe('Page', function () {
       expect(uaData['platformVersion']).toBe('3.1');
       expect(request.headers['user-agent']).toBe('MockBrowser');
     });
+    it('should restore original', async () => {
+      const {page, server} = await getTestState();
+
+      const userAgent = await page.evaluate(() => {
+        return navigator.userAgent;
+      });
+
+      await page.setUserAgent('foobar');
+      const [requestWithOverride] = await Promise.all([
+        server.waitForRequest('/empty.html'),
+        page.goto(server.EMPTY_PAGE),
+      ]);
+      expect(requestWithOverride.headers['user-agent']).toBe('foobar');
+
+      await page.setUserAgent('');
+      const [request] = await Promise.all([
+        server.waitForRequest('/empty.html'),
+        page.goto(server.EMPTY_PAGE),
+      ]);
+      expect(request.headers['user-agent']).toBe(userAgent);
+      const userAgentRestored = await page.evaluate(() => {
+        return navigator.userAgent;
+      });
+      expect(userAgentRestored).toBe(userAgent);
+    });
   });
 
   describe('Page.setContent', function () {