From 496658f02945b53096483f36cb3d64556cff045e Mon Sep 17 00:00:00 2001 From: jrandolf <101637635+jrandolf@users.noreply.github.com> Date: Wed, 25 Jan 2023 23:56:33 -0800 Subject: [PATCH] fix: use LazyArg for puppeteer utilities (#9575) This PR fixes the following edge case: - `const oldPromise = world.puppeteerUtil`. - setContext occurs but context is immediately destroyed, i.e. `world.#puppeteerUtil === oldPromise` is not resolved. - clearContext occurs due to destruction, i.e. `world.#puppeteerUtil` is replaced (`world.#puppeteerUtil !== oldPromise`). - `oldPromise` never resolves. --- packages/puppeteer-core/src/common/Frame.ts | 15 ++++++++++---- .../src/common/IsolatedWorld.ts | 20 ++++++++----------- packages/puppeteer-core/src/common/LazyArg.ts | 8 +++++++- .../puppeteer-core/src/common/QueryHandler.ts | 14 +++++++++++-- .../puppeteer-core/src/common/WaitTask.ts | 13 +++++++++--- packages/puppeteer-core/src/common/types.ts | 13 ------------ test/src/injected.spec.ts | 5 ++++- 7 files changed, 52 insertions(+), 36 deletions(-) diff --git a/packages/puppeteer-core/src/common/Frame.ts b/packages/puppeteer-core/src/common/Frame.ts index cbd1097a..a4e56d16 100644 --- a/packages/puppeteer-core/src/common/Frame.ts +++ b/packages/puppeteer-core/src/common/Frame.ts @@ -34,6 +34,7 @@ import {Page} from '../api/Page.js'; import {getQueryHandlerAndSelector} from './QueryHandler.js'; import {EvaluateFunc, HandleFor, NodeFor} from './types.js'; import {importFS} from './util.js'; +import {LazyArg} from './LazyArg.js'; /** * @public @@ -804,8 +805,9 @@ export class Frame { type = type ?? 'text/javascript'; + const puppeteerWorld = this.worlds[PUPPETEER_WORLD]; return this.worlds[MAIN_WORLD].transferHandle( - await this.worlds[PUPPETEER_WORLD].evaluateHandle( + await puppeteerWorld.evaluateHandle( async ({createDeferredPromise}, {url, id, type, content}) => { const promise = createDeferredPromise(); const script = document.createElement('script'); @@ -839,7 +841,9 @@ export class Frame { await promise; return script; }, - await this.worlds[PUPPETEER_WORLD].puppeteerUtil, + LazyArg.create(() => { + return puppeteerWorld.puppeteerUtil; + }), {...options, type, content} ) ); @@ -887,8 +891,9 @@ export class Frame { options.content = content; } + const puppeteerWorld = this.worlds[PUPPETEER_WORLD]; return this.worlds[MAIN_WORLD].transferHandle( - await this.worlds[PUPPETEER_WORLD].evaluateHandle( + await puppeteerWorld.evaluateHandle( async ({createDeferredPromise}, {url, content}) => { const promise = createDeferredPromise(); let element: HTMLStyleElement | HTMLLinkElement; @@ -923,7 +928,9 @@ export class Frame { await promise; return element; }, - await this.worlds[PUPPETEER_WORLD].puppeteerUtil, + LazyArg.create(() => { + return puppeteerWorld.puppeteerUtil; + }), options ) ); diff --git a/packages/puppeteer-core/src/common/IsolatedWorld.ts b/packages/puppeteer-core/src/common/IsolatedWorld.ts index 13a22555..51af6705 100644 --- a/packages/puppeteer-core/src/common/IsolatedWorld.ts +++ b/packages/puppeteer-core/src/common/IsolatedWorld.ts @@ -28,7 +28,7 @@ import {JSHandle} from './JSHandle.js'; import {LazyArg} from './LazyArg.js'; import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js'; import {TimeoutSettings} from './TimeoutSettings.js'; -import {EvaluateFunc, HandleFor, InnerLazyParams, NodeFor} from './types.js'; +import {EvaluateFunc, HandleFor, NodeFor} from './types.js'; import {createJSHandle, debugError, pageBindingInitString} from './util.js'; import {TaskManager, WaitTask} from './WaitTask.js'; import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js'; @@ -138,8 +138,11 @@ export class IsolatedWorld { } clearContext(): void { + // Only create a new promise if the old one was resolved. + if (this.#puppeteerUtil.resolved()) { + this.#puppeteerUtil = createDeferredPromise(); + } this.#document = undefined; - this.#puppeteerUtil = createDeferredPromise(); this.#context = createDeferredPromise(); } @@ -517,13 +520,8 @@ export class IsolatedWorld { root, timeout, }, - new LazyArg(async () => { - try { - // In case CDP fails. - return await this.puppeteerUtil; - } catch { - return undefined; - } + LazyArg.create(() => { + return this.puppeteerUtil; }), queryOne.toString(), selector, @@ -547,9 +545,7 @@ export class IsolatedWorld { waitForFunction< Params extends unknown[], - Func extends EvaluateFunc> = EvaluateFunc< - InnerLazyParams - > + Func extends EvaluateFunc = EvaluateFunc >( pageFunction: Func | string, options: { diff --git a/packages/puppeteer-core/src/common/LazyArg.ts b/packages/puppeteer-core/src/common/LazyArg.ts index c182663a..8599e82c 100644 --- a/packages/puppeteer-core/src/common/LazyArg.ts +++ b/packages/puppeteer-core/src/common/LazyArg.ts @@ -18,8 +18,14 @@ * @internal */ export class LazyArg { + static create = (callback: () => Promise): T => { + // We do type coercion here because we don't want to introduce LazyArgs to + // the type system. + return new LazyArg(callback) as unknown as T; + }; + #get: () => Promise; - constructor(get: () => Promise) { + private constructor(get: () => Promise) { this.#get = get; } diff --git a/packages/puppeteer-core/src/common/QueryHandler.ts b/packages/puppeteer-core/src/common/QueryHandler.ts index da9a4530..b8cba2d7 100644 --- a/packages/puppeteer-core/src/common/QueryHandler.ts +++ b/packages/puppeteer-core/src/common/QueryHandler.ts @@ -15,11 +15,13 @@ */ import PuppeteerUtil from '../injected/injected.js'; +import {assert} from '../util/assert.js'; import {ariaHandler} from './AriaQueryHandler.js'; import {ElementHandle} from './ElementHandle.js'; import {Frame} from './Frame.js'; import {WaitForSelectorOptions} from './IsolatedWorld.js'; import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js'; +import {LazyArg} from './LazyArg.js'; /** * @public @@ -102,7 +104,11 @@ function createPuppeteerQueryHandler( const jsHandle = await element.evaluateHandle( queryOne, selector, - await element.executionContext()._world!.puppeteerUtil + LazyArg.create(() => { + const world = element.executionContext()._world; + assert(world); + return world.puppeteerUtil; + }) ); const elementHandle = jsHandle.asElement(); if (elementHandle) { @@ -148,7 +154,11 @@ function createPuppeteerQueryHandler( const jsHandle = await element.evaluateHandle( queryAll, selector, - await element.executionContext()._world!.puppeteerUtil + LazyArg.create(() => { + const world = element.executionContext()._world; + assert(world); + return world.puppeteerUtil; + }) ); const properties = await jsHandle.getProperties(); await jsHandle.dispose(); diff --git a/packages/puppeteer-core/src/common/WaitTask.ts b/packages/puppeteer-core/src/common/WaitTask.ts index 955cc2d8..3fff4d96 100644 --- a/packages/puppeteer-core/src/common/WaitTask.ts +++ b/packages/puppeteer-core/src/common/WaitTask.ts @@ -20,6 +20,7 @@ import {ElementHandle} from './ElementHandle.js'; import {TimeoutError} from './Errors.js'; import {IsolatedWorld} from './IsolatedWorld.js'; import {JSHandle} from './JSHandle.js'; +import {LazyArg} from './LazyArg.js'; import {HandleFor} from './types.js'; /** @@ -114,7 +115,9 @@ export class WaitTask { return fun(...args) as Promise; }); }, - await this.#world.puppeteerUtil, + LazyArg.create(() => { + return this.#world.puppeteerUtil; + }), this.#fn, ...this.#args ); @@ -127,7 +130,9 @@ export class WaitTask { return fun(...args) as Promise; }, root || document); }, - await this.#world.puppeteerUtil, + LazyArg.create(() => { + return this.#world.puppeteerUtil; + }), this.#root, this.#fn, ...this.#args @@ -141,7 +146,9 @@ export class WaitTask { return fun(...args) as Promise; }, ms); }, - await this.#world.puppeteerUtil, + LazyArg.create(() => { + return this.#world.puppeteerUtil; + }), this.#polling, this.#fn, ...this.#args diff --git a/packages/puppeteer-core/src/common/types.ts b/packages/puppeteer-core/src/common/types.ts index ed30e735..e51abde4 100644 --- a/packages/puppeteer-core/src/common/types.ts +++ b/packages/puppeteer-core/src/common/types.ts @@ -16,7 +16,6 @@ import {JSHandle} from './JSHandle.js'; import {ElementHandle} from './ElementHandle.js'; -import {LazyArg} from './LazyArg.js'; /** * @public @@ -38,18 +37,6 @@ export type HandleOr = HandleFor | JSHandle | T; */ export type FlattenHandle = T extends HandleOr ? U : never; -/** - * @internal - */ -export type FlattenLazyArg = T extends LazyArg ? U : T; - -/** - * @internal - */ -export type InnerLazyParams = { - [K in keyof T]: FlattenLazyArg; -}; - /** * @public */ diff --git a/test/src/injected.spec.ts b/test/src/injected.spec.ts index cde47fcc..3b51fa9f 100644 --- a/test/src/injected.spec.ts +++ b/test/src/injected.spec.ts @@ -16,6 +16,7 @@ import expect from 'expect'; import {PUPPETEER_WORLD} from 'puppeteer-core/internal/common/IsolatedWorlds.js'; +import {LazyArg} from 'puppeteer-core/internal/common/LazyArg.js'; import { getTestState, setupTestBrowserHooks, @@ -45,7 +46,9 @@ describe('PuppeteerUtil tests', function () { ({createFunction}, fnString) => { return createFunction(fnString)(4); }, - await world.puppeteerUtil, + LazyArg.create(() => { + return world.puppeteerUtil; + }), (() => { return 4; }).toString()