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.
puppeteer · Jan 26, 2023 · 496658f · 496658f
1 parent 0b71700
commit 496658f
Show file tree

Hide file tree

Showing 7 changed files with 52 additions and 36 deletions.
diff --git 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<void>();
           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<void>();
           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
@@ -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<InnerLazyParams<Params>> = EvaluateFunc<
-      InnerLazyParams<Params>
-    >
+    Func extends EvaluateFunc<Params> = EvaluateFunc<Params>
   >(
     pageFunction: Func | string,
     options: {

diff --git a/packages/puppeteer-core/src/common/LazyArg.ts b/packages/puppeteer-core/src/common/LazyArg.ts
@@ -18,8 +18,14 @@
  * @internal
  */
 export class LazyArg<T> {
+  static create = <T>(callback: () => Promise<T>): 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<T>;
-  constructor(get: () => Promise<T>) {
+  private constructor(get: () => Promise<T>) {
     this.#get = get;
   }
 

diff --git 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
@@ -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<T = unknown> {
                 return fun(...args) as Promise<T>;
               });
             },
-            await this.#world.puppeteerUtil,
+            LazyArg.create(() => {
+              return this.#world.puppeteerUtil;
+            }),
             this.#fn,
             ...this.#args
           );
@@ -127,7 +130,9 @@ export class WaitTask<T = unknown> {
                 return fun(...args) as Promise<T>;
               }, 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<T = unknown> {
                 return fun(...args) as Promise<T>;
               }, 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
@@ -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<T> = HandleFor<T> | JSHandle<T> | T;
  */
 export type FlattenHandle<T> = T extends HandleOr<infer U> ? U : never;
 
-/**
- * @internal
- */
-export type FlattenLazyArg<T> = T extends LazyArg<infer U> ? U : T;
-
-/**
- * @internal
- */
-export type InnerLazyParams<T extends unknown[]> = {
-  [K in keyof T]: FlattenLazyArg<T[K]>;
-};
-
 /**
  * @public
  */

diff --git 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()