chore: add injection framework (#8862)

* chore: add injection framework

src/common/Frame.ts

@@ -214,7 +214,7 @@ export class Frame {
     this.#client = client;
     this.worlds = {
       [MAIN_WORLD]: new IsolatedWorld(this),
-      [PUPPETEER_WORLD]: new IsolatedWorld(this),
+      [PUPPETEER_WORLD]: new IsolatedWorld(this, true),
     };
   }
 

src/common/IsolatedWorld.ts

@@ -15,11 +15,9 @@
  */
 
 import {Protocol} from 'devtools-protocol';
+import {source as injectedSource} from '../generated/injected.js';
 import {assert} from '../util/assert.js';
-import {
-  createDeferredPromise,
-  DeferredPromise,
-} from '../util/DeferredPromise.js';
+import {createDeferredPromise} from '../util/DeferredPromise.js';
 import {CDPSession} from './Connection.js';
 import {ElementHandle} from './ElementHandle.js';
 import {TimeoutError} from './Errors.js';
@@ -117,8 +115,9 @@ export interface IsolatedWorldChart {
  */
 export class IsolatedWorld {
   #frame: Frame;
+  #injected: boolean;
   #document?: ElementHandle<Document>;
-  #contextPromise: DeferredPromise<ExecutionContext> = createDeferredPromise();
+  #contextPromise = createDeferredPromise<ExecutionContext>();
   #detached = false;
 
   // Set of bindings that have been registered in the current context.
@@ -140,10 +139,11 @@ export class IsolatedWorld {
     return `${name}_${contextId}`;
   };
 
-  constructor(frame: Frame) {
+  constructor(frame: Frame, injected = false) {
     // Keep own reference to client because it might differ from the FrameManager's
     // client for OOP iframes.
     this.#frame = frame;
+    this.#injected = injected;
     this.#client.on('Runtime.bindingCalled', this.#onBindingCalled);
   }
 
@@ -169,10 +169,9 @@ export class IsolatedWorld {
   }
 
   setContext(context: ExecutionContext): void {
-    assert(
-      this.#contextPromise,
-      `ExecutionContext ${context._contextId} has already been set.`
-    );
+    if (this.#injected) {
+      context.evaluate(injectedSource).catch(debugError);
+    }
     this.#ctxBindings.clear();
     this.#contextPromise.resolve(context);
     for (const waitTask of this._waitTasks) {

src/injected/injected.ts

@@ -1 +1,12 @@
-export * from './Poller.js';
+import * as Poller from './Poller.js';
+import * as util from './util.js';
+
+Object.assign(
+  self,
+  Object.freeze({
+    InjectedUtil: {
+      ...Poller,
+      ...util,
+    },
+  })
+);

src/injected/util.ts

@@ -0,0 +1,18 @@
+const createdFunctions = new Map<string, (...args: unknown[]) => unknown>();
+
+/**
+ * Creates a function from a string.
+ */
+export const createFunction = (
+  functionValue: string
+): ((...args: unknown[]) => unknown) => {
+  let fn = createdFunctions.get(functionValue);
+  if (fn) {
+    return fn;
+  }
+  fn = new Function(`return ${functionValue}`)() as (
+    ...args: unknown[]
+  ) => unknown;
+  createdFunctions.set(functionValue, fn);
+  return fn;
+};

src/templates/injected.ts.tmpl

@@ -0,0 +1,9 @@
+import * as Poller from '../injected/Poller.js';
+import * as util from '../injected/util.js';
+
+declare global {
+  const InjectedUtil: Readonly<typeof Poller & typeof util>;
+}
+
+/** @internal */
+export const source = SOURCE_CODE;

src/tsconfig.esm.json

@@ -8,5 +8,6 @@
   "references": [
     {"path": "../vendor/tsconfig.esm.json"},
     {"path": "../compat/esm/tsconfig.json"}
-  ]
+  ],
+  "exclude": ["injected/injected.ts"]
 }

test/src/injected.spec.ts

@@ -0,0 +1,40 @@
+/**
+ * Copyright 2022 Google Inc. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import expect from 'expect';
+import '../../lib/cjs/puppeteer/generated/injected.js';
+import {PUPPETEER_WORLD} from '../../lib/cjs/puppeteer/common/IsolatedWorld.js';
+import {
+  getTestState,
+  setupTestBrowserHooks,
+  setupTestPageAndContextHooks,
+} from './mocha-utils.js';
+
+describe('InjectedUtil tests', function () {
+  setupTestBrowserHooks();
+  setupTestPageAndContextHooks();
+
+  it('should work', async () => {
+    const {page} = getTestState();
+
+    const handle = await page
+      .mainFrame()
+      .worlds[PUPPETEER_WORLD].evaluate(() => {
+        return typeof InjectedUtil === 'object';
+      });
+    expect(handle).toBeTruthy();
+  });
+});

test/src/network.spec.ts

@@ -851,10 +851,13 @@ describe('network', function () {
           res.end();
         });
         await page.goto(httpsServer.PREFIX + '/setcookie.html');
-
+        const url = httpsServer.CROSS_PROCESS_PREFIX + '/setcookie.html';
         const response = await new Promise<HTTPResponse>(resolve => {
-          page.on('response', resolve);
-          const url = httpsServer.CROSS_PROCESS_PREFIX + '/setcookie.html';
+          page.on('response', response => {
+            if (response.url() === url) {
+              resolve(response);
+            }
+          });
           page.evaluate(src => {
             const xhr = new XMLHttpRequest();
             xhr.open('GET', src);

utils/generate_sources.ts

@@ -17,27 +17,37 @@ import {job} from './internal/job.js';
     .build();
 
   await job('', async ({name, inputs, outputs}) => {
+    const input = inputs.find(input => {
+      return input.endsWith('injected.ts');
+    })!;
+    const template = await readFile(
+      inputs.find(input => {
+        return input.includes('injected.ts.tmpl');
+      })!,
+      'utf8'
+    );
     const tmp = await mkdtemp(name);
     await esbuild.build({
-      entryPoints: [inputs[0]!],
+      entryPoints: [input],
       bundle: true,
       outdir: tmp,
       format: 'cjs',
       platform: 'browser',
       target: 'ES2019',
     });
-    const baseName = path.basename(inputs[0]!);
+    const baseName = path.basename(input);
     const content = await readFile(
       path.join(tmp, baseName.replace('.ts', '.js')),
       'utf-8'
     );
-    const scriptContent = `/** @internal */
-export const source = ${JSON.stringify(content)};
-`;
+    const scriptContent = template.replace(
+      'SOURCE_CODE',
+      JSON.stringify(content)
+    );
     await writeFile(outputs[0]!, scriptContent);
-    await rimraf.sync(tmp);
+    rimraf.sync(tmp);
   })
-    .inputs(['src/injected/**.ts'])
+    .inputs(['src/templates/injected.ts.tmpl', 'src/injected/**.ts'])
     .outputs(['src/generated/injected.ts'])
     .build();
 

utils/internal/job.ts

@@ -59,8 +59,8 @@ class JobBuilder {
       );
 
       if (
-        outputStats.reduce(reduceMaxTime, 0) >=
-        inputStats.reduce(reduceMinTime, Infinity)
+        outputStats.reduce(reduceMinTime, Infinity) >
+        inputStats.reduce(reduceMaxTime, 0)
       ) {
         shouldRun = false;
       }