feat: init-script injection

init_script/hello_world.js

@@ -0,0 +1,4 @@
+(() => {
+    const root = typeof globalThis !== "undefined" ? globalThis : (typeof global !== "undefined" ? global : this);
+    root["__helloWorld"] = "Hello, World!";
+})();

src/cli.ts

@@ -5,13 +5,15 @@ type CliOptions = {
     cdpPort: number;
     debugMain: boolean;
     debugFrida: boolean;
+    initScript: string;
 };
 
 // default debugging port, do not change
 const DEBUG_PORT = 9421;
 // CDP port, change to whatever you like
 // use this port by navigating to devtools://devtools/bundled/inspector.html?ws=127.0.0.1:${CDP_PORT}
 const CDP_PORT = 62000;
+const INIT_SCRIPT = "init_script/hello_world.js";
 
 const print_help = () => {
     console.log(`Usage: npx ts-node src/index.ts [options]
@@ -21,6 +23,7 @@ Options:
   --cdp-port <port>    CDP proxy server port (default: ${CDP_PORT})
   --debug-main         Output main process debug messages
   --debug-frida        Output Frida client messages
+  --init-script <path> Path to the init script (default: ${INIT_SCRIPT})
   -h, --help           Show this help message`);
 };
 
@@ -48,6 +51,7 @@ const parse_cli_options = (): CliOptions => {
             "cdp-port": { type: "string" },
             "debug-main": { type: "boolean" },
             "debug-frida": { type: "boolean" },
+            "init-script": { type: "string" },
             help: { type: "boolean", short: "h" },
         },
         allowPositionals: false,
@@ -63,6 +67,7 @@ const parse_cli_options = (): CliOptions => {
         cdpPort: parse_port("--cdp-port", values["cdp-port"], CDP_PORT),
         debugMain: values["debug-main"] ?? false,
         debugFrida: values["debug-frida"] ?? false,
+        initScript: values["init-script"] ?? INIT_SCRIPT,
     };
 };
 

src/index.ts

@@ -20,14 +20,148 @@ const bufferToHexString = (buffer: ArrayBuffer) => {
         .join("");
 };
 
