chore: refactor exposed function for BiDi (#11976)

packages/puppeteer-core/src/bidi/ExposedFunction.ts

@@ -6,29 +6,24 @@
 
 import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';
 
+import {EventEmitter} from '../common/EventEmitter.js';
 import type {Awaitable, FlattenHandle} from '../common/types.js';
 import {debugError} from '../common/util.js';
-import {assert} from '../util/assert.js';
+import {DisposableStack} from '../util/disposable.js';
-import {Deferred} from '../util/Deferred.js';
 import {interpolateFunction, stringifyFunction} from '../util/Function.js';
 
 import type {Connection} from './core/Connection.js';
-import {BidiDeserializer} from './Deserializer.js';
+import {BidiElementHandle} from './ElementHandle.js';
 import type {BidiFrame} from './Frame.js';
-import {BidiSerializer} from './Serializer.js';
+import {BidiJSHandle} from './JSHandle.js';
 
-type SendArgsChannel<Args> = (value: [id: number, args: Args]) => void;
+type CallbackChannel<Args, Ret> = (
-type SendResolveChannel<Ret> = (
+  value: [
-  value: [id: number, resolve: (ret: FlattenHandle<Awaited<Ret>>) => void]
+    resolve: (ret: FlattenHandle<Awaited<Ret>>) => void,
+    reject: (error: unknown) => void,
+    args: Args,
+  ]
 ) => void;
-type SendRejectChannel = (
-  value: [id: number, reject: (error: unknown) => void]
-) => void;
-
-interface RemotePromiseCallbacks {
-  resolve: Deferred<Bidi.Script.RemoteValue>;
-  reject: Deferred<Bidi.Script.RemoteValue>;
-}
 
 /**
  * @internal
@@ -38,65 +33,53 @@ export class ExposeableFunction<Args extends unknown[], Ret> {
 
   readonly name;
   readonly #apply;
+  readonly #isolate;
 
-  readonly #channels;
+  readonly #channel;
-  readonly #callerInfos = new Map<
-    string,
-    Map<number, RemotePromiseCallbacks>
-  >();
 
-  #preloadScriptId?: Bidi.Script.PreloadScript;
+  #scripts: Array<[BidiFrame, Bidi.Script.PreloadScript]> = [];
+  #disposables = new DisposableStack();
 
   constructor(
     frame: BidiFrame,
     name: string,
-    apply: (...args: Args) => Awaitable<Ret>
+    apply: (...args: Args) => Awaitable<Ret>,
+    isolate = false
   ) {
     this.#frame = frame;
     this.name = name;
     this.#apply = apply;
+    this.#isolate = isolate;
 
-    this.#channels = {
+    this.#channel = `__puppeteer__${this.#frame._id}_page_exposeFunction_${this.name}`;
-      args: `__puppeteer__${this.#frame._id}_page_exposeFunction_${this.name}_args`,
-      resolve: `__puppeteer__${this.#frame._id}_page_exposeFunction_${this.name}_resolve`,
-      reject: `__puppeteer__${this.#frame._id}_page_exposeFunction_${this.name}_reject`,
-    };
   }
 
   async expose(): Promise<void> {
     const connection = this.#connection;
-    const channelArguments = this.#channelArguments;
+    const channel = {
+      type: 'channel' as const,
+      value: {
+        channel: this.#channel,
+        ownership: Bidi.Script.ResultOwnership.Root,
+      },
+    };
 
-    // TODO(jrandolf): Implement cleanup with removePreloadScript.
+    const connectionEmitter = this.#disposables.use(
-    connection.on(
+      new EventEmitter(connection)
-      Bidi.ChromiumBidi.Script.EventNames.Message,
-      this.#handleArgumentsMessage
     );
-    connection.on(
+    connectionEmitter.on(
       Bidi.ChromiumBidi.Script.EventNames.Message,
-      this.#handleResolveMessage
+      this.#handleMessage
-    );
-    connection.on(
-      Bidi.ChromiumBidi.Script.EventNames.Message,
-      this.#handleRejectMessage
     );
 
     const functionDeclaration = stringifyFunction(
       interpolateFunction(
-        (
+        (callback: CallbackChannel<Args, Ret>) => {
-          sendArgs: SendArgsChannel<Args>,
-          sendResolve: SendResolveChannel<Ret>,
-          sendReject: SendRejectChannel
-        ) => {
-          let id = 0;
           Object.assign(globalThis, {
             [PLACEHOLDER('name') as string]: function (...args: Args) {
               return new Promise<FlattenHandle<Awaited<Ret>>>(
                 (resolve, reject) => {
-                  sendArgs([id, args]);
+                  callback([resolve, reject, args]);
-                  sendResolve([id, resolve]);
-                  sendReject([id, reject]);
-                  ++id;
                 }
               );
             },
@@ -106,179 +89,140 @@ export class ExposeableFunction<Args extends unknown[], Ret> {
       )
     );
 
-    const {result} = await connection.send('script.addPreloadScript', {
+    const frames = [this.#frame];
-      functionDeclaration,
+    for (const frame of frames) {
-      arguments: channelArguments,
+      frames.push(...frame.childFrames());
-      contexts: [this.#frame.page().mainFrame()._id],
+    }
-    });
-    this.#preloadScriptId = result.script;
 
     await Promise.all(
-      this.#frame
+      frames.map(async frame => {
-        .page()
+        const realm = this.#isolate ? frame.isolatedRealm() : frame.mainRealm();
-        .frames()
+        try {
-        .map(async frame => {
+          this.#scripts.push([
-          return await connection.send('script.callFunction', {
+            frame,
-            functionDeclaration,
+            await frame.browsingContext.addPreloadScript(functionDeclaration, {
-            arguments: channelArguments,
+              arguments: [channel],
-            awaitPromise: false,
+              sandbox: realm.sandbox,
-            target: frame.mainRealm().realm.target,
+            }),
+          ]);
+        } catch (error) {
+          // If it errors, the frame probably doesn't support adding preload
+          // scripts. We fail gracefully.
+          debugError(error);
+        }
+
+        try {
+          await realm.realm.callFunction(functionDeclaration, false, {
+            arguments: [channel],
           });
+        } catch (error) {
+          // If it errors, the frame probably doesn't support call function. We
+          // fail gracefully.
+          debugError(error);
+        }
       })
     );
   }
 
-  #handleArgumentsMessage = async (params: Bidi.Script.MessageParameters) => {
+  get #connection(): Connection {
-    if (params.channel !== this.#channels.args) {
+    return this.#frame.page().browser().connection;
+  }
+
+  #handleMessage = async (params: Bidi.Script.MessageParameters) => {
+    if (params.channel !== this.#channel) {
       return;
     }
-    const connection = this.#connection;
+    const realm = this.#getRealm(params.source);
-    const {callbacks, remoteValue} = this.#getCallbacksAndRemoteValue(params);
+    if (!realm) {
-    const args = remoteValue.value?.[1];
+      // Unrelated message.
-    assert(args);
+      return;
-    try {
+    }
-      const result = await this.#apply(...BidiDeserializer.deserialize(args));
+
-      await connection.send('script.callFunction', {
+    using dataHandle = BidiJSHandle.from<
-        functionDeclaration: stringifyFunction(([_, resolve]: any, result) => {
+      [
-          resolve(result);
+        resolve: (ret: FlattenHandle<Awaited<Ret>>) => void,
-        }),
+        reject: (error: unknown) => void,
-        arguments: [
+        args: Args,
-          (await callbacks.resolve.valueOrThrow()) as Bidi.Script.LocalValue,
+      ]
-          BidiSerializer.serialize(result),
+    >(params.data, realm);
-        ],
+
-        awaitPromise: false,
+    using argsHandle = await dataHandle.evaluateHandle(([, , args]) => {
-        target: {
+      return args;
-          realm: params.source.realm,
-        },
     });
+
+    using stack = new DisposableStack();
+    const args = [];
+    for (const [index, handle] of await argsHandle.getProperties()) {
+      stack.use(handle);
+
+      // Element handles are passed as is.
+      if (handle instanceof BidiElementHandle) {
+        args[+index] = handle;
+        stack.use(handle);
+        continue;
+      }
+
+      // Everything else is passed as the JS value.
+      args[+index] = handle.jsonValue();
+    }
+
+    let result;
+    try {
+      result = await this.#apply(...((await Promise.all(args)) as Args));
     } catch (error) {
       try {
         if (error instanceof Error) {
-          await connection.send('script.callFunction', {
+          await dataHandle.evaluate(
-            functionDeclaration: stringifyFunction(
+            ([, reject], name, message, stack) => {
-              (
-                [_, reject]: [unknown, (error: Error) => void],
-                name: string,
-                message: string,
-                stack?: string
-              ) => {
               const error = new Error(message);
               error.name = name;
               if (stack) {
                 error.stack = stack;
               }
               reject(error);
-              }
-            ),
-            arguments: [
-              (await callbacks.reject.valueOrThrow()) as Bidi.Script.LocalValue,
-              BidiSerializer.serialize(error.name),
-              BidiSerializer.serialize(error.message),
-              BidiSerializer.serialize(error.stack),
-            ],
-            awaitPromise: false,
-            target: {
-              realm: params.source.realm,
             },
-          });
+            error.name,
+            error.message,
+            error.stack
+          );
         } else {
-          await connection.send('script.callFunction', {
+          await dataHandle.evaluate(([, reject], error) => {
-            functionDeclaration: stringifyFunction(
-              (
-                [_, reject]: [unknown, (error: unknown) => void],
-                error: unknown
-              ) => {
             reject(error);
-              }
+          }, error);
-            ),
-            arguments: [
-              (await callbacks.reject.valueOrThrow()) as Bidi.Script.LocalValue,
-              BidiSerializer.serialize(error),
-            ],
-            awaitPromise: false,
-            target: {
-              realm: params.source.realm,
-            },
-          });
         }
       } catch (error) {
         debugError(error);
       }
-    }
-  };
-
-  get #connection(): Connection {
-    return this.#frame.page().browser().connection;
-  }
-
-  get #channelArguments() {
-    return [
-      {
-        type: 'channel' as const,
-        value: {
-          channel: this.#channels.args,
-          ownership: Bidi.Script.ResultOwnership.Root,
-        },
-      },
-      {
-        type: 'channel' as const,
-        value: {
-          channel: this.#channels.resolve,
-          ownership: Bidi.Script.ResultOwnership.Root,
-        },
-      },
-      {
-        type: 'channel' as const,
-        value: {
-          channel: this.#channels.reject,
-          ownership: Bidi.Script.ResultOwnership.Root,
-        },
-      },
-    ];
-  }
-
-  #handleResolveMessage = (params: Bidi.Script.MessageParameters) => {
-    if (params.channel !== this.#channels.resolve) {
       return;
     }
-    const {callbacks, remoteValue} = this.#getCallbacksAndRemoteValue(params);
+
-    callbacks.resolve.resolve(remoteValue);
+    try {
+      await dataHandle.evaluate(([resolve], result) => {
+        resolve(result);
+      }, result);
+    } catch (error) {
+      debugError(error);
+    }
   };
 
-  #handleRejectMessage = (params: Bidi.Script.MessageParameters) => {
+  #getRealm(source: Bidi.Script.Source) {
-    if (params.channel !== this.#channels.reject) {
+    const frame = this.#findFrame(source.context as string);
+    if (!frame) {
+      // Unrelated message.
       return;
     }
-    const {callbacks, remoteValue} = this.#getCallbacksAndRemoteValue(params);
+    return frame.realm(source.realm);
-    callbacks.reject.resolve(remoteValue);
-  };
-
-  #getCallbacksAndRemoteValue(params: Bidi.Script.MessageParameters) {
-    const {data, source} = params;
-    assert(data.type === 'array');
-    assert(data.value);
-
-    const callerIdRemote = data.value[0];
-    assert(callerIdRemote);
-    assert(callerIdRemote.type === 'number');
-    assert(typeof callerIdRemote.value === 'number');
-
-    let bindingMap = this.#callerInfos.get(source.realm);
-    if (!bindingMap) {
-      bindingMap = new Map();
-      this.#callerInfos.set(source.realm, bindingMap);
   }
 
-    const callerId = callerIdRemote.value;
+  #findFrame(id: string) {
-    let callbacks = bindingMap.get(callerId);
+    const frames = [this.#frame];
-    if (!callbacks) {
+    for (const frame of frames) {
-      callbacks = {
+      if (frame._id === id) {
-        resolve: new Deferred(),
+        return frame;
-        reject: new Deferred(),
-      };
-      bindingMap.set(callerId, callbacks);
       }
-    return {callbacks, remoteValue: data};
+      frames.push(...frame.childFrames());
+    }
+    return;
   }
 
   [Symbol.dispose](): void {
@@ -286,10 +230,11 @@ export class ExposeableFunction<Args extends unknown[], Ret> {
   }
 
   async [Symbol.asyncDispose](): Promise<void> {
-    if (this.#preloadScriptId) {
+    this.#disposables.dispose();
-      await this.#connection.send('script.removePreloadScript', {
+    await Promise.all(
-        script: this.#preloadScriptId,
+      this.#scripts.map(async ([frame, script]) => {
-      });
+        await frame.browsingContext.removePreloadScript(script);
-    }
+      })
+    );
   }
 }

packages/puppeteer-core/src/bidi/Frame.ts

@@ -224,14 +224,23 @@ export class BidiFrame extends Frame {
     return this.page()._timeoutSettings;
   }
 
-  override mainRealm(): BidiRealm {
+  override mainRealm(): BidiFrameRealm {
     return this.realms.default;
   }
 
-  override isolatedRealm(): BidiRealm {
+  override isolatedRealm(): BidiFrameRealm {
     return this.realms.internal;
   }
 
+  realm(id: string): BidiRealm | undefined {
+    for (const realm of Object.values(this.realms)) {
+      if (realm.realm.id === id) {
+        return realm;
+      }
+    }
+    return;
+  }
+
   override page(): BidiPage {
     let parent = this.#parent;
     while (parent instanceof BidiFrame) {

test/TestExpectations.json

@@ -2871,12 +2871,6 @@
     "parameters": ["cdp", "firefox"],
     "expectations": ["SKIP"]
   },
-  {
-    "testIdPattern": "[page.spec] Page Page.exposeFunction should be callable from-inside evaluateOnNewDocument",
-    "platforms": ["darwin", "linux", "win32"],
-    "parameters": ["chrome", "webDriverBiDi"],
-    "expectations": ["FAIL"]
-  },
   {
     "testIdPattern": "[page.spec] Page Page.metrics metrics event fired on console.timeStamp",
     "platforms": ["darwin", "linux", "win32"],

test/src/page.spec.ts

@@ -15,6 +15,7 @@ import type {HTTPRequest} from 'puppeteer-core/internal/api/HTTPRequest.js';
 import type {Metrics, Page} from 'puppeteer-core/internal/api/Page.js';
 import type {CdpPage} from 'puppeteer-core/internal/cdp/Page.js';
 import type {ConsoleMessage} from 'puppeteer-core/internal/common/ConsoleMessage.js';
+import {Deferred} from 'puppeteer-core/internal/util/Deferred.js';
 import sinon from 'sinon';
 
 import {getTestState, setupTestBrowserHooks} from './mocha-utils.js';
@@ -1030,15 +1031,15 @@ describe('Page', function () {
     it('should be callable from-inside evaluateOnNewDocument', async () => {
       const {page} = await getTestState();
 
-      let called = false;
+      const called = new Deferred<void>();
       await page.exposeFunction('woof', function () {
-        called = true;
+        called.resolve();
       });
       await page.evaluateOnNewDocument(() => {
         return (globalThis as any).woof();
       });
       await page.reload();
-      expect(called).toBe(true);
+      await called.valueOrThrow();
     });
     it('should survive navigation', async () => {
       const {page, server} = await getTestState();