decky-frontend-lib/src/utils/patcher.ts

// TODO: implement storing patches as an option so we can offer unpatchAll selectively
// Return this in a replacePatch to call the original method (can still modify args).
export let callOriginal = Symbol("DECKY_CALL_ORIGINAL");

export interface PatchOptions {
    singleShot?: boolean
}

type GenericPatchHandler = (args: any[], ret?: any) => any;

export interface Patch {
    original: Function;
    property: string;
    object: any;
    patchedFunction: any;
    hasUnpatched: boolean;
    handler: GenericPatchHandler;

    unpatch: () => void
};

// let patches = new Set<Patch>();

export function beforePatch(object: any, property: string, handler: (args: any[]) => any, options: PatchOptions = {}): Patch {
    const orig = object[property];
    object[property] = function (...args: any[]) {
        handler.call(this, args);
        const ret = patch.original.call(this, ...args);
        if (options.singleShot) {
            patch.unpatch();
        }
        return ret;
    }
    const patch = processPatch(object, property, handler, object[property], orig);
    return patch;
}

export function afterPatch(object: any, property: string, handler: (args: any[], ret: any) => any, options: PatchOptions = {}): Patch {
    const orig = object[property];
    object[property] = function (...args: any[]) {
        let ret = patch.original.call(this, ...args);
        ret = handler.call(this, args, ret);
        if (options.singleShot) {
            patch.unpatch();
        }
        return ret;
    }
    const patch = processPatch(object, property, handler, object[property], orig);
    return patch;
}

export function replacePatch(object: any, property: string, handler: (args: any[]) => any, options: PatchOptions = {}): Patch {
    const orig = object[property];
    object[property] = function (...args: any[]) {
        const ret = handler.call(this, args);
        if (ret == callOriginal) return patch.original.call(this, ...args);
        if (options.singleShot) {
            patch.unpatch();
        }
        return ret;
    };
    const patch = processPatch(object, property, handler, object[property], orig);
    return patch;
}

function processPatch(object: any, property: any, handler: GenericPatchHandler, patchedFunction: any, original: any): Patch {
    // Assign all props of original function to new one
    Object.assign(object[property], original);
    // Allow toString webpack filters to continue to work
    object[property].toString = () => original.toString();

    // HACK: for compatibility, remove when all plugins are using new patcher
    Object.defineProperty(object[property], "__deckyOrig", {
        get: () => patch.original,
        set: (val: any) => patch.original = val
    })

    // Build a Patch object of this patch
    const patch: Patch = {
        object,
        property,
        handler,
        patchedFunction,
        original,
        hasUnpatched: false,
        unpatch: () => unpatch(patch)
    };

    object[property].__deckyPatch = patch;

    return patch;
}

function unpatch(patch: Patch): void {
    const { object, property, handler, patchedFunction, original } = patch;
    if (patch.hasUnpatched) throw new Error("Function is already unpatched.")
    let realProp = property;
    let realObject = object;
    console.debug("[Patcher] unpatching", {realObject, realProp, object, property, handler, patchedFunction, original, isEqual: realObject[realProp] === patchedFunction})

    // If another patch has been applied to this function after this one, move down until we find the correct patch
    while (realObject[realProp] && realObject[realProp] !== patchedFunction) {
        realObject = realObject[realProp].__deckyPatch;
        realProp = "original";
        console.debug("[Patcher] moved to next", {realObject, realProp, object, property, handler, patchedFunction, original, isEqual: realObject[realProp] === patchedFunction})
    }

    realObject[realProp] = realObject[realProp].__deckyPatch.original

    patch.hasUnpatched = true;
    console.debug("[Patcher] unpatched", {realObject, realProp, object, property, handler, patchedFunction, original, isEqual: realObject[realProp] === patchedFunction})
}