From 2225f126cf94f756bc19987b0a5faf4cc27e0075 Mon Sep 17 00:00:00 2001
From: Pavel Feldman <pavel.feldman@gmail.com>
Date: Mon, 7 Sep 2020 11:37:22 -0700
Subject: [PATCH] chore: expose internal api for extending injected script

---
 src/client/frame.ts                |  5 +++++
 src/debug/debugController.ts       |  8 +-------
 src/dispatchers/frameDispatcher.ts |  4 ++++
 src/protocol/channels.ts           | 11 +++++++++++
 src/protocol/protocol.yml          |  8 ++++++++
 src/protocol/validator.ts          |  4 ++++
 src/server/dom.ts                  |  7 -------
 src/server/frames.ts               |  8 ++++++++
 8 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/src/client/frame.ts b/src/client/frame.ts
index e8e43f70adf448..24b68c88a3ff4b 100644
--- a/src/client/frame.ts
+++ b/src/client/frame.ts
@@ -410,6 +410,11 @@ export class Frame extends ChannelOwner<channels.FrameChannel, channels.FrameIni
       return (await this._channel.title()).value;
     });
   }
+
+  async _extendInjectedScript<Arg>(source: string, arg?: Arg): Promise<JSHandle> {
+    const result = await this._channel.extendInjectedScript({ source, arg: serializeArgument(arg) });
+    return JSHandle.from(result.handle);
+  }
 }
 
 export function verifyLoadState(name: string, waitUntil: LifecycleEvent): LifecycleEvent {
diff --git a/src/debug/debugController.ts b/src/debug/debugController.ts
index 156bd6bb922f3b..a4d827805cdff5 100644
--- a/src/debug/debugController.ts
+++ b/src/debug/debugController.ts
@@ -22,16 +22,10 @@ import { Progress } from '../server/progress';
 import { isDebugMode } from '../utils/utils';
 import * as debugScriptSource from '../generated/debugScriptSource';
 
-const debugScriptSymbol = Symbol('debugScript');
-
 export class DebugController implements InstrumentingAgent {
   private async ensureInstalledInFrame(frame: frames.Frame) {
     try {
-      const mainContext = await frame._mainContext();
-      if ((mainContext as any)[debugScriptSymbol])
-        return;
-      (mainContext as any)[debugScriptSymbol] = true;
-      await mainContext.extendInjectedScript(debugScriptSource.source);
+      await frame.extendInjectedScript(debugScriptSource.source);
     } catch (e) {
     }
   }
diff --git a/src/dispatchers/frameDispatcher.ts b/src/dispatchers/frameDispatcher.ts
index 52426c11613235..326d1f26741cda 100644
--- a/src/dispatchers/frameDispatcher.ts
+++ b/src/dispatchers/frameDispatcher.ts
@@ -178,4 +178,8 @@ export class FrameDispatcher extends Dispatcher<Frame, channels.FrameInitializer
   async title(): Promise<channels.FrameTitleResult> {
     return { value: await this._frame.title() };
   }
+
+  async extendInjectedScript(params: channels.FrameExtendInjectedScriptParams): Promise<channels.FrameExtendInjectedScriptResult> {
+    return { handle: createHandle(this._scope, await this._frame.extendInjectedScript(params.source, parseArgument(params.arg))) };
+  }
 }
diff --git a/src/protocol/channels.ts b/src/protocol/channels.ts
index 3fe9d0cec0c753..f4a8052b661fe8 100644
--- a/src/protocol/channels.ts
+++ b/src/protocol/channels.ts
@@ -1111,6 +1111,7 @@ export interface FrameChannel extends Channel {
   uncheck(params: FrameUncheckParams): Promise<FrameUncheckResult>;
   waitForFunction(params: FrameWaitForFunctionParams): Promise<FrameWaitForFunctionResult>;
   waitForSelector(params: FrameWaitForSelectorParams): Promise<FrameWaitForSelectorResult>;
+  extendInjectedScript(params: FrameExtendInjectedScriptParams): Promise<FrameExtendInjectedScriptResult>;
 }
 export type FrameLoadstateEvent = {
   add?: 'load' | 'domcontentloaded' | 'networkidle',
@@ -1510,6 +1511,16 @@ export type FrameWaitForSelectorOptions = {
 export type FrameWaitForSelectorResult = {
   element?: ElementHandleChannel,
 };
+export type FrameExtendInjectedScriptParams = {
+  source: string,
+  arg: SerializedArgument,
+};
+export type FrameExtendInjectedScriptOptions = {
+
+};
+export type FrameExtendInjectedScriptResult = {
+  handle: JSHandleChannel,
+};
 
 // ----------- Worker -----------
 export type WorkerInitializer = {
diff --git a/src/protocol/protocol.yml b/src/protocol/protocol.yml
index c298019170a063..eb59688f1e485a 100644
--- a/src/protocol/protocol.yml
+++ b/src/protocol/protocol.yml
@@ -1242,6 +1242,14 @@ Frame:
       returns:
         element: ElementHandle?
 
+    extendInjectedScript:
+      experimental: True
+      parameters:
+        source: string
+        arg: SerializedArgument
+      returns:
+        handle: JSHandle
+
   events:
 
     loadstate:
diff --git a/src/protocol/validator.ts b/src/protocol/validator.ts
index 892a24ff76bfe2..f6dffa12180fa0 100644
--- a/src/protocol/validator.ts
+++ b/src/protocol/validator.ts
@@ -604,6 +604,10 @@ export function createScheme(tChannel: (name: string) => Validator): Scheme {
     timeout: tOptional(tNumber),
     state: tOptional(tEnum(['attached', 'detached', 'visible', 'hidden'])),
   });
+  scheme.FrameExtendInjectedScriptParams = tObject({
+    source: tString,
+    arg: tType('SerializedArgument'),
+  });
   scheme.WorkerEvaluateExpressionParams = tObject({
     expression: tString,
     isFunction: tBoolean,
diff --git a/src/server/dom.ts b/src/server/dom.ts
index 5f2b55c33a7e8e..766ebdddd46b6f 100644
--- a/src/server/dom.ts
+++ b/src/server/dom.ts
@@ -90,13 +90,6 @@ export class FrameExecutionContext extends js.ExecutionContext {
     return this._injectedScriptPromise;
   }
 
-  async extendInjectedScript(source: string, params?: any) {
-    const injectedScriptHandle = await this.injectedScript();
-    return injectedScriptHandle.evaluate((injectedScript, {source, params}) => {
-      return injectedScript.extend(source, params);
-    }, { source, params });
-  }
-
   async doSlowMo() {
     return this.frame._page._doSlowMo();
   }
diff --git a/src/server/frames.ts b/src/server/frames.ts
index 636e42dcc73ac9..fcdcc6cdfc71cc 100644
--- a/src/server/frames.ts
+++ b/src/server/frames.ts
@@ -975,6 +975,14 @@ export class Frame extends EventEmitter {
       clearTimeout(this._networkIdleTimer);
     this._networkIdleTimer = undefined;
   }
+
+  async extendInjectedScript(source: string, arg?: any): Promise<js.JSHandle> {
+    const mainContext = await this._mainContext();
+    const injectedScriptHandle = await mainContext.injectedScript();
+    return injectedScriptHandle.evaluate((injectedScript, {source, arg}) => {
+      return injectedScript.extend(source, arg);
+    }, { source, arg });
+  }
 }
 
 class RerunnableTask {