-const debug_server = (options: CliOptions, logger: Logger) => {
+const tryParseJson = (value: string) => {
+    try {
+        return JSON.parse(value);
+    } catch {
+        return null;
+    }
+};
+
+const APPSERVICE_CONTEXT_PROBE_EXPRESSION = `(() => ({
+    ok: true,
+    hasLocation: typeof location !== "undefined",
+    hasWx: typeof wx !== "undefined",
+    hasRequire: typeof require !== "undefined",
+    href: typeof location !== "undefined" ? location.href : null,
+}))()`;
+
+const debug_server = (
+    options: CliOptions,
+    logger: Logger,
+    initScriptSource: string,
+) => {
     const wss = new WebSocketServer({ port: options.debugPort });
     logger.info(
         `[server] debug server running on ws://localhost:${options.debugPort}`,
     );
     logger.info(`[server] debug server waiting for miniapp to connect...`);
 
     let messageCounter = 0;
+    let nextProbeScriptInjectionId = 100_000;
+    const pendingProbeScriptInjectionResults = new Map<
+        number,
+        {
+            executionContextId: number;
+            frameId: string | null;
+            sessionId: string | null;
+        }
+    >();
+    const probeScriptInjectedContexts = new Set<string>();
+    let nextInitScriptInjectionId = 200_000;
+    const pendingInitScriptInjectionResults = new Map<
+        number,
+        {
+            executionContextId: number;
+            frameId: string | null;
+            sessionId: string | null;
+        }
+    >();
+    const initScriptInjectedContexts = new Set<string>();
+    let candidate:
+        | {
+              executionContextId: number;
+              frameId: string | null;
+              sessionId: string | null;
+          }
+        | null = null;
+
+    const sendCdpCommand = (command: Record<string, unknown>) => {
+        debugMessageEmitter.emit("proxymessage", JSON.stringify(command));
+    };
+
+    const injectWxAppServiceProbeScript = (
+        executionContextId: number,
+        frameId: string | null,
+        sessionId: string | null,
+    ) => {
+        const dedupeKey = `${sessionId ?? "root"}:${executionContextId}`;
+        if (probeScriptInjectedContexts.has(dedupeKey)) {
+            return;
+        }
+
+        probeScriptInjectedContexts.add(dedupeKey);
+        const id = nextProbeScriptInjectionId++;
+        pendingProbeScriptInjectionResults.set(id, {
+            executionContextId,
+            frameId,
+            sessionId,
+        });
+
+        const command: Record<string, unknown> = {
+            id,
+            method: "Runtime.evaluate",
+            params: {
+                expression: APPSERVICE_CONTEXT_PROBE_EXPRESSION,
+                contextId: executionContextId,
+                returnByValue: true,
+                silent: true,
+            },
+        };
+
+        if (sessionId) {
+            command.sessionId = sessionId;
+        }
+
+        sendCdpCommand(command);
+    };
+
+    const injectInitScript = (
+        executionContextId: number,
+        frameId: string | null,
+        sessionId: string | null,
+    ) => {
+        const dedupeKey = `${sessionId ?? "root"}:${executionContextId}`;
+        if (initScriptInjectedContexts.has(dedupeKey)) {
+            return;
+        }
+
+        initScriptInjectedContexts.add(dedupeKey);
+        const id = nextInitScriptInjectionId++;
+        pendingInitScriptInjectionResults.set(id, {
+            executionContextId,
+            frameId,
+            sessionId,
+        });
+
+        const command: Record<string, unknown> = {
+            id,
+            method: "Runtime.evaluate",
+            params: {
+                expression: initScriptSource,
+                contextId: executionContextId,
+                includeCommandLineAPI: true,
+                awaitPromise: false,
+                returnByValue: true,
+            },
+        };
+
+        if (sessionId) {
+            command.sessionId = sessionId;
+        }
+
+        sendCdpCommand(command);
+    };
+
+    const resetContextSelection = () => {
+        nextProbeScriptInjectionId = 100_000;
+        nextInitScriptInjectionId = 200_000;
+        pendingProbeScriptInjectionResults.clear();
+        pendingInitScriptInjectionResults.clear();
+        probeScriptInjectedContexts.clear();
+        initScriptInjectedContexts.clear();
+        candidate = null;
+    };
 
     const onMessage = (message: ArrayBuffer) => {
         logger.main_debug(
@@ -51,12 +185,92 @@ const debug_server = (options: CliOptions, logger: Logger) => {
         }
 
         if (unwrappedData.category === "chromeDevtoolsResult") {
+            const parsedPayload = tryParseJson(unwrappedData.data.payload);
+            const toCheckProbeResult =
+                parsedPayload && typeof parsedPayload.id === "number"
+                    ? pendingProbeScriptInjectionResults.get(parsedPayload.id)
+                    : undefined;
+            const toCheckInitScriptInjectionResult =
+                parsedPayload && typeof parsedPayload.id === "number"
+                    ? pendingInitScriptInjectionResults.get(parsedPayload.id)
+                    : undefined;
+
+            if (toCheckProbeResult) {
+                pendingProbeScriptInjectionResults.delete(parsedPayload.id);
+                const probeValue = parsedPayload.result?.result?.value ?? null;
+
+                if (
+                    candidate === null &&
+                    probeValue &&
+                    typeof probeValue === "object" &&
+                    probeValue.ok === true &&
+                    probeValue.hasLocation === false &&
+                    probeValue.hasWx === true &&
+                    probeValue.hasRequire === true
+                ) {
+                    candidate = {
+                        executionContextId:
+                            toCheckProbeResult.executionContextId,
+                        frameId: toCheckProbeResult.frameId,
+                        sessionId: toCheckProbeResult.sessionId,
+                    };
+                    logger.info("[miniapp] identified wechat app-service candidate", {
+                        sessionId: candidate.sessionId,
+                        executionContextId:
+                            candidate.executionContextId,
+                        frameId: candidate.frameId,
+                    });
+                    injectInitScript(
+                        candidate.executionContextId,
+                        candidate.frameId,
+                        candidate.sessionId,
+                    );
+                }
+            }
+
+            if (toCheckInitScriptInjectionResult) {
+                pendingInitScriptInjectionResults.delete(parsedPayload.id);
+                const exceptionDetails =
+                    parsedPayload.result?.exceptionDetails ?? null;
+
+                logger.info("[miniapp] init-script injection result", {
+                    sessionId: toCheckInitScriptInjectionResult.sessionId,
+                    executionContextId:
+                        toCheckInitScriptInjectionResult.executionContextId,
+                    frameId: toCheckInitScriptInjectionResult.frameId,
+                    success: exceptionDetails === null,
+                    exceptionDetails,
+                });
+            }
+
+
+            const context = parsedPayload?.params?.context;
+            if (
+                parsedPayload?.method === "Runtime.executionContextCreated" &&
+                (parsedPayload.sessionId === undefined || parsedPayload.sessionId === null) &&
+                context?.origin === "https://servicewechat.com" &&
+                context?.auxData?.type === "default"
+            ) {
+                logger.info("[miniapp] need to inject probe script", {
+                    method: parsedPayload?.method,
+                    sessionId: parsedPayload?.sessionId,
+                    contextOrigin: context?.origin,
+                    contextAuxDataType: context?.auxData?.type,
+                });
+                injectWxAppServiceProbeScript(
+                    parsedPayload.params.context.id,
+                    parsedPayload.params.context.auxData?.frameId ?? null,
+                    parsedPayload.sessionId ?? null,
+                );
+            }
+
             // need to proxy to CDP client
             debugMessageEmitter.emit("cdpmessage", unwrappedData.data.payload);
         }
     };
 
     wss.on("connection", (ws: WebSocket) => {
+        resetContextSelection();
         logger.info("[miniapp] miniapp client connected");
         ws.on("message", onMessage);
         ws.on("error", (err) => {
@@ -179,14 +393,7 @@ const frida_server = async (options: CliOptions, logger: Logger) => {
     const session = await localDevice.attach(Number(wmpfPid));
 
     // find hook script
-    const projectRoot = path.join(
-        path.dirname(
-            (require.main && require.main.filename) ||
-                (process.mainModule && process.mainModule.filename) ||
-                process.cwd(),
-        ),
-        "..",
-    );
+    const projectRoot = getProjectRoot();
     let scriptContent: string | null = null;
     try {
         scriptContent = (
@@ -237,10 +444,35 @@ const frida_server = async (options: CliOptions, logger: Logger) => {
     logger.info(`[frida] you can now open any miniapps`);
 };
 
+const getProjectRoot = () =>
+    path.join(
+        path.dirname(
+            (require.main && require.main.filename) ||
+                (process.mainModule && process.mainModule.filename) ||
+                process.cwd(),
+        ),
+        "..",
+    );
+
+const load_init_script_source = async (options: CliOptions) => {
+    const projectRoot = getProjectRoot();
+
+    try {
+        return (
+            await promises.readFile(path.join(projectRoot, options.initScript))
+        ).toString();
+    } catch {
+        throw new Error(
+            `[main] init script not found: ${options.initScript}`,
+        );
+    }
+};
+
 const main = async () => {
     const options = parse_cli_options();
     const logger = create_logger(options);
-    debug_server(options, logger);
+    const initScriptSource = await load_init_script_source(options);
+    debug_server(options, logger, initScriptSource);
     proxy_server(options, logger);
     frida_server(options, logger);
 };