chore: add internal method for utility context bindings (#4566)

* internal binding extracted from dnd patch

* refactor it into the page

* dgozman comments 1

src/server/browserContext.ts

@@ -184,14 +184,15 @@ export abstract class BrowserContext extends EventEmitter {
   }
 
   async exposeBinding(name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource): Promise<void> {
+    const identifier = PageBinding.identifier(name, 'main');
+    if (this._pageBindings.has(identifier))
+      throw new Error(`Function "${name}" has been already registered`);
     for (const page of this.pages()) {
-      if (page._pageBindings.has(name))
+      if (page.getBinding(name, 'main'))
         throw new Error(`Function "${name}" has been already registered in one of the pages`);
     }
-    if (this._pageBindings.has(name))
+    const binding = new PageBinding(name, playwrightBinding, needsHandle, 'main');
-      throw new Error(`Function "${name}" has been already registered`);
+    this._pageBindings.set(identifier, binding);
-    const binding = new PageBinding(name, playwrightBinding, needsHandle);
-    this._pageBindings.set(name, binding);
     this._doExposeBinding(binding);
   }
 

src/server/chromium/crPage.ts

@@ -198,8 +198,8 @@ export class CRPage implements PageDelegate {
     return this._go(+1);
   }
 
-  async evaluateOnNewDocument(source: string): Promise<void> {
+  async evaluateOnNewDocument(source: string, world: types.World = 'main'): Promise<void> {
-    await this._forAllFrameSessions(frame => frame._evaluateOnNewDocument(source));
+    await this._forAllFrameSessions(frame => frame._evaluateOnNewDocument(source, world));
   }
 
   async closePage(runBeforeUnload: boolean): Promise<void> {
@@ -406,7 +406,7 @@ class FrameSession {
             worldName: UTILITY_WORLD_NAME,
           });
           for (const binding of this._crPage._browserContext._pageBindings.values())
-            frame._evaluateExpression(binding.source, false, {}).catch(e => {});
+            frame._evaluateExpression(binding.source, false, {}, binding.world).catch(e => {});
         }
         const isInitialEmptyPage = this._isMainFrame() && this._page.mainFrame().url() === ':';
         if (isInitialEmptyPage) {
@@ -455,14 +455,12 @@ class FrameSession {
     promises.push(this._updateOffline(true));
     promises.push(this._updateHttpCredentials(true));
     promises.push(this._updateEmulateMedia(true));
-    for (const binding of this._crPage._browserContext._pageBindings.values())
+    for (const binding of this._crPage._page.allBindings())
-      promises.push(this._initBinding(binding));
-    for (const binding of this._crPage._page._pageBindings.values())
       promises.push(this._initBinding(binding));
     for (const source of this._crPage._browserContext._evaluateOnNewDocumentSources)
-      promises.push(this._evaluateOnNewDocument(source));
+      promises.push(this._evaluateOnNewDocument(source, 'main'));
     for (const source of this._crPage._page._evaluateOnNewDocumentSources)
-      promises.push(this._evaluateOnNewDocument(source));
+      promises.push(this._evaluateOnNewDocument(source, 'main'));
     if (this._isMainFrame() && this._crPage._browserContext._options.recordVideo) {
       const size = this._crPage._browserContext._options.recordVideo.size || this._crPage._browserContext._options.viewport || { width: 1280, height: 720 };
       const screencastId = createGuid();
@@ -565,11 +563,14 @@ class FrameSession {
     if (!frame)
       return;
     const delegate = new CRExecutionContext(this._client, contextPayload);
-    const context = new dom.FrameExecutionContext(delegate, frame);
+    let worldName: types.World|null = null;
     if (contextPayload.auxData && !!contextPayload.auxData.isDefault)
-      frame._contextCreated('main', context);
+      worldName = 'main';
     else if (contextPayload.name === UTILITY_WORLD_NAME)
-      frame._contextCreated('utility', context);
+      worldName = 'utility';
+    const context = new dom.FrameExecutionContext(delegate, frame, worldName);
+    if (worldName)
+      frame._contextCreated(worldName, context);
     this._contextIdToContext.set(contextPayload.id, context);
   }
 
@@ -683,9 +684,10 @@ class FrameSession {
   }
 
   async _initBinding(binding: PageBinding) {
+    const worldName = binding.world === 'utility' ? UTILITY_WORLD_NAME : undefined;
     await Promise.all([
-      this._client.send('Runtime.addBinding', { name: binding.name }),
+      this._client.send('Runtime.addBinding', { name: binding.name, executionContextName: worldName }),
-      this._client.send('Page.addScriptToEvaluateOnNewDocument', { source: binding.source })
+      this._client.send('Page.addScriptToEvaluateOnNewDocument', { source: binding.source, worldName })
     ]);
   }
 
@@ -693,7 +695,7 @@ class FrameSession {
     const context = this._contextIdToContext.get(event.executionContextId)!;
     const pageOrError = await this._crPage.pageOrError();
     if (!(pageOrError instanceof Error))
-      this._page._onBindingCalled(event.payload, context);
+      await this._page._onBindingCalled(event.payload, context);
   }
 
   _onDialog(event: Protocol.Page.javascriptDialogOpeningPayload) {
@@ -877,8 +879,9 @@ class FrameSession {
     await this._client.send('Page.setInterceptFileChooserDialog', { enabled }).catch(e => {}); // target can be closed.
   }
 
-  async _evaluateOnNewDocument(source: string): Promise<void> {
+  async _evaluateOnNewDocument(source: string, world: types.World): Promise<void> {
-    await this._client.send('Page.addScriptToEvaluateOnNewDocument', { source });
+    const worldName = world === 'utility' ? UTILITY_WORLD_NAME : undefined;
+    await this._client.send('Page.addScriptToEvaluateOnNewDocument', { source, worldName });
   }
 
   async _getContentFrame(handle: dom.ElementHandle): Promise<frames.Frame | null> {

src/server/dom.ts

@@ -28,11 +28,12 @@ import { FatalDOMError, RetargetableDOMError } from './common/domErrors';
 export class FrameExecutionContext extends js.ExecutionContext {
   readonly frame: frames.Frame;
   private _injectedScriptPromise?: Promise<js.JSHandle>;
-  private _debugScriptPromise?: Promise<js.JSHandle | undefined>;
+  readonly world: types.World | null;
 
-  constructor(delegate: js.ExecutionContextDelegate, frame: frames.Frame) {
+  constructor(delegate: js.ExecutionContextDelegate, frame: frames.Frame, world: types.World|null) {
     super(delegate);
     this.frame = frame;
+    this.world = world;
   }
 
   adoptIfNeeded(handle: js.JSHandle): Promise<js.JSHandle> | null {

src/server/firefox/ffBrowser.ts

@@ -298,6 +298,8 @@ export class FFBrowserContext extends BrowserContext {
   }
 
   async _doExposeBinding(binding: PageBinding) {
+    if (binding.world !== 'main')
+      throw new Error('Only main context bindings are supported in Firefox.');
     await this._browser._connection.send('Browser.addBinding', { browserContextId: this._browserContextId, name: binding.name, script: binding.source });
   }
 

src/server/firefox/ffPage.ts

@@ -134,11 +134,14 @@ export class FFPage implements PageDelegate {
     if (!frame)
       return;
     const delegate = new FFExecutionContext(this._session, executionContextId);
-    const context = new dom.FrameExecutionContext(delegate, frame);
+    let worldName: types.World|null = null;
     if (auxData.name === UTILITY_WORLD_NAME)
-      frame._contextCreated('utility', context);
+      worldName = 'utility';
     else if (!auxData.name)
-      frame._contextCreated('main', context);
+      worldName = 'main';
+    const context = new dom.FrameExecutionContext(delegate, frame, worldName);
+    if (worldName)
+      frame._contextCreated(worldName, context);
     this._contextIdToContext.set(executionContextId, context);
   }
 
@@ -290,6 +293,8 @@ export class FFPage implements PageDelegate {
   }
 
   async exposeBinding(binding: PageBinding) {
+    if (binding.world !== 'main')
+      throw new Error('Only main context bindings are supported in Firefox.');
     await this._session.send('Page.addBinding', { name: binding.name, script: binding.source });
   }
 

src/server/page.ts

@@ -133,7 +133,7 @@ export class Page extends EventEmitter {
   readonly _timeoutSettings: TimeoutSettings;
   readonly _delegate: PageDelegate;
   readonly _state: PageState;
-  readonly _pageBindings = new Map<string, PageBinding>();
+  private readonly _pageBindings = new Map<string, PageBinding>();
   readonly _evaluateOnNewDocumentSources: string[] = [];
   readonly _screenshotter: Screenshotter;
   readonly _frameManager: frames.FrameManager;
@@ -258,12 +258,13 @@ export class Page extends EventEmitter {
   }
 
   async exposeBinding(name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource) {
-    if (this._pageBindings.has(name))
+    const identifier = PageBinding.identifier(name, 'main');
+    if (this._pageBindings.has(identifier))
       throw new Error(`Function "${name}" has been already registered`);
-    if (this._browserContext._pageBindings.has(name))
+    if (this._browserContext._pageBindings.has(identifier))
       throw new Error(`Function "${name}" has been already registered in the browser context`);
-    const binding = new PageBinding(name, playwrightBinding, needsHandle);
+    const binding = new PageBinding(name, playwrightBinding, needsHandle, 'main');
-    this._pageBindings.set(name, binding);
+    this._pageBindings.set(identifier, binding);
     await this._delegate.exposeBinding(binding);
   }
 
@@ -462,6 +463,15 @@ export class Page extends EventEmitter {
       return;
     this._browserContext.addVisitedOrigin(new URL(url).origin);
   }
+
+  allBindings() {
+    return [...this._browserContext._pageBindings.values(), ...this._pageBindings.values()];
+  }
+
+  getBinding(name: string, world: types.World) {
+    const identifier = PageBinding.identifier(name, world);
+    return this._pageBindings.get(identifier) || this._browserContext._pageBindings.get(identifier);
+  }
 }
 
 export class Worker extends EventEmitter {
@@ -504,26 +514,31 @@ export class PageBinding {
   readonly playwrightFunction: frames.FunctionWithSource;
   readonly source: string;
   readonly needsHandle: boolean;
+  readonly world: types.World;
 
-  constructor(name: string, playwrightFunction: frames.FunctionWithSource, needsHandle: boolean) {
+  constructor(name: string, playwrightFunction: frames.FunctionWithSource, needsHandle: boolean, world: types.World) {
     this.name = name;
     this.playwrightFunction = playwrightFunction;
     this.source = `(${addPageBinding.toString()})(${JSON.stringify(name)}, ${needsHandle})`;
     this.needsHandle = needsHandle;
+    this.world = world;
+  }
+
+  static identifier(name: string, world: types.World) {
+    return world + ':' + name;
   }
 
   static async dispatch(page: Page, payload: string, context: dom.FrameExecutionContext) {
     const {name, seq, args} = JSON.parse(payload);
     try {
-      let binding = page._pageBindings.get(name);
+      assert(context.world);
-      if (!binding)
+      const binding = page.getBinding(name, context.world)!;
-        binding = page._browserContext._pageBindings.get(name);
       let result: any;
-      if (binding!.needsHandle) {
+      if (binding.needsHandle) {
         const handle = await context.evaluateHandleInternal(takeHandle, { name, seq }).catch(e => null);
-        result = await binding!.playwrightFunction({ frame: context.frame, page, context: page._browserContext }, handle);
+        result = await binding.playwrightFunction({ frame: context.frame, page, context: page._browserContext }, handle);
       } else {
-        result = await binding!.playwrightFunction({ frame: context.frame, page, context: page._browserContext }, ...args);
+        result = await binding.playwrightFunction({ frame: context.frame, page, context: page._browserContext }, ...args);
       }
       context.evaluateInternal(deliverResult, { name, seq, result }).catch(e => debugLogger.log('error', e));
     } catch (error) {

src/server/webkit/wkPage.ts

@@ -454,11 +454,14 @@ export class WKPage implements PageDelegate {
     if (!frame)
       return;
     const delegate = new WKExecutionContext(this._session, contextPayload.id);
-    const context = new dom.FrameExecutionContext(delegate, frame);
+    let worldName: types.World|null = null;
     if (contextPayload.type === 'normal')
-      frame._contextCreated('main', context);
+      worldName = 'main';
     else if (contextPayload.type === 'user' && contextPayload.name === UTILITY_WORLD_NAME)
-      frame._contextCreated('utility', context);
+      worldName = 'utility';
+    const context = new dom.FrameExecutionContext(delegate, frame, worldName);
+    if (worldName)
+      frame._contextCreated(worldName, context);
     if (contextPayload.type === 'normal' && frame === this._page.mainFrame())
       this._mainFrameContextId = contextPayload.id;
     this._contextIdToContext.set(contextPayload.id, context);
@@ -679,11 +682,15 @@ export class WKPage implements PageDelegate {
   }
 
   async exposeBinding(binding: PageBinding): Promise<void> {
+    if (binding.world !== 'main')
+      throw new Error('Only main context bindings are supported in WebKit.');
     await this._updateBootstrapScript();
     await this._evaluateBindingScript(binding);
   }
 
   private async _evaluateBindingScript(binding: PageBinding): Promise<void> {
+    if (binding.world !== 'main')
+      throw new Error('Only main context bindings are supported in WebKit.');
     const script = this._bindingToScript(binding);
     await Promise.all(this._page.frames().map(frame => frame._evaluateExpression(script, false, {}).catch(e => {})));
   }
@@ -698,9 +705,7 @@ export class WKPage implements PageDelegate {
 
   private _calculateBootstrapScript(): string {
     const scripts: string[] = [];
-    for (const binding of this._browserContext._pageBindings.values())
+    for (const binding of this._page.allBindings())
-      scripts.push(this._bindingToScript(binding));
-    for (const binding of this._page._pageBindings.values())
       scripts.push(this._bindingToScript(binding));
     scripts.push(...this._browserContext._evaluateOnNewDocumentSources);
     scripts.push(...this._page._evaluateOnNewDocumentSources);