From f23289ae5470354d4391700f53e169ca025c72be Mon Sep 17 00:00:00 2001 From: yofukashino Date: Sat, 7 Sep 2024 09:08:03 +0530 Subject: [PATCH] concept --- src/renderer/modules/injector.ts | 234 +++++++++++++++++++++++++++++++ 1 file changed, 234 insertions(+) diff --git a/src/renderer/modules/injector.ts b/src/renderer/modules/injector.ts index 8be38446b..8683205af 100644 --- a/src/renderer/modules/injector.ts +++ b/src/renderer/modules/injector.ts @@ -10,8 +10,11 @@ import { addButton } from "../coremods/messagePopover"; enum InjectionTypes { Before, + BeforeAsync, Instead, + InsteadAsync, After, + AfterAsync, } /** @@ -25,6 +28,17 @@ export type BeforeCallback = ( self: ObjectExports, ) => A | undefined | void; +/** + * Code to run before the original function asynchronously + * @param args Arguments passed to the original function + * @param self The module the injected function is on + * @returns New arguments to pass to the original function, or undefined to leave them unchanged + */ +export type BeforeCallbackAsync = ( + args: A, + self: ObjectExports, +) => Promise; + /** * Code to run instead of the original function * @param args Arguments passed to the original function @@ -38,6 +52,19 @@ export type InsteadCallback = ( self: ObjectExports, ) => R | void; +/** + * Code to run instead of the original function asynchronously + * @param args Arguments passed to the original function + * @param orig The original function + * @param self The module the injected function is on + * @returns New result to return + */ +export type InsteadCallbackAsync = ( + args: A, + orig: (...args: A) => R, + self: ObjectExports, +) => R | void; + /** * Code to run after the original function * @param args Arguments passed to the original function @@ -51,6 +78,19 @@ export type AfterCallback = ( self: ObjectExports, ) => R | undefined | void; +/** + * Code to run after the original function asynchronously + * @param args Arguments passed to the original function + * @param res Result of the original function + * @param self The module the injected function is on + * @returns New result to return, or undefined to leave it unchanged + */ +export type AfterCallbackAsync = ( + args: A, + res: R, + self: ObjectExports, +) => Promise; + interface ObjectInjections { injections: Map< string, @@ -63,8 +103,22 @@ interface ObjectInjections { original: Map; } +interface ObjectInjectionsAsync { + injections: Map< + string, + { + before: Set; + instead: Set; + after: Set; + } + >; + original: Map; +} + const injections = new WeakMap(); +const injectionsAsync = new WeakMap(); + function replaceMethod, U extends keyof T & string>( obj: T, funcName: U, @@ -136,6 +190,77 @@ function replaceMethod, U extends keyof T & str } } +function replaceMethodAsync, U extends keyof T & string>( + obj: T, + funcName: U, +): void { + if (typeof obj[funcName] !== "function") { + throw new TypeError(`Value of '${funcName}' in object is not a function`); + } + + let objInjections: ObjectInjectionsAsync; + if (injectionsAsync.has(obj)) { + objInjections = injectionsAsync.get(obj)!; + } else { + objInjections = { + injections: new Map(), + original: new Map(), + }; + injectionsAsync.set(obj, objInjections); + } + + if (!objInjections.injections.has(funcName)) { + objInjections.injections.set(funcName, { + before: new Set(), + instead: new Set(), + after: new Set(), + }); + } + + if (!objInjections.original.has(funcName)) { + objInjections.original.set(funcName, obj[funcName]); + + const originalFunc = obj[funcName]; + + // @ts-expect-error https://github.com/microsoft/TypeScript/issues/48992 + obj[funcName] = async function (...args: unknown[]): unknown { + const injectionsForProp = objInjections.injections.get(funcName)!; + + for (const b of injectionsForProp.before) { + const newArgs = await b.call(this, args, this); + if (Array.isArray(newArgs)) { + args = newArgs; + } + } + + let res: unknown; + + if (injectionsForProp.instead.size === 0) { + res = await originalFunc.apply(this, args); + } else { + for (const i of injectionsForProp.instead) { + const newResult = await i.call(this, args, originalFunc, this); + if (newResult !== void 0) { + res = newResult; + } + } + } + + for (const a of injectionsForProp.after) { + const newResult = await a.call(this, args, res, this); + if (newResult !== void 0) { + res = newResult; + } + } + + return res; + }; + + Object.defineProperties(obj[funcName], Object.getOwnPropertyDescriptors(originalFunc)); + obj[funcName].toString = originalFunc.toString.bind(originalFunc); + } +} + function inject, U extends keyof T & string>( obj: T, funcName: U, @@ -163,6 +288,33 @@ function inject, U extends keyof T & string>( return () => void setRef.deref()?.delete(cb); } +function injectAsync, U extends keyof T & string>( + obj: T, + funcName: U, + cb: BeforeCallbackAsync | InsteadCallbackAsync | AfterCallbackAsync, + type: InjectionTypes, +): () => void { + replaceMethodAsync(obj, funcName); + const methodInjections = injectionsAsync.get(obj)!.injections.get(funcName)!; + let set: Set; + switch (type) { + case InjectionTypes.BeforeAsync: + set = methodInjections.before; + break; + case InjectionTypes.InsteadAsync: + set = methodInjections.instead; + break; + case InjectionTypes.AfterAsync: + set = methodInjections.after; + break; + default: + throw new Error(`Invalid injection type: ${type}`); + } + set.add(cb); + const setRef = new WeakRef(set); + return () => void setRef.deref()?.delete(cb); +} + function before< T extends Record, U extends keyof T & string, @@ -172,6 +324,15 @@ function before< return inject(obj, funcName, cb as BeforeCallback, InjectionTypes.Before); } +function beforeAsync< + T extends Record, + U extends keyof T & string, + A extends unknown[] = Parameters, +>(obj: T, funcName: U, cb: BeforeCallbackAsync): () => void { + // @ts-expect-error 'unknown[]' is assignable to the constraint of type 'A', but 'A' could be instantiated with a different subtype of constraint 'unknown[]'. + return injectAsync(obj, funcName, cb as BeforeCallbackAsync, InjectionTypes.BeforeAsync); +} + function instead< T extends Record, U extends keyof T & string, @@ -182,6 +343,16 @@ function instead< return inject(obj, funcName, cb, InjectionTypes.Instead); } +function insteadAsync< + T extends Record, + U extends keyof T & string, + A extends unknown[] = Parameters, + R = ReturnType, +>(obj: T, funcName: U, cb: InsteadCallbackAsync): () => void { + // @ts-expect-error 'unknown[]' is assignable to the constraint of type 'A', but 'A' could be instantiated with a different subtype of constraint 'unknown[]'. + return inject(obj, funcName, cb, InjectionTypes.InsteadAsync); +} + function after< T extends Record, U extends keyof T & string, @@ -192,6 +363,16 @@ function after< return inject(obj, funcName, cb, InjectionTypes.After); } +function afterAsync< + T extends Record, + U extends keyof T & string, + A extends unknown[] = Parameters, + R = ReturnType, +>(obj: T, funcName: U, cb: AfterCallbackAsync): () => void { + // @ts-expect-error 'unknown[]' is assignable to the constraint of type 'A', but 'A' could be instantiated with a different subtype of constraint 'unknown[]'. + return inject(obj, funcName, cb, InjectionTypes.AfterAsync); +} + /** * Inject code into Discord's webpack modules. * @@ -238,6 +419,23 @@ export class Injector { return uninjector; } + /** + * Run code before a native module asynchronously + * @param obj Module to inject to + * @param funcName Function name on that module to inject + * @param cb Code to run + * @returns Uninject function + */ + public beforeAsync< + T extends Record, + U extends keyof T & string, + A extends unknown[] = Parameters, + >(obj: T, funcName: U, cb: BeforeCallbackAsync): () => void { + const uninjector = beforeAsync(obj, funcName, cb); + this.#uninjectors.add(uninjector); + return uninjector; + } + /** * Run code instead of a native module * @param obj Module to inject to @@ -256,6 +454,24 @@ export class Injector { return uninjector; } + /** + * Run code instead of a native module asynchronously + * @param obj Module to inject to + * @param funcName Function name on that module to inject + * @param cb Code to run + * @returns Uninject function + */ + public insteadAsync< + T extends Record, + U extends keyof T & string, + A extends unknown[] = Parameters, + R = ReturnType, + >(obj: T, funcName: U, cb: InsteadCallbackAsync): () => void { + const uninjector = insteadAsync(obj, funcName, cb); + this.#uninjectors.add(uninjector); + return uninjector; + } + /** * Run code after a native module * @param obj Module to inject to @@ -274,6 +490,24 @@ export class Injector { return uninjector; } + /** + * Run code after a native module asynchronously + * @param obj Module to inject to + * @param funcName Function name on that module to inject + * @param cb Code to run + * @returns Uninject function + */ + public afterAsync< + T extends Record, + U extends keyof T & string, + A extends unknown[] = Parameters, + R = ReturnType, + >(obj: T, funcName: U, cb: AfterCallbackAsync): () => void { + const uninjector = afterAsync(obj, funcName, cb); + this.#uninjectors.add(uninjector); + return uninjector; + } + /** * A few utils to add stuff in frequent modules. */