refactor: extract toggling of Network Interception (#12266)

puppeteer · Apr 12, 2024 · 402b4a4 · 402b4a4
1 parent df986ef
commit 402b4a4
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 25 deletions.
diff --git a/packages/puppeteer-core/src/bidi/Page.ts b/packages/puppeteer-core/src/bidi/Page.ts
@@ -511,18 +511,13 @@ export class BidiPage extends Page {
     return [...this.#workers];
   }
 
-  #interception?: string;
+  #userInterception?: string;
   override async setRequestInterception(enable: boolean): Promise<void> {
-    if (enable && !this.#interception) {
-      this.#interception = await this.#frame.browsingContext.addIntercept({
-        phases: [Bidi.Network.InterceptPhase.BeforeRequestSent],
-      });
-    } else if (!enable && this.#interception) {
-      await this.#frame.browsingContext.userContext.browser.removeIntercept(
-        this.#interception
-      );
-      this.#interception = undefined;
-    }
+    this.#userInterception = await this.#toggleInterception(
+      [Bidi.Network.InterceptPhase.BeforeRequestSent],
+      this.#userInterception,
+      enable
+    );
   }
 
   /**
@@ -531,17 +526,31 @@ export class BidiPage extends Page {
   _credentials: Credentials | null = null;
   #authInterception?: string;
   override async authenticate(credentials: Credentials | null): Promise<void> {
-    if (credentials && !this.#authInterception) {
-      this.#authInterception = await this.#frame.browsingContext.addIntercept({
-        phases: [Bidi.Network.InterceptPhase.AuthRequired],
+    this.#authInterception = await this.#toggleInterception(
+      [Bidi.Network.InterceptPhase.AuthRequired],
+      this.#authInterception,
+      Boolean(credentials)
+    );
+
+    this._credentials = credentials;
+  }
+
+  async #toggleInterception(
+    phases: [Bidi.Network.InterceptPhase, ...Bidi.Network.InterceptPhase[]],
+    interception: string | undefined,
+    expected: boolean
+  ): Promise<string | undefined> {
+    if (expected && !interception) {
+      return await this.#frame.browsingContext.addIntercept({
+        phases,
       });
-    } else if (!credentials && this.#authInterception) {
+    } else if (!expected && interception) {
       await this.#frame.browsingContext.userContext.browser.removeIntercept(
-        this.#authInterception
+        interception
       );
-      this.#authInterception = undefined;
+      return;
     }
-    this._credentials = credentials;
+    return interception;
   }
 
   override setDragInterception(): never {

diff --git a/packages/puppeteer-core/src/cdp/NetworkManager.ts b/packages/puppeteer-core/src/cdp/NetworkManager.ts
@@ -140,18 +140,16 @@ export class NetworkManager extends EventEmitter<NetworkManagerEvents> {
     );
   }
 
-  async setExtraHTTPHeaders(
-    extraHTTPHeaders: Record<string, string>
-  ): Promise<void> {
-    this.#extraHTTPHeaders = {};
-    for (const key of Object.keys(extraHTTPHeaders)) {
-      const value = extraHTTPHeaders[key];
+  async setExtraHTTPHeaders(headers: Record<string, string>): Promise<void> {
+    const extraHTTPHeaders: Record<string, string> = {};
+    for (const [key, value] of Object.entries(headers)) {
       assert(
         isString(value),
         `Expected value of header "${key}" to be String, but "${typeof value}" is found.`
       );
-      this.#extraHTTPHeaders[key.toLowerCase()] = value;
+      extraHTTPHeaders[key.toLowerCase()] = value;
     }
+    this.#extraHTTPHeaders = extraHTTPHeaders;
 
     await this.#applyToAllClients(this.#applyExtraHTTPHeaders.bind(this));
   }