From 3c2b5fc5e16289b41c37cc4241fa30f36adb0c3e Mon Sep 17 00:00:00 2001
From: toonlink <toonlink@riseup.net>
Date: Tue, 20 Aug 2024 18:05:35 -0400
Subject: [PATCH] fix(unpatch): ensure unpatch does not hold a reference to
 Patch object

---
 package.json        |  2 +-
 src/getPatchFunc.ts | 14 ++++++++++----
 src/index.ts        |  4 ++--
 src/patcher.ts      |  0
 src/shared.ts       |  2 --
 src/unpatch.ts      | 10 +++++-----
 6 files changed, 18 insertions(+), 14 deletions(-)
 create mode 100644 src/patcher.ts

diff --git a/package.json b/package.json
index 1bf9a57..4b834da 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "spitroast",
-  "version": "2.0.0",
+  "version": "2.0.1",
   "description": "A simple JavaScript function patcher.",
   "main": "dist/cjs.js",
   "module": "dist/esm/index.js",
diff --git a/src/getPatchFunc.ts b/src/getPatchFunc.ts
index 5d5896d..6b2dadb 100644
--- a/src/getPatchFunc.ts
+++ b/src/getPatchFunc.ts
@@ -14,12 +14,13 @@ export default <CallbackType extends Function>(patchType: PatchType) =>
     callback: CallbackType,
     oneTime = false
   ) => {
-    if (typeof funcParent[funcName] !== "function")
+    let origFunc = funcParent[funcName];
+
+    if (typeof origFunc !== "function")
       throw new Error(
         `${funcName} is not a function in ${funcParent.constructor.name}`
       );
 
-    let origFunc = funcParent[funcName];
     let funcPatch = patchedFunctions.get(origFunc);
 
     if (!funcPatch) {
@@ -43,7 +44,11 @@ export default <CallbackType extends Function>(patchType: PatchType) =>
             : Reflect.get(target, prop, receiver),
       });
 
-      const runHook: any = (ctxt: unknown, args: unknown[], construct: boolean) => hook(replaceProxy, origFunc, args, ctxt, construct);
+      const runHook: any = (
+        ctxt: unknown,
+        args: unknown[],
+        construct: boolean
+      ) => hook(replaceProxy, origFunc, args, ctxt, construct);
 
       patchedFunctions.set(replaceProxy, funcPatch);
 
@@ -58,7 +63,8 @@ export default <CallbackType extends Function>(patchType: PatchType) =>
     }
 
     const hookId = Symbol();
-    const unpatchThisPatch = () => unpatch(funcPatch, hookId, patchType);
+    const funcPatchRef = new WeakRef(funcPatch);
+    const unpatchThisPatch = () => unpatch(funcPatchRef, hookId, patchType);
 
     if (oneTime) funcPatch.c.push(unpatchThisPatch);
     funcPatch[patchType].set(hookId, callback);
diff --git a/src/index.ts b/src/index.ts
index 87f8b34..670cd7b 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,5 +1,4 @@
 import getPatchFunc from "./getPatchFunc";
-import { unpatchAll } from "./unpatch";
 
 type BeforeCallback = (args: any[]) => void | undefined | any[];
 type InsteadCallback = (args: any[], origFunc: Function) => any;
@@ -9,4 +8,5 @@ const before = getPatchFunc<BeforeCallback>("b");
 const instead = getPatchFunc<InsteadCallback>("i");
 const after = getPatchFunc<AfterCallback>("a");
 
-export { instead, before, after, unpatchAll };
+export { instead, before, after };
+export { resetPatches as unpatchAll } from "./shared";
\ No newline at end of file
diff --git a/src/patcher.ts b/src/patcher.ts
new file mode 100644
index 0000000..e69de29
diff --git a/src/shared.ts b/src/shared.ts
index 255428f..18248d9 100644
--- a/src/shared.ts
+++ b/src/shared.ts
@@ -20,8 +20,6 @@ export type Patch = {
   i: Map<symbol, Function>;
 };
 
-export const patcherContext = { currentContext: { SHOULD_UNPATCH: false } };
-
 export let patchedFunctions: WeakMap<Function, Patch>;
 export let resetPatches = () =>
   (patchedFunctions = new WeakMap<Function, Patch>());
diff --git a/src/unpatch.ts b/src/unpatch.ts
index f04c1bd..9f3b5a6 100644
--- a/src/unpatch.ts
+++ b/src/unpatch.ts
@@ -1,7 +1,9 @@
 import { PatchType, Patch, patchTypes, patchedFunctions } from "./shared";
 
-export function unpatch(patch: Patch, hookId: symbol, type: PatchType) {
-  if (!patch[type].delete(hookId)) return false;
+export function unpatch(patchRef: WeakRef<Patch>, hookId: symbol, type: PatchType) {
+  const patch = patchRef.deref();
+
+  if (!patch || !patch[type].delete(hookId)) return false;
   const funcParent = patch.p.deref();
 
   // If there are no more hooks for every type, remove the patch
@@ -24,6 +26,4 @@ export function unpatch(patch: Patch, hookId: symbol, type: PatchType) {
   }
 
   return true;
-}
-
-export { resetPatches as unpatchAll } from "./shared";
+}
\ No newline at end of file