refactor: move binding installation into the execution context (#12398)

packages/puppeteer-core/src/cdp/ExecutionContext.ts

@@ -11,29 +11,36 @@ import type {ElementHandle} from '../api/ElementHandle.js';
 import type {JSHandle} from '../api/JSHandle.js';
 import {LazyArg} from '../common/LazyArg.js';
 import {scriptInjector} from '../common/ScriptInjector.js';
-import type {EvaluateFunc, HandleFor} from '../common/types.js';
+import type {BindingPayload, EvaluateFunc, HandleFor} from '../common/types.js';
 import {
   PuppeteerURL,
   SOURCE_URL_REGEX,
+  debugError,
   getSourcePuppeteerURLIfAvailable,
   getSourceUrlComment,
   isString,
 } from '../common/util.js';
 import type PuppeteerUtil from '../injected/injected.js';
 import {AsyncIterableUtil} from '../util/AsyncIterableUtil.js';
+import {disposeSymbol} from '../util/disposable.js';
 import {stringifyFunction} from '../util/Function.js';
+import {Mutex} from '../util/Mutex.js';
 
 import {ARIAQueryHandler} from './AriaQueryHandler.js';
 import {Binding} from './Binding.js';
 import {CdpElementHandle} from './ElementHandle.js';
 import type {IsolatedWorld} from './IsolatedWorld.js';
 import {CdpJSHandle} from './JSHandle.js';
-import {createEvaluationError, valueFromRemoteObject} from './utils.js';
+import {
+  addPageBinding,
+  createEvaluationError,
+  valueFromRemoteObject,
+} from './utils.js';
 
 /**
  * @internal
  */
-export class ExecutionContext {
+export class ExecutionContext implements Disposable {
   _client: CDPSession;
   _world: IsolatedWorld;
   _contextId: number;
@@ -50,8 +57,91 @@ export class ExecutionContext {
     if (contextPayload.name) {
       this._contextName = contextPayload.name;
     }
+    this._client.on('Runtime.bindingCalled', this.#onBindingCalled);
   }
 
+  // Set of bindings that have been registered in the current context.
+  #contextBindings = new Set<string>();
+
+  // Contains mapping from functions that should be bound to Puppeteer functions.
+  #bindings = new Map<string, Binding>();
+
+  // If multiple waitFor are set up asynchronously, we need to wait for the
+  // first one to set up the binding in the page before running the others.
+  #mutex = new Mutex();
+  async _addBindingToContext(name: string): Promise<void> {
+    if (this.#contextBindings.has(name)) {
+      return;
+    }
+
+    using _ = await this.#mutex.acquire();
+    try {
+      await this._client.send(
+        'Runtime.addBinding',
+        this._contextName
+          ? {
+              name,
+              executionContextName: this._contextName,
+            }
+          : {
+              name,
+              executionContextId: this._contextId,
+            }
+      );
+
+      await this.evaluate(addPageBinding, 'internal', name);
+
+      this.#contextBindings.add(name);
+    } catch (error) {
+      // We could have tried to evaluate in a context which was already
+      // destroyed. This happens, for example, if the page is navigated while
+      // we are trying to add the binding
+      if (error instanceof Error) {
+        // Destroyed context.
+        if (error.message.includes('Execution context was destroyed')) {
+          return;
+        }
+        // Missing context.
+        if (error.message.includes('Cannot find context with specified id')) {
+          return;
+        }
+      }
+
+      debugError(error);
+    }
+  }
+
+  #onBindingCalled = async (
+    event: Protocol.Runtime.BindingCalledEvent
+  ): Promise<void> => {
+    let payload: BindingPayload;
+    try {
+      payload = JSON.parse(event.payload);
+    } catch {
+      // The binding was either called by something in the page or it was
+      // called before our wrapper was initialized.
+      return;
+    }
+    const {type, name, seq, args, isTrivial} = payload;
+    if (type !== 'internal') {
+      return;
+    }
+    if (!this.#contextBindings.has(name)) {
+      return;
+    }
+
+    try {
+      if (event.executionContextId !== this._contextId) {
+        return;
+      }
+
+      const binding = this.#bindings.get(name);
+      await binding?.run(this, seq, args, isTrivial);
+    } catch (err) {
+      debugError(err);
+    }
+  };
+
   #bindingsInstalled = false;
   #puppeteerUtil?: Promise<JSHandle<PuppeteerUtil>>;
   get puppeteerUtil(): Promise<JSHandle<PuppeteerUtil>> {
@@ -97,8 +187,8 @@ export class ExecutionContext {
   async #installGlobalBinding(binding: Binding) {
     try {
       if (this._world) {
-        this._world._bindings.set(binding.name, binding);
+        this.#bindings.set(binding.name, binding);
-        await this._world._addBindingToContext(this, binding.name);
+        await this._addBindingToContext(binding.name);
       }
     } catch {
       // If the binding cannot be added, then either the browser doesn't support
@@ -357,6 +447,10 @@ export class ExecutionContext {
       return {value: arg};
     }
   }
+
+  [disposeSymbol](): void {
+    this._client.off('Runtime.bindingCalled', this.#onBindingCalled);
+  }
 }
 
 const rewriteError = (error: Error): Protocol.Runtime.EvaluateResponse => {

packages/puppeteer-core/src/cdp/IsolatedWorld.ts

@@ -11,19 +11,16 @@ import type {ElementHandle} from '../api/ElementHandle.js';
 import type {JSHandle} from '../api/JSHandle.js';
 import {Realm} from '../api/Realm.js';
 import type {TimeoutSettings} from '../common/TimeoutSettings.js';
-import type {BindingPayload, EvaluateFunc, HandleFor} from '../common/types.js';
+import type {EvaluateFunc, HandleFor} from '../common/types.js';
-import {debugError, withSourcePuppeteerURLIfNone} from '../common/util.js';
+import {withSourcePuppeteerURLIfNone} from '../common/util.js';
 import {Deferred} from '../util/Deferred.js';
 import {disposeSymbol} from '../util/disposable.js';
-import {Mutex} from '../util/Mutex.js';
 
-import type {Binding} from './Binding.js';
 import {CdpElementHandle} from './ElementHandle.js';
 import {ExecutionContext} from './ExecutionContext.js';
 import type {CdpFrame} from './Frame.js';
 import type {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 import {CdpJSHandle} from './JSHandle.js';
-import {addPageBinding} from './utils.js';
 import type {CdpWebWorker} from './WebWorker.js';
 
 /**
@@ -49,16 +46,6 @@ export interface IsolatedWorldChart {
 export class IsolatedWorld extends Realm {
   #context = Deferred.create<ExecutionContext>();
 
-  // Set of bindings that have been registered in the current context.
-  #contextBindings = new Set<string>();
-
-  // Contains mapping from functions that should be bound to Puppeteer functions.
-  #bindings = new Map<string, Binding>();
-
-  get _bindings(): Map<string, Binding> {
-    return this.#bindings;
-  }
-
   readonly #frameOrWorker: CdpFrame | CdpWebWorker;
 
   constructor(
@@ -74,9 +61,7 @@ export class IsolatedWorld extends Realm {
     return this.#frameOrWorker;
   }
 
-  frameUpdated(): void {
+  frameUpdated(): void {}
-    this.client.on('Runtime.bindingCalled', this.#onBindingCalled);
-  }
 
   get client(): CDPSession {
     return this.#frameOrWorker.client;
@@ -92,7 +77,10 @@ export class IsolatedWorld extends Realm {
   }
 
   setContext(context: ExecutionContext): void {
-    this.#contextBindings.clear();
+    const existingContext = this.#context.value();
+    if (existingContext instanceof ExecutionContext) {
+      existingContext[disposeSymbol]();
+    }
     this.#context.resolve(context);
     void this.taskManager.rerunAll();
   }
@@ -146,86 +134,6 @@ export class IsolatedWorld extends Realm {
     return await context.evaluate(pageFunction, ...args);
   }
 
-  // If multiple waitFor are set up asynchronously, we need to wait for the
-  // first one to set up the binding in the page before running the others.
-  #mutex = new Mutex();
-  async _addBindingToContext(
-    context: ExecutionContext,
-    name: string
-  ): Promise<void> {
-    if (this.#contextBindings.has(name)) {
-      return;
-    }
-
-    using _ = await this.#mutex.acquire();
-    try {
-      await context._client.send(
-        'Runtime.addBinding',
-        context._contextName
-          ? {
-              name,
-              executionContextName: context._contextName,
-            }
-          : {
-              name,
-              executionContextId: context._contextId,
-            }
-      );
-
-      await context.evaluate(addPageBinding, 'internal', name);
-
-      this.#contextBindings.add(name);
-    } catch (error) {
-      // We could have tried to evaluate in a context which was already
-      // destroyed. This happens, for example, if the page is navigated while
-      // we are trying to add the binding
-      if (error instanceof Error) {
-        // Destroyed context.
-        if (error.message.includes('Execution context was destroyed')) {
-          return;
-        }
-        // Missing context.
-        if (error.message.includes('Cannot find context with specified id')) {
-          return;
-        }
-      }
-
-      debugError(error);
-    }
-  }
-
-  #onBindingCalled = async (
-    event: Protocol.Runtime.BindingCalledEvent
-  ): Promise<void> => {
-    let payload: BindingPayload;
-    try {
-      payload = JSON.parse(event.payload);
-    } catch {
-      // The binding was either called by something in the page or it was
-      // called before our wrapper was initialized.
-      return;
-    }
-    const {type, name, seq, args, isTrivial} = payload;
-    if (type !== 'internal') {
-      return;
-    }
-    if (!this.#contextBindings.has(name)) {
-      return;
-    }
-
-    try {
-      const context = await this.#context.valueOrThrow();
-      if (event.executionContextId !== context._contextId) {
-        return;
-      }
-
-      const binding = this._bindings.get(name);
-      await binding?.run(context, seq, args, isTrivial);
-    } catch (err) {
-      debugError(err);
-    }
-  };
-
   override async adoptBackendNode(
     backendNodeId?: Protocol.DOM.BackendNodeId
   ): Promise<JSHandle<Node>> {
@@ -282,7 +190,10 @@ export class IsolatedWorld extends Realm {
   }
 
   [disposeSymbol](): void {
+    const existingContext = this.#context.value();
+    if (existingContext instanceof ExecutionContext) {
+      existingContext[disposeSymbol]();
+    }
     super[disposeSymbol]();
-    this.client.off('Runtime.bindingCalled', this.#onBindingCalled);
   }
 }