tor-browser

ExposedFunction.ts (6935B)

---

/**
* @license
* Copyright 2023 Google Inc.
* SPDX-License-Identifier: Apache-2.0
*/

import * as Bidi from 'chromium-bidi/lib/cjs/protocol/protocol.js';

import {EventEmitter} from '../common/EventEmitter.js';
import type {Awaitable, FlattenHandle} from '../common/types.js';
import {debugError} from '../common/util.js';
import {DisposableStack} from '../util/disposable.js';
import {interpolateFunction, stringifyFunction} from '../util/Function.js';

import type {Connection} from './core/Connection.js';
import {BidiElementHandle} from './ElementHandle.js';
import type {BidiFrame} from './Frame.js';
import {BidiJSHandle} from './JSHandle.js';

type CallbackChannel<Args, Ret> = (
 value: [
   resolve: (ret: FlattenHandle<Awaited<Ret>>) => void,
   reject: (error: unknown) => void,
   args: Args,
 ],
) => void;

/**
* @internal
*/
export class ExposableFunction<Args extends unknown[], Ret> {
 static async from<Args extends unknown[], Ret>(
   frame: BidiFrame,
   name: string,
   apply: (...args: Args) => Awaitable<Ret>,
   isolate = false,
 ): Promise<ExposableFunction<Args, Ret>> {
   const func = new ExposableFunction(frame, name, apply, isolate);
   await func.#initialize();
   return func;
 }

 readonly #frame;

 readonly name;
 readonly #apply;
 readonly #isolate;

 readonly #channel;

 #scripts: Array<[BidiFrame, Bidi.Script.PreloadScript]> = [];
 #disposables = new DisposableStack();

 constructor(
   frame: BidiFrame,
   name: string,
   apply: (...args: Args) => Awaitable<Ret>,
   isolate = false,
 ) {
   this.#frame = frame;
   this.name = name;
   this.#apply = apply;
   this.#isolate = isolate;

   this.#channel = `__puppeteer__${this.#frame._id}_page_exposeFunction_${this.name}`;
 }

 async #initialize() {
   const connection = this.#connection;
   const channel = {
     type: 'channel' as const,
     value: {
       channel: this.#channel,
       ownership: Bidi.Script.ResultOwnership.Root,
     },
   };

   const connectionEmitter = this.#disposables.use(
     new EventEmitter(connection),
   );
   connectionEmitter.on(
     Bidi.ChromiumBidi.Script.EventNames.Message,
     this.#handleMessage,
   );

   const functionDeclaration = stringifyFunction(
     interpolateFunction(
       (callback: CallbackChannel<Args, Ret>) => {
         Object.assign(globalThis, {
           [PLACEHOLDER('name') as string]: function (...args: Args) {
             return new Promise<FlattenHandle<Awaited<Ret>>>(
               (resolve, reject) => {
                 callback([resolve, reject, args]);
               },
             );
           },
         });
       },
       {name: JSON.stringify(this.name)},
     ),
   );

   const frames = [this.#frame];
   for (const frame of frames) {
     frames.push(...frame.childFrames());
   }

   await Promise.all(
     frames.map(async frame => {
       const realm = this.#isolate ? frame.isolatedRealm() : frame.mainRealm();
       try {
         const [script] = await Promise.all([
           frame.browsingContext.addPreloadScript(functionDeclaration, {
             arguments: [channel],
             sandbox: realm.sandbox,
           }),
           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.
         debugError(error);
       }
     }),
   );
 }

 get #connection(): Connection {
   return this.#frame.page().browser().connection;
 }

 #handleMessage = async (params: Bidi.Script.MessageParameters) => {
   if (params.channel !== this.#channel) {
     return;
   }
   const realm = this.#getRealm(params.source);
   if (!realm) {
     // Unrelated message.
     return;
   }

   using dataHandle = BidiJSHandle.from<
     [
       resolve: (ret: FlattenHandle<Awaited<Ret>>) => void,
       reject: (error: unknown) => void,
       args: Args,
     ]
   >(params.data, realm);

   using stack = new DisposableStack();
   const args = [];

   let result;
   try {
     using argsHandle = await dataHandle.evaluateHandle(([, , args]) => {
       return args;
     });

     for (const [index, handle] of await argsHandle.getProperties()) {
       stack.use(handle);

       // Element handles are passed as is.
       if (handle instanceof BidiElementHandle) {
         args[+index] = handle;
         stack.use(handle);
         continue;
       }

       // Everything else is passed as the JS value.
       args[+index] = handle.jsonValue();
     }
     result = await this.#apply(...((await Promise.all(args)) as Args));
   } catch (error) {
     try {
       if (error instanceof Error) {
         await dataHandle.evaluate(
           ([, reject], name, message, stack) => {
             const error = new Error(message);
             error.name = name;
             if (stack) {
               error.stack = stack;
             }
             reject(error);
           },
           error.name,
           error.message,
           error.stack,
         );
       } else {
         await dataHandle.evaluate(([, reject], error) => {
           reject(error);
         }, error);
       }
     } catch (error) {
       debugError(error);
     }
     return;
   }

   try {
     await dataHandle.evaluate(([resolve], result) => {
       resolve(result);
     }, result);
   } catch (error) {
     debugError(error);
   }
 };

 #getRealm(source: Bidi.Script.Source) {
   const frame = this.#findFrame(source.context as string);
   if (!frame) {
     // Unrelated message.
     return;
   }
   return frame.realm(source.realm);
 }

 #findFrame(id: string) {
   const frames = [this.#frame];
   for (const frame of frames) {
     if (frame._id === id) {
       return frame;
     }
     frames.push(...frame.childFrames());
   }
   return;
 }

 [Symbol.dispose](): void {
   void this[Symbol.asyncDispose]().catch(debugError);
 }

 async [Symbol.asyncDispose](): Promise<void> {
   this.#disposables.dispose();
   await Promise.all(
     this.#scripts.map(async ([frame, script]) => {
       const realm = this.#isolate ? frame.isolatedRealm() : frame.mainRealm();
       try {
         await Promise.all([
           realm.evaluate(name => {
             delete (globalThis as any)[name];
           }, this.name),
           ...frame.childFrames().map(childFrame => {
             return childFrame.evaluate(name => {
               delete (globalThis as any)[name];
             }, this.name);
           }),
           frame.browsingContext.removePreloadScript(script),
         ]);
       } catch (error) {
         debugError(error);
       }
     }),
   );
 }
}