refactor: implement reverse argument binding (#9651)

packages/puppeteer-core/src/common/Binding.ts

@@ -0,0 +1,119 @@
+import {JSHandle} from '../api/JSHandle.js';
+import {isErrorLike} from '../util/ErrorLike.js';
+import {ExecutionContext} from './ExecutionContext.js';
+import {debugError} from './util.js';
+
+/**
+ * @internal
+ */
+export class Binding {
+  #name: string;
+  #fn: (...args: unknown[]) => unknown;
+  constructor(name: string, fn: (...args: unknown[]) => unknown) {
+    this.#name = name;
+    this.#fn = fn;
+  }
+
+  /**
+   *
+   * @param context - Context to run the binding in; the context should have
+   * the binding added to it beforehand.
+   * @param id - ID of the call. This should come from the CDP
+   * `onBindingCalled` response.
+   * @param args - Plain arguments from CDP.
+   */
+  async run(
+    context: ExecutionContext,
+    id: number,
+    args: unknown[],
+    isTrivial: boolean
+  ): Promise<void> {
+    const garbage = [];
+    try {
+      if (!isTrivial) {
+        // Getting non-trivial arguments.
+        const handles = await context.evaluateHandle(
+          (name, seq) => {
+            // @ts-expect-error Code is evaluated in a different context.
+            return globalThis[name].args.get(seq);
+          },
+          this.#name,
+          id
+        );
+        try {
+          const properties = await handles.getProperties();
+          for (const [index, handle] of properties) {
+            // This is not straight-forward since some arguments can stringify, but
+            // aren't plain objects so add subtypes when the use-case arises.
+            if (index in args) {
+              switch (handle.remoteObject().subtype) {
+                case 'node':
+                  args[+index] = handle;
+                  break;
+                default:
+                  garbage.push(handle.dispose());
+              }
+            } else {
+              garbage.push(handle.dispose());
+            }
+          }
+        } finally {
+          await handles.dispose();
+        }
+      }
+
+      await context.evaluate(
+        (name, seq, result) => {
+          // @ts-expect-error Code is evaluated in a different context.
+          const callbacks = globalThis[name].callbacks;
+          callbacks.get(seq).resolve(result);
+          callbacks.delete(seq);
+        },
+        this.#name,
+        id,
+        await this.#fn(...args)
+      );
+
+      for (const arg of args) {
+        if (arg instanceof JSHandle) {
+          garbage.push(arg.dispose());
+        }
+      }
+    } catch (error) {
+      if (isErrorLike(error)) {
+        await context
+          .evaluate(
+            (name, seq, message, stack) => {
+              const error = new Error(message);
+              error.stack = stack;
+              // @ts-expect-error Code is evaluated in a different context.
+              const callbacks = globalThis[name].callbacks;
+              callbacks.get(seq).reject(error);
+              callbacks.delete(seq);
+            },
+            this.#name,
+            id,
+            error.message,
+            error.stack
+          )
+          .catch(debugError);
+      } else {
+        await context
+          .evaluate(
+            (name, seq, error) => {
+              // @ts-expect-error Code is evaluated in a different context.
+              const callbacks = globalThis[name].callbacks;
+              callbacks.get(seq).reject(error);
+              callbacks.delete(seq);
+            },
+            this.#name,
+            id,
+            error
+          )
+          .catch(debugError);
+      }
+    } finally {
+      await Promise.all(garbage);
+    }
+  }
+}

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

@@ -28,16 +28,18 @@ import {MAIN_WORLD, PUPPETEER_WORLD} from './IsolatedWorlds.js';
 import {LifecycleWatcher, PuppeteerLifeCycleEvent} from './LifecycleWatcher.js';
 import {TimeoutSettings} from './TimeoutSettings.js';
 import {
+  BindingPayload,
   EvaluateFunc,
   EvaluateFuncWith,
   HandleFor,
   InnerLazyParams,
   NodeFor,
 } from './types.js';
-import {createJSHandle, debugError, pageBindingInitString} from './util.js';
+import {addPageBinding, createJSHandle, debugError} from './util.js';
 import {TaskManager, WaitTask} from './WaitTask.js';
 
 import type {ElementHandle} from '../api/ElementHandle.js';
+import {Binding} from './Binding.js';
 import {LazyArg} from './LazyArg.js';
 
 /**
@@ -95,24 +97,20 @@ export class IsolatedWorld {
   #detached = false;
 
   // Set of bindings that have been registered in the current context.
-  #ctxBindings = new Set<string>();
+  #contextBindings = new Set<string>();
 
   // Contains mapping from functions that should be bound to Puppeteer functions.
-  #boundFunctions = new Map<string, Function>();
+  #bindings = new Map<string, Binding>();
   #taskManager = new TaskManager();
 
   get taskManager(): TaskManager {
     return this.#taskManager;
   }
 
-  get _boundFunctions(): Map<string, Function> {
+  get _bindings(): Map<string, Binding> {
-    return this.#boundFunctions;
+    return this.#bindings;
   }
 
-  static #bindingIdentifier = (name: string, contextId: number) => {
-    return `${name}_${contextId}`;
-  };
-
   constructor(frame: Frame) {
     // Keep own reference to client because it might differ from the FrameManager's
     // client for OOP iframes.
@@ -142,7 +140,7 @@ export class IsolatedWorld {
   }
 
   setContext(context: ExecutionContext): void {
-    this.#ctxBindings.clear();
+    this.#contextBindings.clear();
     this.#context.resolve(context);
     this.#taskManager.rerunAll();
   }
@@ -354,71 +352,50 @@ export class IsolatedWorld {
 
   // 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.
-  #settingUpBinding: Promise<void> | null = null;
+  #mutex = new Mutex();
-
   async _addBindingToContext(
     context: ExecutionContext,
     name: string
   ): Promise<void> {
-    // Previous operation added the binding so we are done.
+    if (this.#contextBindings.has(name)) {
-    if (
-      this.#ctxBindings.has(
-        IsolatedWorld.#bindingIdentifier(name, context._contextId)
-      )
-    ) {
       return;
     }
-    // Wait for other operation to finish
-    if (this.#settingUpBinding) {
-      await this.#settingUpBinding;
-      return this._addBindingToContext(context, name);
-    }
 
-    const bind = async (name: string) => {
+    await this.#mutex.acquire();
-      const expression = pageBindingInitString('internal', name);
+    try {
-      try {
+      await context._client.send('Runtime.addBinding', {
-        // TODO: In theory, it would be enough to call this just once
+        name,
-        await context._client.send('Runtime.addBinding', {
+        executionContextName: context._contextName,
-          name,
+      });
-          executionContextName: context._contextName,
+
-        });
+      await context.evaluate(addPageBinding, 'internal', name);
-        await context.evaluate(expression);
+
-      } catch (error) {
+      this.#contextBindings.add(name);
-        // We could have tried to evaluate in a context which was already
+    } catch (error) {
-        // destroyed. This happens, for example, if the page is navigated while
+      // We could have tried to evaluate in a context which was already
-        // we are trying to add the binding
+      // destroyed. This happens, for example, if the page is navigated while
-        if (error instanceof Error) {
+      // we are trying to add the binding
-          // Destroyed context.
+      if (error instanceof Error) {
-          if (error.message.includes('Execution context was destroyed')) {
+        // Destroyed context.
-            return;
+        if (error.message.includes('Execution context was destroyed')) {
-          }
+          return;
-          // Missing context.
+        }
-          if (error.message.includes('Cannot find context with specified id')) {
+        // Missing context.
-            return;
+        if (error.message.includes('Cannot find context with specified id')) {
-          }
+          return;
         }
-
-        debugError(error);
-        return;
       }
-      this.#ctxBindings.add(
-        IsolatedWorld.#bindingIdentifier(name, context._contextId)
-      );
-    };
 
-    this.#settingUpBinding = bind(name);
+      debugError(error);
-    await this.#settingUpBinding;
+    } finally {
-    this.#settingUpBinding = null;
+      this.#mutex.release();
+    }
   }
 
   #onBindingCalled = async (
     event: Protocol.Runtime.BindingCalledEvent
   ): Promise<void> => {
-    let payload: {type: string; name: string; seq: number; args: unknown[]};
+    let payload: BindingPayload;
-    if (!this.hasContext()) {
-      return;
-    }
-    const context = await this.executionContext();
     try {
       payload = JSON.parse(event.payload);
     } catch {
@@ -426,46 +403,21 @@ export class IsolatedWorld {
       // called before our wrapper was initialized.
       return;
     }
-    const {type, name, seq, args} = payload;
+    const {type, name, seq, args, isTrivial} = payload;
-    if (
+    if (type !== 'internal') {
-      type !== 'internal' ||
-      !this.#ctxBindings.has(
-        IsolatedWorld.#bindingIdentifier(name, context._contextId)
-      )
-    ) {
       return;
     }
-    if (context._contextId !== event.executionContextId) {
+    if (!this.#contextBindings.has(name)) {
       return;
     }
-    try {
+
-      const fn = this._boundFunctions.get(name);
+    const context = await this.#context;
-      if (!fn) {
+    if (event.executionContextId !== context._contextId) {
-        throw new Error(`Bound function $name is not found`);
+      return;
-      }
-      const result = await fn(...args);
-      await context.evaluate(
-        (name: string, seq: number, result: unknown) => {
-          // @ts-expect-error Code is evaluated in a different context.
-          const callbacks = self[name].callbacks;
-          callbacks.get(seq).resolve(result);
-          callbacks.delete(seq);
-        },
-        name,
-        seq,
-        result
-      );
-    } catch (error) {
-      // The WaitTask may already have been resolved by timing out, or the
-      // execution context may have been destroyed.
-      // In both caes, the promises above are rejected with a protocol error.
-      // We can safely ignores these, as the WaitTask is re-installed in
-      // the next execution context if needed.
-      if ((error as Error).message.includes('Protocol error')) {
-        return;
-      }
-      debugError(error);
     }
+
+    const binding = this._bindings.get(name);
+    await binding?.run(context, seq, args, isTrivial);
   };
 
   async _waitForSelectorInPage(
@@ -598,3 +550,31 @@ export class IsolatedWorld {
     return result;
   }
 }
+
+class Mutex {
+  #locked = false;
+  #acquirers: Array<() => void> = [];
+
+  // This is FIFO.
+  acquire(): Promise<void> {
+    if (!this.#locked) {
+      this.#locked = true;
+      return Promise.resolve();
+    }
+    let resolve!: () => void;
+    const promise = new Promise<void>(res => {
+      resolve = res;
+    });
+    this.#acquirers.push(resolve);
+    return promise;
+  }
+
+  release(): void {
+    const resolve = this.#acquirers.shift();
+    if (!resolve) {
+      this.#locked = false;
+      return;
+    }
+    resolve();
+  }
+}

packages/puppeteer-core/src/common/Page.ts

@@ -18,6 +18,8 @@ import {Protocol} from 'devtools-protocol';
 import type {Readable} from 'stream';
 import type {Browser} from '../api/Browser.js';
 import type {BrowserContext} from '../api/BrowserContext.js';
+import {ElementHandle} from '../api/ElementHandle.js';
+import {JSHandle} from '../api/JSHandle.js';
 import {
   GeolocationOptions,
   MediaFeature,
@@ -36,6 +38,7 @@ import {
 } from '../util/DeferredPromise.js';
 import {isErrorLike} from '../util/ErrorLike.js';
 import {Accessibility} from './Accessibility.js';
+import {Binding} from './Binding.js';
 import {
   CDPSession,
   CDPSessionEmittedEvents,
@@ -44,7 +47,6 @@ import {
 import {ConsoleMessage, ConsoleMessageType} from './ConsoleMessage.js';
 import {Coverage} from './Coverage.js';
 import {Dialog} from './Dialog.js';
-import {ElementHandle} from '../api/ElementHandle.js';
 import {EmulationManager} from './EmulationManager.js';
 import {FileChooser} from './FileChooser.js';
 import {
@@ -59,7 +61,6 @@ import {HTTPResponse} from './HTTPResponse.js';
 import {Keyboard, Mouse, MouseButton, Touchscreen} from './Input.js';
 import {WaitForSelectorOptions} from './IsolatedWorld.js';
 import {MAIN_WORLD} from './IsolatedWorlds.js';
-import {JSHandle} from '../api/JSHandle.js';
 import {
   Credentials,
   NetworkConditions,
@@ -72,7 +73,13 @@ import {TargetManagerEmittedEvents} from './TargetManager.js';
 import {TaskQueue} from './TaskQueue.js';
 import {TimeoutSettings} from './TimeoutSettings.js';
 import {Tracing} from './Tracing.js';
-import {EvaluateFunc, EvaluateFuncWith, HandleFor, NodeFor} from './types.js';
+import {
+  BindingPayload,
+  EvaluateFunc,
+  EvaluateFuncWith,
+  HandleFor,
+  NodeFor,
+} from './types.js';
 import {
   createJSHandle,
   debugError,
@@ -83,9 +90,6 @@ import {
   importFS,
   isNumber,
   isString,
-  pageBindingDeliverErrorString,
-  pageBindingDeliverErrorValueString,
-  pageBindingDeliverResultString,
   pageBindingInitString,
   releaseObject,
   valueFromRemoteObject,
@@ -140,7 +144,7 @@ export class CDPPage extends Page {
   #frameManager: FrameManager;
   #emulationManager: EmulationManager;
   #tracing: Tracing;
-  #pageBindings = new Map<string, Function>();
+  #bindings = new Map<string, Binding>();
   #coverage: Coverage;
   #javascriptEnabled = true;
   #viewport: Viewport | null;
@@ -648,23 +652,29 @@ export class CDPPage extends Page {
     name: string,
     pptrFunction: Function | {default: Function}
   ): Promise<void> {
-    if (this.#pageBindings.has(name)) {
+    if (this.#bindings.has(name)) {
       throw new Error(
         `Failed to add page binding with name ${name}: window['${name}'] already exists!`
       );
     }
 
-    let exposedFunction: Function;
+    let binding: Binding;
     switch (typeof pptrFunction) {
       case 'function':
-        exposedFunction = pptrFunction;
+        binding = new Binding(
+          name,
+          pptrFunction as (...args: unknown[]) => unknown
+        );
         break;
       default:
-        exposedFunction = pptrFunction.default;
+        binding = new Binding(
+          name,
+          pptrFunction.default as (...args: unknown[]) => unknown
+        );
         break;
     }
 
-    this.#pageBindings.set(name, exposedFunction);
+    this.#bindings.set(name, binding);
 
     const expression = pageBindingInitString('exposedFun', name);
     await this.#client.send('Runtime.addBinding', {name: name});
@@ -772,7 +782,7 @@ export class CDPPage extends Page {
   async #onBindingCalled(
     event: Protocol.Runtime.BindingCalledEvent
   ): Promise<void> {
-    let payload: {type: string; name: string; seq: number; args: unknown[]};
+    let payload: BindingPayload;
     try {
       payload = JSON.parse(event.payload);
     } catch {
@@ -780,34 +790,21 @@ export class CDPPage extends Page {
       // called before our wrapper was initialized.
       return;
     }
-    const {type, name, seq, args} = payload;
+    const {type, name, seq, args, isTrivial} = payload;
-    if (type !== 'exposedFun' || !this.#pageBindings.has(name)) {
+    if (type !== 'exposedFun') {
       return;
     }
-    let expression = null;
+
-    try {
+    const context = this.#frameManager.executionContextById(
-      const pageBinding = this.#pageBindings.get(name);
+      event.executionContextId,
-      assert(pageBinding);
+      this.#client
-      const result = await pageBinding(...args);
+    );
-      expression = pageBindingDeliverResultString(name, seq, result);
+    if (!context) {
-    } catch (error) {
+      return;
-      if (isErrorLike(error)) {
-        expression = pageBindingDeliverErrorString(
-          name,
-          seq,
-          error.message,
-          error.stack
-        );
-      } else {
-        expression = pageBindingDeliverErrorValueString(name, seq, error);
-      }
     }
-    this.#client
+
-      .send('Runtime.evaluate', {
+    const binding = this.#bindings.get(name);
-        expression,
+    await binding?.run(context, seq, args, isTrivial);
-        contextId: event.executionContextId,
-      })
-      .catch(debugError);
   }
 
   #addConsoleMessage(

packages/puppeteer-core/src/common/WaitTask.ts

@@ -14,12 +14,13 @@
  * limitations under the License.
  */
 
+import {ElementHandle} from '../api/ElementHandle.js';
+import {JSHandle} from '../api/JSHandle.js';
 import type {Poller} from '../injected/Poller.js';
 import {createDeferredPromise} from '../util/DeferredPromise.js';
-import {ElementHandle} from '../api/ElementHandle.js';
+import {Binding} from './Binding.js';
 import {TimeoutError} from './Errors.js';
 import {IsolatedWorld} from './IsolatedWorld.js';
-import {JSHandle} from '../api/JSHandle.js';
 import {LazyArg} from './LazyArg.js';
 import {HandleFor} from './types.js';
 
@@ -84,7 +85,10 @@ export class WaitTask<T = unknown> {
 
     if (this.#bindings.size !== 0) {
       for (const [name, fn] of this.#bindings) {
-        this.#world._boundFunctions.set(name, fn);
+        this.#world._bindings.set(
+          name,
+          new Binding(name, fn as (...args: unknown[]) => unknown)
+        );
       }
     }
 

packages/puppeteer-core/src/common/types.ts

@@ -14,9 +14,23 @@
  * limitations under the License.
  */
 
-import {JSHandle} from '../api/JSHandle.js';
+import type {JSHandle} from '../api/JSHandle.js';
-import {ElementHandle} from '../api/ElementHandle.js';
+import type {ElementHandle} from '../api/ElementHandle.js';
-import {LazyArg} from './LazyArg.js';
+import type {LazyArg} from './LazyArg.js';
+
+/**
+ * @internal
+ */
+export type BindingPayload = {
+  type: string;
+  name: string;
+  seq: number;
+  args: unknown[];
+  /**
+   * Determines whether the arguments of the payload are trivial.
+   */
+  isTrivial: boolean;
+};
 
 /**
  * @public

packages/puppeteer-core/src/common/util.ts

@@ -270,84 +270,59 @@ export function evaluationString(
 /**
  * @internal
  */
-export function pageBindingInitString(type: string, name: string): string {
+export function addPageBinding(type: string, name: string): void {
-  function addPageBinding(type: string, name: string): void {
+  // This is the CDP binding.
-    // This is the CDP binding.
+  // @ts-expect-error: In a different context.
-    // @ts-expect-error: In a different context.
+  const callCDP = globalThis[name];
-    const callCDP = self[name];
 
-    // We replace the CDP binding with a Puppeteer binding.
+  // We replace the CDP binding with a Puppeteer binding.
-    Object.assign(self, {
+  Object.assign(globalThis, {
-      [name](...args: unknown[]): Promise<unknown> {
+    [name](...args: unknown[]): Promise<unknown> {
-        // This is the Puppeteer binding.
+      // This is the Puppeteer binding.
-        // @ts-expect-error: In a different context.
+      // @ts-expect-error: In a different context.
-        const callPuppeteer = self[name];
+      const callPuppeteer = globalThis[name];
-        callPuppeteer.callbacks ??= new Map();
+      callPuppeteer.args ??= new Map();
-        const seq = (callPuppeteer.lastSeq ?? 0) + 1;
+      callPuppeteer.callbacks ??= new Map();
-        callPuppeteer.lastSeq = seq;
+
-        callCDP(JSON.stringify({type, name, seq, args}));
+      const seq = (callPuppeteer.lastSeq ?? 0) + 1;
-        return new Promise((resolve, reject) => {
+      callPuppeteer.lastSeq = seq;
-          callPuppeteer.callbacks.set(seq, {resolve, reject});
+      callPuppeteer.args.set(seq, args);
+
+      callCDP(
+        JSON.stringify({
+          type,
+          name,
+          seq,
+          args,
+          isTrivial: !args.some(value => {
+            return value instanceof Node;
+          }),
+        })
+      );
+
+      return new Promise((resolve, reject) => {
+        callPuppeteer.callbacks.set(seq, {
+          resolve(value: unknown) {
+            callPuppeteer.args.delete(seq);
+            resolve(value);
+          },
+          reject(value?: unknown) {
+            callPuppeteer.args.delete(seq);
+            reject(value);
+          },
         });
-      },
+      });
-    });
+    },
-  }
+  });
+}
+
+/**
+ * @internal
+ */
+export function pageBindingInitString(type: string, name: string): string {
   return evaluationString(addPageBinding, type, name);
 }
 
-/**
- * @internal
- */
-export function pageBindingDeliverResultString(
-  name: string,
-  seq: number,
-  result: unknown
-): string {
-  function deliverResult(name: string, seq: number, result: unknown): void {
-    (window as any)[name].callbacks.get(seq).resolve(result);
-    (window as any)[name].callbacks.delete(seq);
-  }
-  return evaluationString(deliverResult, name, seq, result);
-}
-
-/**
- * @internal
- */
-export function pageBindingDeliverErrorString(
-  name: string,
-  seq: number,
-  message: string,
-  stack?: string
-): string {
-  function deliverError(
-    name: string,
-    seq: number,
-    message: string,
-    stack?: string
-  ): void {
-    const error = new Error(message);
-    error.stack = stack;
-    (window as any)[name].callbacks.get(seq).reject(error);
-    (window as any)[name].callbacks.delete(seq);
-  }
-  return evaluationString(deliverError, name, seq, message, stack);
-}
-
-/**
- * @internal
- */
-export function pageBindingDeliverErrorValueString(
-  name: string,
-  seq: number,
-  value: unknown
-): string {
-  function deliverErrorValue(name: string, seq: number, value: unknown): void {
-    (window as any)[name].callbacks.get(seq).reject(value);
-    (window as any)[name].callbacks.delete(seq);
-  }
-  return evaluationString(deliverErrorValue, name, seq, value);
-}
-
 /**
  * @internal
  */