diff --git a/src/renderer/modules/webpack/filters.ts b/src/renderer/modules/webpack/filters.ts index f09f644fb..0cba5dee3 100644 --- a/src/renderer/modules/webpack/filters.ts +++ b/src/renderer/modules/webpack/filters.ts @@ -1,4 +1,4 @@ -import { getExportsForProps } from "./get-modules"; +import { getExportsForProps, getFunctionForPrototype } from "./get-modules"; import { sourceStrings } from "./patch-load"; import type { RawModule } from "../../../types"; @@ -10,6 +10,14 @@ export const byProps =

(...props: P[]) => { return (m: RawModule) => typeof getExportsForProps(m.exports, props) !== "undefined"; }; +/** + * Get a module that has all the given prototypes on one of its functions + * @param prototypes List of prototype names + */ +export const byPrototype =

(...prototypes: P[]) => { + return (m: RawModule) => typeof getFunctionForPrototype(m.exports, prototypes) !== "undefined"; +}; + /** * Get a module whose source code matches the given string or RegExp * @param match String or RegExp to match in the module's source code diff --git a/src/renderer/modules/webpack/get-modules.ts b/src/renderer/modules/webpack/get-modules.ts index e2a282486..8216387c5 100644 --- a/src/renderer/modules/webpack/get-modules.ts +++ b/src/renderer/modules/webpack/get-modules.ts @@ -1,4 +1,4 @@ -import type { Filter, GetModuleOptions, RawModule } from "src/types"; +import type { AnyFunction, Filter, GetModuleOptions, RawModule, WithPrototype } from "src/types"; import { wpRequire } from "./patch-load"; import { logError } from "./util"; @@ -22,27 +22,40 @@ export function getExports(m: RawModule): T | undefined { } /** - * Iterates over an object and its top-level children that could have properties + * Iterates over an object and its top-level and second-level (if specified) children that could have properties * @param m Object (module exports) to iterate over + * @param secondLevel Boolean on whether to iterate over second level children in object */ -function* iterateModuleExports(m: unknown): IterableIterator> { +function* iterateModuleExports( + m: unknown, + secondLevel?: boolean, +): IterableIterator | AnyFunction> { // if m is null or not an object/function, then it will obviously not have the props - // if if has no props, then it definitely has no children either - if (m && (typeof m === "object" || typeof m === "function")) { - yield m as Record; - for (const key in m) { - try { + // if it has no props, then it definitely has no children either + try { + if (m && (typeof m === "object" || typeof m === "function")) { + yield m as Record; + for (const key in m) { // This could throw an error ("illegal invocation") if val === DOMTokenList.prototype // and key === "length" // There could be other cases too, hence this try-catch instead of a specific exclusion const val = (m as Record)[key]; if (val && (typeof val === "object" || typeof val === "function")) { - yield val as Record; + yield val as Record | AnyFunction; + if (secondLevel && typeof val === "object") { + for (const subKey in val) { + const subVal = (val as Record)[subKey]; + if (subVal && (typeof val === "object" || typeof val === "function")) { + yield subVal as Record | AnyFunction; + continue; + } + } + } } - } catch { - // ignore this export } } + } catch { + // ignore this export } } @@ -65,6 +78,25 @@ export function getExportsForProps( } } +/** + * Find an object in a module that has all the given properties. You will usually not need this function. + * @param m Module to search + * @param props Array of prototype names + * @returns Function that contains all the given prototypes (and any others), or undefined if not found + */ +export function getFunctionForPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(m: unknown, prototypes: P[]): T | undefined { + // Loop over the module and its exports at the top level + // Return the first thing that has all the indicated props + for (const exported of iterateModuleExports(m, true)) { + if (prototypes.every((p) => (exported as AnyFunction)?.prototype?.[p])) { + return exported as T; + } + } +} + // This doesn't have anywhere else to go export function getById(id: number, raw?: false): T | undefined; diff --git a/src/renderer/modules/webpack/helpers.ts b/src/renderer/modules/webpack/helpers.ts index a9f0abe2d..8226a8e3c 100644 --- a/src/renderer/modules/webpack/helpers.ts +++ b/src/renderer/modules/webpack/helpers.ts @@ -1,5 +1,11 @@ -import type { GetModuleOptions, RawModule, WaitForOptions } from "src/types"; -import { getExportsForProps, getModule } from "./get-modules"; +import type { + AnyFunction, + GetModuleOptions, + RawModule, + WaitForOptions, + WithPrototype, +} from "src/types"; +import { getExportsForProps, getFunctionForPrototype, getModule } from "./get-modules"; import * as filters from "./filters"; import Flux, { Store } from "../common/flux"; import { waitForModule } from "./lazy"; @@ -165,6 +171,131 @@ export async function waitForProps( return getExportsForProps(result as T, props)!; } +// Get by prototype + +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options?: { all?: false; raw?: false }): T | undefined; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options: { all: true; raw?: false }): T[]; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options: { all?: false; raw: true }): RawModule | undefined; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options: { all: true; raw: true }): Array>; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options?: { all: true; raw?: boolean }): T[] | Array>; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options: { all?: false; raw?: boolean }): T | RawModule | undefined; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>( + prototypes: P[], + options: { all?: boolean; raw: true }, +): RawModule | Array> | undefined; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P, options: { all?: boolean; raw?: false }): T | T[] | undefined; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>( + prototypes: P[], + options?: { all?: boolean; raw?: boolean }, +): T | T[] | RawModule | Array> | undefined; +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey[] = Array>, +>(...prototypes: P): T | undefined; + +/** + * Equivalent to `getModule(filters.byPrototypes(...props), options)` + * + * @see {@link filters.byPrototype} + * @see {@link getModule} + */ +export function getByPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>( + ...args: [P[], GetModuleOptions] | P[] +): T | T[] | RawModule | Array> | undefined { + const prototypes = (typeof args[0] === "string" ? args : args[0]) as P[]; + const raw = typeof args[0] === "string" ? false : (args[1] as GetModuleOptions)?.raw; + + const result = + typeof args.at(-1) === "object" + ? getModule(filters.byPrototype(...prototypes), args.at(-1) as GetModuleOptions) + : getModule(filters.byPrototype(...prototypes)); + + if (raw || typeof result === "undefined") { + return result as RawModule | undefined; + } + + if (result instanceof Array) { + // @ts-expect-error TypeScript isn't going to infer types based on the raw variable, so this is fine + return result.map((m) => getFunctionForPrototypes(m, prototypes)); + } + + return getFunctionForPrototype(result, prototypes); +} + +// Wait for prototypes + +export function waitForPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options: WaitForOptions & { raw?: false }): Promise; +export function waitForPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options: WaitForOptions & { raw: true }): Promise>; +export function waitForPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(prototypes: P[], options?: WaitForOptions): Promise>; +export function waitForPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(...prototypes: P[]): Promise; + +/** + * Like {@link getByPrototype} but waits for the module to be loaded. + * + * @see {@link getByPrototype} + * @see {@link waitForModule} + */ +export async function waitForPrototype< + T extends AnyFunction, + P extends PropertyKey = keyof WithPrototype, +>(...args: [P[], WaitForOptions] | P[]): Promise> { + const prototypes = (typeof args[0] === "string" ? args : args[0]) as P[]; + const raw = typeof args[0] === "string" ? false : (args[1] as WaitForOptions)?.raw; + + const result = await (typeof args.at(-1) === "object" + ? waitForModule(filters.byPrototype(...prototypes), args.at(-1) as WaitForOptions) + : waitForModule(filters.byPrototype(...prototypes))); + + if (raw) { + return result as RawModule; + } + + // We know this will always exist since filters.byPrototypes will always return a module that has the function with prototypes + return getFunctionForPrototype(result as T, prototypes)!; +} + // Get by value export function getByValue( diff --git a/src/renderer/modules/webpack/index.ts b/src/renderer/modules/webpack/index.ts index 719e010bb..f74263727 100644 --- a/src/renderer/modules/webpack/index.ts +++ b/src/renderer/modules/webpack/index.ts @@ -2,7 +2,7 @@ export { waitForModule } from "./lazy"; export { getFunctionBySource, getFunctionKeyBySource } from "./inner-search"; -export { getById, getExportsForProps, getModule } from "./get-modules"; +export { getById, getExportsForProps, getFunctionForPrototype, getModule } from "./get-modules"; /** * Filter functions to use with {@link getModule} diff --git a/src/types/webpack.ts b/src/types/webpack.ts index 77d9b670c..84623c553 100644 --- a/src/types/webpack.ts +++ b/src/types/webpack.ts @@ -9,6 +9,8 @@ export type ModuleExports = export type ModuleExportsWithProps

= Record & Record; +export type WithPrototype = T extends { prototype: infer P } ? P : never; + export interface RawModule { id: number; loaded: boolean;