From 627bd003996d5b29ba95a4ee0130c54a721fcba0 Mon Sep 17 00:00:00 2001
From: jrandolf <101637635+jrandolf@users.noreply.github.com>
Date: Mon, 26 Feb 2024 11:42:51 +0100
Subject: [PATCH] chore: implement `removeExposedFunction` (#11984)

---
 .../src/bidi/ExposedFunction.ts               | 44 ++++++++++++-------
 packages/puppeteer-core/src/bidi/Frame.ts     | 18 +++++---
 packages/puppeteer-core/src/bidi/Page.ts      |  5 +--
 test/TestExpectations.json                    |  6 ---
 4 files changed, 43 insertions(+), 30 deletions(-)

diff --git a/packages/puppeteer-core/src/bidi/ExposedFunction.ts b/packages/puppeteer-core/src/bidi/ExposedFunction.ts
index 6a6f42656c3..f6e1304a550 100644
--- a/packages/puppeteer-core/src/bidi/ExposedFunction.ts
+++ b/packages/puppeteer-core/src/bidi/ExposedFunction.ts
@@ -29,6 +29,17 @@ type CallbackChannel<Args, Ret> = (
  * @internal
  */
 export class ExposeableFunction<Args extends unknown[], Ret> {
+  static async from<Args extends unknown[], Ret>(
+    frame: BidiFrame,
+    name: string,
+    apply: (...args: Args) => Awaitable<Ret>,
+    isolate = false
+  ): Promise<ExposeableFunction<Args, Ret>> {
+    const func = new ExposeableFunction(frame, name, apply, isolate);
+    await func.#initialize();
+    return func;
+  }
+
   readonly #frame;
 
   readonly name;
@@ -54,7 +65,7 @@ export class ExposeableFunction<Args extends unknown[], Ret> {
     this.#channel = `__puppeteer__${this.#frame._id}_page_exposeFunction_${this.name}`;
   }
 
-  async expose(): Promise<void> {
+  async #initialize() {
     const connection = this.#connection;
     const channel = {
       type: 'channel' as const,
@@ -98,23 +109,16 @@ export class ExposeableFunction<Args extends unknown[], Ret> {
       frames.map(async frame => {
         const realm = this.#isolate ? frame.isolatedRealm() : frame.mainRealm();
         try {
-          this.#scripts.push([
-            frame,
-            await frame.browsingContext.addPreloadScript(functionDeclaration, {
+          const [script] = await Promise.all([
+            frame.browsingContext.addPreloadScript(functionDeclaration, {
               arguments: [channel],
               sandbox: realm.sandbox,
             }),
+            realm.realm.callFunction(functionDeclaration, false, {
+              arguments: [channel],
+            }),
           ]);
-        } 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],
-          });
+          this.#scripts.push([frame, script]);
         } catch (error) {
           // If it errors, the frame probably doesn't support call function. We
           // fail gracefully.
@@ -233,7 +237,17 @@ export class ExposeableFunction<Args extends unknown[], Ret> {
     this.#disposables.dispose();
     await Promise.all(
       this.#scripts.map(async ([frame, script]) => {
-        await frame.browsingContext.removePreloadScript(script);
+        const realm = this.#isolate ? frame.isolatedRealm() : frame.mainRealm();
+        try {
+          await Promise.all([
+            realm.evaluate(name => {
+              delete (globalThis as any)[name];
+            }, this.name),
+            frame.browsingContext.removePreloadScript(script),
+          ]);
+        } catch (error) {
+          debugError(error);
+        }
       })
     );
   }
diff --git a/packages/puppeteer-core/src/bidi/Frame.ts b/packages/puppeteer-core/src/bidi/Frame.ts
index 1320702c9cc..0eb3e34df75 100644
--- a/packages/puppeteer-core/src/bidi/Frame.ts
+++ b/packages/puppeteer-core/src/bidi/Frame.ts
@@ -408,14 +408,20 @@ export class BidiFrame extends Frame {
         `Failed to add page binding with name ${name}: globalThis['${name}'] already exists!`
       );
     }
-    const exposeable = new ExposeableFunction(this, name, apply);
+    const exposeable = await ExposeableFunction.from(this, name, apply);
     this.#exposedFunctions.set(name, exposeable);
-    try {
-      await exposeable.expose();
-    } catch (error) {
-      this.#exposedFunctions.delete(name);
-      throw error;
+  }
+
+  async removeExposedFunction(name: string): Promise<void> {
+    const exposedFunction = this.#exposedFunctions.get(name);
+    if (!exposedFunction) {
+      throw new Error(
+        `Failed to remove page binding with name ${name}: window['${name}'] does not exists!`
+      );
     }
+
+    this.#exposedFunctions.delete(name);
+    await exposedFunction[Symbol.asyncDispose]();
   }
 
   override waitForSelector<Selector extends string>(
diff --git a/packages/puppeteer-core/src/bidi/Page.ts b/packages/puppeteer-core/src/bidi/Page.ts
index a69c7720db3..4e4279ecb61 100644
--- a/packages/puppeteer-core/src/bidi/Page.ts
+++ b/packages/puppeteer-core/src/bidi/Page.ts
@@ -577,9 +577,8 @@ export class BidiPage extends Page {
     throw new UnsupportedOperation();
   }
 
-  override removeExposedFunction(): never {
-    // TODO: Quick win?
-    throw new UnsupportedOperation();
+  override async removeExposedFunction(name: string): Promise<void> {
+    await this.#frame.removeExposedFunction(name);
   }
 
   override authenticate(): never {
diff --git a/test/TestExpectations.json b/test/TestExpectations.json
index 4d1368f1333..33beaadf343 100644
--- a/test/TestExpectations.json
+++ b/test/TestExpectations.json
@@ -674,12 +674,6 @@
     "parameters": ["firefox"],
     "expectations": ["SKIP"]
   },
-  {
-    "testIdPattern": "[page.spec] Page Page.removeExposedFunction should work",
-    "platforms": ["darwin", "linux", "win32"],
-    "parameters": ["webDriverBiDi"],
-    "expectations": ["SKIP"]
-  },
   {
     "testIdPattern": "[page.spec] Page Page.setBypassCSP *",
     "platforms": ["darwin", "linux", "win32"],