feat(addInitScript): allow passing functions as init script arguments (#41921)

Author: pavelfeldman

docs/src/api/class-browsercontext.md

@@ -476,6 +476,9 @@ Path to the JavaScript file. If `path` is a relative path, then it is resolved r
 
 Script to be evaluated in all pages in the browser context. Optional.
 
+### option: BrowserContext.addInitScript.exposeFunctions = %%-js-init-script-expose-functions-%%
+* since: v1.62
+
 ## method: BrowserContext.backgroundPages
 * since: v1.11
 * deprecated: Background pages have been removed from Chromium together with Manifest V2 extensions.

docs/src/api/class-page.md

@@ -651,6 +651,9 @@ Path to the JavaScript file. If `path` is a relative path, then it is resolved r
 
 Script to be evaluated in all pages in the browser context. Optional.
 
+### option: Page.addInitScript.exposeFunctions = %%-js-init-script-expose-functions-%%
+* since: v1.62
+
 ## async method: Page.addScriptTag
 * since: v1.8
 - returns: <[ElementHandle]>

docs/src/api/params.md

@@ -590,6 +590,12 @@ Function to be evaluated in the page context.
 
 When set to `true`, functions passed inside [`param: arg`] are exposed in the page and can be called from the page function. Calling one returns a [Promise] of its result. Under the hood, each function is exposed via [`method: Page.exposeFunction`], so it is technically accessible from all frames and worlds of the page. Exposed functions are cleared upon the top-level navigation. Defaults to `false`, in which case functions are not serializable and passing one throws an error.
 
+## js-init-script-expose-functions
+* langs: js
+- `exposeFunctions` <[boolean]>
+
+When set to `true`, functions passed inside [`param: arg`] are exposed in the page and can be called from the init script. Calling one returns a [Promise] of its result. Under the hood, each function is exposed via [`method: Page.exposeFunction`], so it is technically accessible from all frames and worlds of the page. Unlike functions passed to [`method: Page.evaluate`], functions passed to an init script are exposed in every new document, so they survive navigations. Defaults to `false`, in which case functions are not serializable and are silently dropped.
+
 ## js-evalonselector-pagefunction
 * langs: js
 - `pageFunction` <[function]\([Element]\)|[string]>

packages/injected/src/bindingsController.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import { serializeAsCallArgument } from '@isomorphic/utilityScriptSerializers';
+import { parseEvaluationResultValue, serializeAsCallArgument } from '@isomorphic/utilityScriptSerializers';
 
 import type { SerializedValue } from '@isomorphic/utilityScriptSerializers';
 
@@ -68,6 +68,12 @@ export class BindingsController {
     return promise;
   }
 
+  parseInitScriptArg(value: SerializedValue): any {
+    // Functions serialized as { fn } deserialize into wrappers
+    // that route the call through this controller.
+    return parseEvaluationResultValue(value);
+  }
+
   removeBinding(bindingName: string) {
     const data = this._bindings.get(bindingName);
     if (data)

packages/isomorphic/utilityScriptSerializers.ts

@@ -39,7 +39,7 @@ export type SerializedValue =
     { ta: { b: string, k: TypedArrayKind } } |
     { ab: { b: string } };
 
-type HandleOrValue = { h: number } | { fallThrough: any };
+type HandleOrValue = { h: number } | { fn: string } | { fallThrough: any };
 
 type VisitorInfo = {
   visited: Map<object, number>;

packages/playwright-client/types/types.d.ts

@@ -314,16 +314,17 @@ export interface Page {
    * ```
    *
    * **NOTE** The order of evaluation of multiple scripts installed via
-   * [browserContext.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script)
-   * and [page.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-page#page-add-init-script) is not
-   * defined.
+   * [browserContext.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script)
+   * and [page.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-page#page-add-init-script)
+   * is not defined.
    *
    * @param script Script to be evaluated in the page.
    * @param arg Optional argument to pass to
    * [`script`](https://playwright.dev/docs/api/class-page#page-add-init-script-option-script) (only supported when
    * passing a function).
+   * @param options
    */
-  addInitScript<Arg>(script: PageFunction<Arg, any> | { path?: string, content?: string }, arg?: Arg): Promise<Disposable>;
+  addInitScript<Arg>(script: PageFunction<Arg, any> | { path?: string, content?: string }, arg?: Arg, options?: { exposeFunctions?: boolean }): Promise<Disposable>;
 
   /**
    * **NOTE** Use locator-based [page.locator(selector[, options])](https://playwright.dev/docs/api/class-page#page-locator)
@@ -9090,16 +9091,17 @@ export interface BrowserContext {
    * ```
    *
    * **NOTE** The order of evaluation of multiple scripts installed via
-   * [browserContext.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script)
-   * and [page.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-page#page-add-init-script) is not
-   * defined.
+   * [browserContext.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script)
+   * and [page.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-page#page-add-init-script)
+   * is not defined.
    *
    * @param script Script to be evaluated in all pages in the browser context.
    * @param arg Optional argument to pass to
    * [`script`](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script-option-script)
    * (only supported when passing a function).
+   * @param options
    */
-  addInitScript<Arg>(script: PageFunction<Arg, any> | { path?: string, content?: string }, arg?: Arg): Promise<Disposable>;
+  addInitScript<Arg>(script: PageFunction<Arg, any> | { path?: string, content?: string }, arg?: Arg, options?: { exposeFunctions?: boolean }): Promise<Disposable>;
 
   /**
    * Removes all the listeners of the given type (or all registered listeners if no type given). Allows to wait for
@@ -20795,14 +20797,15 @@ export interface Dialog {
 /**
  * [Disposable](https://playwright.dev/docs/api/class-disposable) is returned from various methods to allow undoing
  * the corresponding action. For example,
- * [page.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-page#page-add-init-script) returns a
- * [Disposable](https://playwright.dev/docs/api/class-disposable) that can be used to remove the init script.
+ * [page.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-page#page-add-init-script)
+ * returns a [Disposable](https://playwright.dev/docs/api/class-disposable) that can be used to remove the init
+ * script.
  */
 export interface Disposable {
   /**
    * Removes the associated resource. For example, removes the init script installed via
-   * [page.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-page#page-add-init-script) or
-   * [browserContext.addInitScript(script[, arg])](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script).
+   * [page.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-page#page-add-init-script) or
+   * [browserContext.addInitScript(script[, arg, options])](https://playwright.dev/docs/api/class-browsercontext#browser-context-add-init-script).
    */
   dispose(): Promise<void>;
 

packages/playwright-core/src/client/browserContext.ts

@@ -37,15 +37,17 @@ import { Events } from './events';
 import { APIRequestContext } from './fetch';
 import { Frame } from './frame';
 import { HarRouter } from './harRouter';
+import { assertEvaluateOptions } from './jsHandle';
 import * as network from './network';
-import { BindingCall, Page } from './page';
+import { BindingCall, Page, addInitScriptWithExposedFunctions } from './page';
 import { Tracing } from './tracing';
 import { Waiter } from './waiter';
 import { WebError } from './webError';
 import { Worker } from './worker';
 import { TimeoutSettings, kNoTimeout } from './timeoutSettings';
 import { mkdirIfNeeded } from './fileUtils';
 
+import type { EvaluateOptions } from './jsHandle';
 import type { BrowserContextOptions, Headers, SetStorageState, StorageState, WaitForEventOptions } from './types';
 import type * as structs from '../../types/structs';
 import type * as api from '../../types/types';
@@ -358,7 +360,10 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel>
     await this._channel.setHTTPCredentials({ httpCredentials: httpCredentials || undefined }, kNoTimeout);
   }
 
-  async addInitScript(script: Function | string | { path?: string, content?: string }, arg?: any) {
+  async addInitScript(script: Function | string | { path?: string, content?: string }, arg?: any, options?: EvaluateOptions) {
+    assertEvaluateOptions(options);
+    if (options?.exposeFunctions)
+      return await addInitScriptWithExposedFunctions(this, script, arg);
     const source = await evaluationScript(script, arg);
     return DisposableObject.from((await this._channel.addInitScript({ source }, kNoTimeout)).disposable);
   }
@@ -369,6 +374,12 @@ export class BrowserContext extends ChannelOwner<channels.BrowserContextChannel>
     return DisposableObject.from(result.disposable);
   }
 
+  async _exposeCallbackBinding(name: string, callback: Function): Promise<DisposableObject> {
+    this._bindings.set(name, (source, ...args) => callback(...args));
+    const result = await this._channel.exposeBinding({ name, noGlobal: true }, kNoTimeout);
+    return DisposableObject.from(result.disposable);
+  }
+
   async exposeFunction(name: string, callback: Function): Promise<DisposableObject> {
     const result = await this._channel.exposeBinding({ name }, kNoTimeout);
     const binding = (source: structs.BindingSource, ...args: any[]) => callback(...args);

packages/playwright-core/src/client/channels.d.ts

@@ -1425,9 +1425,10 @@ export type BrowserContextCookiesResult = {
 };
 export type BrowserContextExposeBindingParams = {
   name: string,
+  noGlobal?: boolean,
 };
 export type BrowserContextExposeBindingOptions = {
-
+  noGlobal?: boolean,
 };
 export type BrowserContextExposeBindingResult = {
   disposable: DisposableChannel,

packages/playwright-core/src/client/clientHelper.ts

@@ -18,6 +18,8 @@
 import fs from 'fs';
 
 import { isString } from '@isomorphic/rtti';
+import { kBindingsControllerProperty, kFunctionBindingPrefix, serializeAsCallArgument } from '@isomorphic/utilityScriptSerializers';
+import { createGuid } from '@utils/crypto';
 
 export function envObjectToArray(env: NodeJS.ProcessEnv): { name: string, value: string }[] {
   const result: { name: string, value: string }[] = [];
@@ -49,6 +51,22 @@ export async function evaluationScript(fun: Function | string | { path?: string,
   throw new Error('Either path or content property must be present');
 }
 
+export async function initScriptSourceWithExposedFunctions(fun: Function, arg: any, expose: (name: string, callback: Function) => Promise<void>): Promise<string> {
+  const exposePromises: Promise<void>[] = [];
+  const serialized = serializeAsCallArgument(arg, value => {
+    if (typeof value === 'function') {
+      const name = kFunctionBindingPrefix + createGuid();
+      exposePromises.push(expose(name, value));
+      return { fn: name };
+    }
+    return { fallThrough: value };
+  });
+  await Promise.all(exposePromises);
+  // Bindings backing the functions are registered through their own init scripts
+  // that are guaranteed to run first, so the controller is available here.
+  return `(${fun.toString()})(globalThis['${kBindingsControllerProperty}'].parseInitScriptArg(${JSON.stringify(serialized)}))`;
+}
+
 export function addSourceUrlToScript(source: string, path: string): string {
   return `${source}\n//# sourceURL=${path.replace(/\n/g, '')}`;
 }

packages/playwright-core/src/client/page.ts

@@ -28,7 +28,7 @@ import { LongStandingScope } from '@isomorphic/manualPromise';
 import { isObject, isRegExp, isString } from '@isomorphic/rtti';
 import { Artifact } from './artifact';
 import { ChannelOwner } from './channelOwner';
-import { evaluationScript } from './clientHelper';
+import { evaluationScript, initScriptSourceWithExposedFunctions } from './clientHelper';
 import { Coverage } from './coverage';
 import { DisposableObject, DisposableStub } from './disposable';
 import { Download } from './download';
@@ -40,7 +40,7 @@ import { Frame, verifyLoadState } from './frame';
 import { HarRouter } from './harRouter';
 import { Keyboard, Mouse, Touchscreen } from './input';
 import { WebStorage } from './webStorage';
-import { assertMaxArguments, parseResult, serializeArgument } from './jsHandle';
+import { assertEvaluateOptions, assertMaxArguments, parseResult, serializeArgument } from './jsHandle';
 import { Request, Response, Route, RouteHandler, WebSocket,  WebSocketRoute, WebSocketRouteHandler, validateHeaders } from './network';
 import { Video } from './video';
 import { Screencast } from './screencast';
@@ -376,10 +376,15 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
     return DisposableObject.from(result.disposable);
   }
 
-  async _exposeEvaluateCallback(name: string, callback: Function) {
+  async _exposeCallbackBinding(name: string, callback: Function): Promise<DisposableObject> {
     this._bindings.set(name, (source, ...args) => callback(...args));
     const result = await this._channel.exposeBinding({ name, noGlobal: true }, kNoTimeout);
-    this._evaluateCallbacks.push({ name, disposable: DisposableObject.from(result.disposable) });
+    return DisposableObject.from(result.disposable);
+  }
+
+  async _exposeEvaluateCallback(name: string, callback: Function) {
+    const disposable = await this._exposeCallbackBinding(name, callback);
+    this._evaluateCallbacks.push({ name, disposable });
   }
 
   _eraseEvaluateCallbacks() {
@@ -549,7 +554,10 @@ export class Page extends ChannelOwner<channels.PageChannel> implements api.Page
     return await this._mainFrame.evaluate(pageFunction, arg, options);
   }
 
-  async addInitScript(script: Function | string | { path?: string, content?: string }, arg?: any) {
+  async addInitScript(script: Function | string | { path?: string, content?: string }, arg?: any, options?: EvaluateOptions) {
+    assertEvaluateOptions(options);
+    if (options?.exposeFunctions)
+      return await addInitScriptWithExposedFunctions(this, script, arg);
     const source = await evaluationScript(script, arg);
     return DisposableObject.from((await this._channel.addInitScript({ source }, kNoTimeout)).disposable);
   }
@@ -947,3 +955,23 @@ function trimUrl(param: any): string | undefined {
   if (isString(param))
     return `"${trimStringWithEllipsis(param, 50)}"`;
 }
+
+export async function addInitScriptWithExposedFunctions(owner: Page | BrowserContext, script: Function | string | { path?: string, content?: string }, arg: any): Promise<DisposableStub> {
+  if (typeof script !== 'function')
+    throw new Error('Passing functions requires the init script to be a function');
+  const callbacks: { name: string, disposable: DisposableObject }[] = [];
+  const source = await owner._wrapApiCall(async () => {
+    return await initScriptSourceWithExposedFunctions(script, arg, async (name, callback) => {
+      const disposable = await owner._exposeCallbackBinding(name, callback);
+      callbacks.push({ name, disposable });
+    });
+  }, { internal: true });
+  const initScriptDisposable = DisposableObject.from((await owner._channel.addInitScript({ source }, kNoTimeout)).disposable);
+  return new DisposableStub(async () => {
+    for (const { name, disposable } of callbacks) {
+      owner._bindings.delete(name);
+      disposable.dispose().catch(() => {});
+    }
+    await initScriptDisposable.dispose();
+  });
+}