From dd6742d08defd5f5c1d6b0bbcc5ba475ed06d107 Mon Sep 17 00:00:00 2001
From: amsyarasyiq <82711525+amsyarasyiq@users.noreply.github.com>
Date: Thu, 23 May 2024 15:17:35 +0800
Subject: [PATCH] feat(patcher): resolvable proxy patch

---
 src/lib/api/patcher.ts  | 28 ++++++++++++++++++++++++++--
 src/lib/ui/safeMode.tsx |  2 +-
 src/metro/finders.ts    |  2 +-
 src/metro/proxy.ts      | 18 ++++++++++--------
 4 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/src/lib/api/patcher.ts b/src/lib/api/patcher.ts
index c804783..8caf4b9 100644
--- a/src/lib/api/patcher.ts
+++ b/src/lib/api/patcher.ts
@@ -5,8 +5,11 @@ import {
     instead as _instead
 } from "spitroast";
 
-function shim<T extends (...args: any) => any>(fn: T) {
-    return function shimmed(this: any, ...args: Parameters<T>) {
+type ParametersAfterOneArgs<F> = F extends (x: any, y: any, ...args: infer P) => unknown ? P : never;
+type ProxyPatchParameters<T, P> = [func: string, parent: [P, (t: P) => any], ...args: ParametersAfterOneArgs<T>];
+
+function shim<T extends (func: string, parent: any, ...args: ParametersAfterOneArgs<T>) => any>(fn: T) {
+    function shimmed(this: any, ...args: Parameters<T>) {
         const proxyInfo = getFindProxyInfo(args[1]);
         if (proxyInfo && !proxyInfo.cache && proxyInfo.indexed) {
             let cancel = false;
@@ -20,7 +23,28 @@ function shim<T extends (...args: any) => any>(fn: T) {
             return () => unpatch();
         }
         return fn.apply(this, args);
+    }
+
+    shimmed.proxy = function proxyPatch<P>(this: any, ...args: ProxyPatchParameters<T, P>) {
+        const [target, resolve] = args[1];
+        const proxyInfo = getFindProxyInfo(target);
+        if (!proxyInfo) {
+            args[1] = resolve(target);
+            return shimmed(...args as unknown as Parameters<T>);
+        }
+
+        let cancel = false;
+        let unpatch = () => cancel = true;
+        proxyInfo.getExports(exp => {
+            if (cancel) return;
+            args[1] = resolve(exp);
+            unpatch = shimmed(...args as unknown as Parameters<T>);
+        });
+
+        return unpatch;
     };
+
+    return shimmed;
 }
 
 export const after = shim(_after);
diff --git a/src/lib/ui/safeMode.tsx b/src/lib/ui/safeMode.tsx
index 0a503af..c9e3627 100644
--- a/src/lib/ui/safeMode.tsx
+++ b/src/lib/ui/safeMode.tsx
@@ -68,7 +68,7 @@ const tabs: Tab[] = [
     { id: "componentStack", title: () => Strings.COMPONENT, trimWhitespace: true },
 ];
 
-export default () => after("render", ErrorBoundary.prototype, function (this: any, _, ret) {
+export default () => after.proxy("render", [ErrorBoundary, r => r.prototype], function (this: any, _, ret) {
     if (!this.state.error) return;
 
     // Not using setState here as we don't want to cause a re-render, we want this to be set in the initial render
diff --git a/src/metro/finders.ts b/src/metro/finders.ts
index 9b0a299..bb28ba2 100644
--- a/src/metro/finders.ts
+++ b/src/metro/finders.ts
@@ -2,7 +2,7 @@ import { getCacherForUniq } from "./caches";
 import { getModules, requireModule } from "./modules";
 import { FilterFn } from "./types";
 
-export function filterExports<A extends unknown[]>(
+function filterExports<A extends unknown[]>(
     moduleExports: any,
     moduleId: number, filter:
         FilterFn<A>
diff --git a/src/metro/proxy.ts b/src/metro/proxy.ts
index 51e22dc..f35f4b4 100644
--- a/src/metro/proxy.ts
+++ b/src/metro/proxy.ts
@@ -1,8 +1,8 @@
 import { proxyLazy } from "@lib/utils/lazy";
 
 import { getMetroCache } from "./caches";
-import { filterExports, findExports } from "./finders";
-import { requireModule, subscribeModule } from "./modules";
+import { findExports } from "./finders";
+import { metroModules, subscribeModule } from "./modules";
 import type { FilterFn } from "./types";
 
 const proxyInfoMap = new WeakMap<{}, FindProxyInfo<any[]>>();
@@ -11,6 +11,7 @@ interface FindProxyInfo<A extends unknown[]> {
     filter: FilterFn<A>;
     indexed: boolean;
     moduleId: number | undefined;
+    getExports(cb: (exports: any) => void): () => void;
     subscribe(cb: (exports: any) => void): () => void;
     unproxy(): any;
     get cache(): any;
@@ -29,12 +30,7 @@ function subscribeModuleForFind(proxy: any, callback: (exports: any) => void) {
     if (!info.indexed) throw new Error("Attempting to subscribe to a non-indexed find");
 
     return subscribeModule(info.moduleId!, () => {
-        const [exp] = filterExports(
-            requireModule(info.moduleId!),
-            info.moduleId!,
-            info.filter
-        );
-        callback(exp);
+        callback(findExports(info.filter));
     });
 }
 
@@ -50,6 +46,12 @@ export function createFindProxy<A extends unknown[]>(filter: FilterFn<A>) {
         filter,
         indexed: !!moduleId,
         moduleId,
+        getExports(cb: (exports: any) => void) {
+            if (moduleId && metroModules[moduleId!]?.isInitialized) {
+                return cb(findExports(filter)), () => { };
+            }
+            return this.subscribe(cb);
+        },
         subscribe(cb: (exports: any) => void) {
             return subscribeModuleForFind(proxy, cb);
         },