refactor: move binding called handling (#12410)

puppeteer · May 8, 2024 · 83a9716 · 83a9716
1 parent d0cd710
commit 83a9716
Show file tree

Hide file tree

Showing 6 changed files with 71 additions and 56 deletions.
diff --git a/packages/puppeteer-core/src/cdp/ExecutionContext.ts b/packages/puppeteer-core/src/cdp/ExecutionContext.ts
@@ -67,13 +67,15 @@ export class ExecutionContext
     /** Emitted when this execution context is disposed. */
     disposed: undefined;
     consoleapicalled: Protocol.Runtime.ConsoleAPICalledEvent;
+    /** Emitted when a binding that is not installed by the ExecutionContext is called. */
+    bindingcalled: Protocol.Runtime.BindingCalledEvent;
   }>
   implements Disposable
 {
-  _client: CDPSession;
-  _world: IsolatedWorld;
-  _contextId: number;
-  _contextName?: string;
+  #client: CDPSession;
+  #world: IsolatedWorld;
+  #id: number;
+  #name?: string;
 
   readonly #disposables = new DisposableStack();
 
@@ -83,16 +85,16 @@ export class ExecutionContext
     world: IsolatedWorld
   ) {
     super();
-    this._client = client;
-    this._world = world;
-    this._contextId = contextPayload.id;
+    this.#client = client;
+    this.#world = world;
+    this.#id = contextPayload.id;
     if (contextPayload.name) {
-      this._contextName = contextPayload.name;
+      this.#name = contextPayload.name;
     }
-    const clientEmitter = this.#disposables.use(new EventEmitter(this._client));
+    const clientEmitter = this.#disposables.use(new EventEmitter(this.#client));
     clientEmitter.on('Runtime.bindingCalled', this.#onBindingCalled.bind(this));
     clientEmitter.on('Runtime.executionContextDestroyed', async event => {
-      if (event.executionContextId === this._contextId) {
+      if (event.executionContextId === this.#id) {
         this[disposeSymbol]();
       }
     });
@@ -118,16 +120,16 @@ export class ExecutionContext
 
     using _ = await this.#mutex.acquire();
     try {
-      await this._client.send(
+      await this.#client.send(
         'Runtime.addBinding',
-        this._contextName
+        this.#name
           ? {
               name: binding.name,
-              executionContextName: this._contextName,
+              executionContextName: this.#name,
             }
           : {
               name: binding.name,
-              executionContextId: this._contextId,
+              executionContextId: this.#id,
             }
       );
 
@@ -166,14 +168,16 @@ export class ExecutionContext
     }
     const {type, name, seq, args, isTrivial} = payload;
     if (type !== 'internal') {
+      this.emit('bindingcalled', event);
       return;
     }
     if (!this.#bindings.has(name)) {
+      this.emit('bindingcalled', event);
       return;
     }
 
     try {
-      if (event.executionContextId !== this._contextId) {
+      if (event.executionContextId !== this.#id) {
         return;
       }
 
@@ -184,8 +188,12 @@ export class ExecutionContext
     }
   }
 
+  get id(): number {
+    return this.#id;
+  }
+
   #onConsoleAPI(event: Protocol.Runtime.ConsoleAPICalledEvent): void {
-    if (event.executionContextId !== this._contextId) {
+    if (event.executionContextId !== this.#id) {
       return;
     }
     this.emit('consoleapicalled', event);
@@ -365,13 +373,13 @@ export class ExecutionContext
     );
 
     if (isString(pageFunction)) {
-      const contextId = this._contextId;
+      const contextId = this.#id;
       const expression = pageFunction;
       const expressionWithSourceUrl = SOURCE_URL_REGEX.test(expression)
         ? expression
         : `${expression}\n${sourceUrlComment}\n`;
 
-      const {exceptionDetails, result: remoteObject} = await this._client
+      const {exceptionDetails, result: remoteObject} = await this.#client
         .send('Runtime.evaluate', {
           expression: expressionWithSourceUrl,
           contextId,
@@ -387,7 +395,7 @@ export class ExecutionContext
 
       return returnByValue
         ? valueFromRemoteObject(remoteObject)
-        : this._world.createCdpHandle(remoteObject);
+        : this.#world.createCdpHandle(remoteObject);
     }
 
     const functionDeclaration = stringifyFunction(pageFunction);
@@ -398,9 +406,9 @@ export class ExecutionContext
       : `${functionDeclaration}\n${sourceUrlComment}\n`;
     let callFunctionOnPromise;
     try {
-      callFunctionOnPromise = this._client.send('Runtime.callFunctionOn', {
+      callFunctionOnPromise = this.#client.send('Runtime.callFunctionOn', {
         functionDeclaration: functionDeclarationWithSourceUrl,
-        executionContextId: this._contextId,
+        executionContextId: this.#id,
         arguments: args.length
           ? await Promise.all(args.map(convertArgument.bind(this)))
           : [],
@@ -424,7 +432,7 @@ export class ExecutionContext
     }
     return returnByValue
       ? valueFromRemoteObject(remoteObject)
-      : this._world.createCdpHandle(remoteObject);
+      : this.#world.createCdpHandle(remoteObject);
 
     async function convertArgument(
       this: ExecutionContext,
@@ -454,7 +462,7 @@ export class ExecutionContext
           ? arg
           : null;
       if (objectHandle) {
-        if (objectHandle.realm !== this._world) {
+        if (objectHandle.realm !== this.#world) {
           throw new Error(
             'JSHandles can be evaluated only in the context they were created!'
           );

diff --git a/packages/puppeteer-core/src/cdp/Frame.ts b/packages/puppeteer-core/src/cdp/Frame.ts
@@ -80,6 +80,10 @@ export class CdpFrame extends Frame {
       'consoleapicalled',
       this.#onMainWorldConsoleApiCalled.bind(this)
     );
+    this.worlds[MAIN_WORLD].emitter.on(
+      'bindingcalled',
+      this.#onMainWorldBindingCalled.bind(this)
+    );
   }
 
   #onMainWorldConsoleApiCalled(
@@ -91,6 +95,13 @@ export class CdpFrame extends Frame {
     ]);
   }
 
+  #onMainWorldBindingCalled(event: Protocol.Runtime.BindingCalledEvent) {
+    this._frameManager.emit(FrameManagerEvent.BindingCalled, [
+      this.worlds[MAIN_WORLD],
+      event,
+    ]);
+  }
+
   /**
    * This is used internally in DevTools.
    *

diff --git a/packages/puppeteer-core/src/cdp/FrameManager.ts b/packages/puppeteer-core/src/cdp/FrameManager.ts
@@ -41,7 +41,6 @@ export class FrameManager extends EventEmitter<FrameManagerEvents> {
   #page: CdpPage;
   #networkManager: NetworkManager;
   #timeoutSettings: TimeoutSettings;
-  #contextIdToContext = new Map<string, ExecutionContext>();
   #isolatedWorlds = new Set<string>();
   #client: CDPSession;
 
@@ -223,22 +222,6 @@ export class FrameManager extends EventEmitter<FrameManagerEvents> {
     }
   }
 
-  executionContextById(
-    contextId: number,
-    session: CDPSession = this.#client
-  ): ExecutionContext {
-    const context = this.getExecutionContextById(contextId, session);
-    assert(context, 'INTERNAL ERROR: missing context with id = ' + contextId);
-    return context;
-  }
-
-  getExecutionContextById(
-    contextId: number,
-    session: CDPSession = this.#client
-  ): ExecutionContext | undefined {
-    return this.#contextIdToContext.get(`${session.id()}:${contextId}`);
-  }
-
   page(): CdpPage {
     return this.#page;
   }
@@ -491,16 +474,6 @@ export class FrameManager extends EventEmitter<FrameManagerEvents> {
       world
     );
     world.setContext(context);
-    const key = `${session.id()}:${contextPayload.id}`;
-    this.#contextIdToContext.set(key, context);
-    context.once('disposed', () => {
-      const key = `${session.id()}:${contextPayload.id}`;
-      const context = this.#contextIdToContext.get(key);
-      if (!context) {
-        return;
-      }
-      this.#contextIdToContext.delete(key);
-    });
   }
 
   #removeFramesRecursively(frame: CdpFrame): void {

diff --git a/packages/puppeteer-core/src/cdp/FrameManagerEvents.ts b/packages/puppeteer-core/src/cdp/FrameManagerEvents.ts
@@ -28,6 +28,7 @@ export namespace FrameManagerEvent {
     'FrameManager.FrameNavigatedWithinDocument'
   );
   export const ConsoleApiCalled = Symbol('FrameManager.ConsoleApiCalled');
+  export const BindingCalled = Symbol('FrameManager.BindingCalled');
 }
 
 /**
@@ -45,4 +46,8 @@ export interface FrameManagerEvents extends Record<EventType, unknown> {
     IsolatedWorld,
     Protocol.Runtime.ConsoleAPICalledEvent,
   ];
+  [FrameManagerEvent.BindingCalled]: [
+    IsolatedWorld,
+    Protocol.Runtime.BindingCalledEvent,
+  ];
 }
diff --git a/packages/puppeteer-core/src/cdp/IsolatedWorld.ts b/packages/puppeteer-core/src/cdp/IsolatedWorld.ts
@@ -54,6 +54,8 @@ type IsolatedWorldEmitter = EventEmitter<{
   disposed: undefined;
   // Emitted when a new console message is logged.
   consoleapicalled: Protocol.Runtime.ConsoleAPICalledEvent;
+  /** Emitted when a binding that is not installed by the ExecutionContext is called. */
+  bindingcalled: Protocol.Runtime.BindingCalledEvent;
 }>;
 
 /**
@@ -89,6 +91,7 @@ export class IsolatedWorld extends Realm {
     this.#context?.[disposeSymbol]();
     context.once('disposed', this.#onContextDisposed.bind(this));
     context.on('consoleapicalled', this.#onContextConsoleApiCalled.bind(this));
+    context.on('bindingcalled', this.#onContextBindingCalled.bind(this));
     this.#context = context;
     this.#emitter.emit('context', context);
     void this.taskManager.rerunAll();
@@ -107,10 +110,18 @@ export class IsolatedWorld extends Realm {
     this.#emitter.emit('consoleapicalled', event);
   }
 
+  #onContextBindingCalled(event: Protocol.Runtime.BindingCalledEvent): void {
+    this.#emitter.emit('bindingcalled', event);
+  }
+
   hasContext(): boolean {
     return !!this.#context;
   }
 
+  get context(): ExecutionContext | undefined {
+    return this.#context;
+  }
+
   #executionContext(): ExecutionContext | undefined {
     if (this.disposed) {
       throw new Error(
@@ -193,7 +204,7 @@ export class IsolatedWorld extends Realm {
     }
     const {object} = await this.client.send('DOM.resolveNode', {
       backendNodeId: backendNodeId,
-      executionContextId: context._contextId,
+      executionContextId: context.id,
     });
     return this.createCdpHandle(object) as JSHandle<Node>;
   }

diff --git a/packages/puppeteer-core/src/cdp/Page.ts b/packages/puppeteer-core/src/cdp/Page.ts
@@ -217,7 +217,6 @@ export class CdpPage extends Page {
         return this.emit(PageEvent.Load, undefined);
       },
     ],
-    ['Runtime.bindingCalled', this.#onBindingCalled.bind(this)],
     ['Page.javascriptDialogOpening', this.#onDialog.bind(this)],
     ['Runtime.exceptionThrown', this.#handleException.bind(this)],
     ['Inspector.targetCrashed', this.#onTargetCrashed.bind(this)],
@@ -259,6 +258,16 @@ export class CdpPage extends Page {
       }
     );
 
+    this.#frameManager.on(
+      FrameManagerEvent.BindingCalled,
+      ([world, event]: [
+        IsolatedWorld,
+        Protocol.Runtime.BindingCalledEvent,
+      ]) => {
+        void this.#onBindingCalled(world, event);
+      }
+    );
+
     for (const [eventName, handler] of this.#networkManagerHandlers) {
       // TODO: Remove any.
       this.#frameManager.networkManager.on(eventName, handler as any);
@@ -803,6 +812,7 @@ export class CdpPage extends Page {
   }
 
   async #onBindingCalled(
+    world: IsolatedWorld,
     event: Protocol.Runtime.BindingCalledEvent
   ): Promise<void> {
     let payload: BindingPayload;
@@ -818,10 +828,7 @@ export class CdpPage extends Page {
       return;
     }
 
-    const context = this.#frameManager.executionContextById(
-      event.executionContextId,
-      this.#primaryTargetClient
-    );
+    const context = world.context;
     if (!context) {
       return;
     }