From ea833daa893f754027d685ceace61b4587114480 Mon Sep 17 00:00:00 2001
From: Joel Einbinder <joel.einbinder@gmail.com>
Date: Tue, 8 Dec 2020 08:38:29 -0800
Subject: [PATCH] chore: fix internal binding (#4598)

---
 src/server/chromium/crPage.ts     |  2 +-
 src/server/dom.ts                 |  2 +-
 src/server/page.ts                |  6 +++---
 test/page-expose-function.spec.ts | 31 +++++++++++++++++++++++++++++++
 4 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/src/server/chromium/crPage.ts b/src/server/chromium/crPage.ts
index aabadb188a..9cf91bf5d9 100644
--- a/src/server/chromium/crPage.ts
+++ b/src/server/chromium/crPage.ts
@@ -128,7 +128,7 @@ export class CRPage implements PageDelegate {
 
   async exposeBinding(binding: PageBinding) {
     await this._forAllFrameSessions(frame => frame._initBinding(binding));
-    await Promise.all(this._page.frames().map(frame => frame._evaluateExpression(binding.source, false, {}).catch(e => {})));
+    await Promise.all(this._page.frames().map(frame => frame._evaluateExpression(binding.source, false, {}, binding.world).catch(e => {})));
   }
 
   async updateExtraHTTPHeaders(): Promise<void> {
diff --git a/src/server/dom.ts b/src/server/dom.ts
index 384e194291..ef84753916 100644
--- a/src/server/dom.ts
+++ b/src/server/dom.ts
@@ -879,7 +879,7 @@ export function waitForSelectorTask(selector: SelectorInfo, state: 'attached' |
 
 export function dispatchEventTask(selector: SelectorInfo, type: string, eventInit: Object): SchedulableTask<undefined> {
   return injectedScript => injectedScript.evaluateHandle((injected, { parsed, type, eventInit }) => {
-    return injected.pollRaf((progress, continuePolling) => {
+    return injected.pollRaf<undefined>((progress, continuePolling) => {
       const element = injected.querySelector(parsed, document);
       if (!element)
         return continuePolling;
diff --git a/src/server/page.ts b/src/server/page.ts
index a63f8c87f2..6163838a71 100644
--- a/src/server/page.ts
+++ b/src/server/page.ts
@@ -257,13 +257,13 @@ export class Page extends EventEmitter {
     this._timeoutSettings.setDefaultTimeout(timeout);
   }
 
-  async exposeBinding(name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource) {
-    const identifier = PageBinding.identifier(name, 'main');
+  async exposeBinding(name: string, needsHandle: boolean, playwrightBinding: frames.FunctionWithSource, world: types.World = 'main') {
+    const identifier = PageBinding.identifier(name, world);
     if (this._pageBindings.has(identifier))
       throw new Error(`Function "${name}" has been already registered`);
     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, 'main');
+    const binding = new PageBinding(name, playwrightBinding, needsHandle, world);
     this._pageBindings.set(identifier, binding);
     await this._delegate.exposeBinding(binding);
   }
diff --git a/test/page-expose-function.spec.ts b/test/page-expose-function.spec.ts
index a8fc016991..4a28ac9251 100644
--- a/test/page-expose-function.spec.ts
+++ b/test/page-expose-function.spec.ts
@@ -16,6 +16,7 @@
  */
 
 import { it, expect } from './fixtures';
+import { attachFrame } from './utils';
 
 it('exposeBinding should work', async ({browser}) => {
   const context = await browser.newContext();
@@ -236,3 +237,33 @@ it('should not result in unhandled rejection', async ({page}) => {
   // Make a round-trip to be sure we did not throw immediately after closing.
   expect(await page.evaluate('1 + 1').catch(e => e)).toBeInstanceOf(Error);
 });
+
+it('should work with internal bindings', (test, { mode, browserName }) => {
+  test.skip(mode !== 'default');
+  test.skip(browserName !== 'chromium');
+}, async ({page, toImpl, server}) => {
+  const implPage: import('../src/server/page').Page = toImpl(page);
+  let foo;
+  await implPage.exposeBinding('foo', false, ({}, arg) => {
+    foo = arg;
+  }, 'utility');
+  expect(await page.evaluate('!!window.foo')).toBe(false);
+  expect(await implPage.mainFrame()._evaluateExpression('!!window.foo', false, {}, 'utility')).toBe(true);
+  expect(foo).toBe(undefined);
+  await implPage.mainFrame()._evaluateExpression('window.foo(123)', false, {}, 'utility');
+  expect(foo).toBe(123);
+
+  // should work after reload
+  await page.goto(server.EMPTY_PAGE);
+  expect(await page.evaluate('!!window.foo')).toBe(false);
+  await implPage.mainFrame()._evaluateExpression('window.foo(456)', false, {}, 'utility');
+  expect(foo).toBe(456);
+
+  // should work inside frames
+  const frame = await attachFrame(page, 'myframe', server.CROSS_PROCESS_PREFIX + '/empty.html');
+  expect(await frame.evaluate('!!window.foo')).toBe(false);
+  const implFrame: import('../src/server/frames').Frame = toImpl(frame);
+  await implFrame._evaluateExpression('window.foo(789)', false, {}, 'utility');
+  expect(foo).toBe(789);
+});
+