From 5a05838b62557cee800e246604150ff471b134b2 Mon Sep 17 00:00:00 2001 From: Alex Rudenko Date: Mon, 6 May 2024 15:28:33 +0200 Subject: [PATCH] refactor: move binding installation into the execution context (#12398) --- .../src/cdp/ExecutionContext.ts | 104 +++++++++++++++- .../puppeteer-core/src/cdp/IsolatedWorld.ts | 111 ++---------------- 2 files changed, 110 insertions(+), 105 deletions(-) diff --git a/packages/puppeteer-core/src/cdp/ExecutionContext.ts b/packages/puppeteer-core/src/cdp/ExecutionContext.ts index 5fcdc1a57aa..0fae73088c6 100644 --- a/packages/puppeteer-core/src/cdp/ExecutionContext.ts +++ b/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(); + + // Contains mapping from functions that should be bound to Puppeteer functions. + #bindings = new Map(); + + // 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 { + 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 => { + 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>; get puppeteerUtil(): Promise> { @@ -97,8 +187,8 @@ export class ExecutionContext { async #installGlobalBinding(binding: Binding) { try { if (this._world) { - this._world._bindings.set(binding.name, binding); - await this._world._addBindingToContext(this, binding.name); + this.#bindings.set(binding.name, binding); + 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 => { diff --git a/packages/puppeteer-core/src/cdp/IsolatedWorld.ts b/packages/puppeteer-core/src/cdp/IsolatedWorld.ts index f4f73624ad1..06355867339 100644 --- a/packages/puppeteer-core/src/cdp/IsolatedWorld.ts +++ b/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 {debugError, withSourcePuppeteerURLIfNone} from '../common/util.js'; +import type {EvaluateFunc, HandleFor} from '../common/types.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(); - // Set of bindings that have been registered in the current context. - #contextBindings = new Set(); - - // Contains mapping from functions that should be bound to Puppeteer functions. - #bindings = new Map(); - - get _bindings(): Map { - return this.#bindings; - } - readonly #frameOrWorker: CdpFrame | CdpWebWorker; constructor( @@ -74,9 +61,7 @@ export class IsolatedWorld extends Realm { return this.#frameOrWorker; } - frameUpdated(): void { - this.client.on('Runtime.bindingCalled', this.#onBindingCalled); - } + frameUpdated(): void {} 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 { - 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 => { - 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> { @@ -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); } }