[exfiltrate] only patch define when necessary

uwu · Apr 1, 2024 · d1e47de · d1e47de
1 parent 9c8a7da
commit d1e47de
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 11 deletions.
diff --git a/packages/shelter/src/exfiltrate.ts b/packages/shelter/src/exfiltrate.ts
@@ -1,14 +1,21 @@
-import { before } from "spitroast";
+import { after } from "spitroast";
+const origDefineProperty = Object.defineProperty;
 const current = new Set<string>();
 
-export default function (prop: string, filter?: (t: any) => boolean) {
+export default function (prop: string, patchDefine: boolean, filter?: (t: any) => boolean) {
   if (current[prop]) throw new Error(`Already exfiltrating ${prop}!`);
 
   const protoKey = Symbol(prop);
   let hitProto = false;
+  let unpatchDefine: () => void;
+
+  const cleanup = () => {
+    unpatchDefine?.();
+    if (!hitProto) delete Object.prototype[prop];
+  };
 
   return new Promise<any>((res) => {
-    Object.defineProperty(Object.prototype, prop, {
+    origDefineProperty(Object.prototype, prop, {
       configurable: true,
       enumerable: false,
       set(v: any) {
@@ -18,26 +25,31 @@ export default function (prop: string, filter?: (t: any) => boolean) {
           return;
         }
 
-        Object.defineProperty(this, prop, {
+        origDefineProperty(this, prop, {
           configurable: true,
           writable: true,
           enumerable: true,
           value: v,
         });
+
+        if (!filter || filter(this)) {
+          cleanup();
+          res(this);
+        }
       },
 
       get() {
         return this[protoKey];
       },
     });
 
-    const unpatch = before("defineProperty", Object, (args) => {
+    if (!patchDefine) return;
+    unpatchDefine = after("defineProperty", Object, (args) => {
       if (args[1] === prop) {
         queueMicrotask(() => {
           if (!filter || filter(args[0])) {
-            unpatch();
+            cleanup();
             res(args[0]);
-            if (!hitProto) delete Object.prototype[prop];
           }
         });
       }

diff --git a/packages/shelter/src/flux/index.ts b/packages/shelter/src/flux/index.ts
@@ -13,13 +13,13 @@ export async function getDispatcher() {
   if (dispatcher) return dispatcher;
 
   // Promises bubble up, this is fine.
-  return (dispatcher = exfiltrate("_dispatcher").then((fluxstore) => fluxstore._dispatcher as Dispatcher));
+  return (dispatcher = exfiltrate("_dispatcher", true).then((fluxstore) => fluxstore._dispatcher as Dispatcher));
 }
 
 export const stores: Record<string, FluxStore | FluxStore[]> = {};
 
 // noinspection JSIgnoredPromiseFromCall
-exfiltrate("_dispatchToken", (store: FluxStore) => {
+exfiltrate("_dispatchToken", true, (store: FluxStore) => {
   const name = store.getName();
   if (!stores[name]) {
     stores[name] = store;

diff --git a/packages/shelter/src/react.ts b/packages/shelter/src/react.ts
@@ -2,5 +2,5 @@ import exfiltrate from "./exfiltrate";
 
 export let React, ReactDOM;
 
-exfiltrate("useId").then((v) => (React = v));
-exfiltrate("findDOMNode").then((v) => (ReactDOM = v));
+exfiltrate("useId", false).then((v) => (React = v));
+exfiltrate("findDOMNode", false).then((v) => (ReactDOM = v));