fix(js-x-ray): properly deep clone and reset probe context in ProbeRunner (#400)

Update workspaces/js-x-ray/test/ProbeRunner.spec.ts

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: fraxken

.changeset/ninety-adults-shout.md

@@ -0,0 +1,5 @@
+---
+"@nodesecure/js-x-ray": patch
+---
+
+Properly deep clone and reset probe context

workspaces/js-x-ray/src/ProbeRunner.ts

@@ -23,6 +23,9 @@ import isSerializeEnv from "./probes/isSerializeEnv.js";
 import type { SourceFile } from "./SourceFile.js";
 import type { OptionalWarningName } from "./warnings.js";
 
+// CONSTANTS
+const kProbeOriginalContext = Symbol.for("ProbeOriginalContext");
+
 export type ProbeReturn = void | null | symbol;
 export type ProbeContextDef = Record<string, any>;
 export type ProbeContext<T extends ProbeContextDef = ProbeContextDef> = {
@@ -50,7 +53,7 @@ export interface Probe<T extends ProbeContextDef = ProbeContextDef> {
   teardown?: (ctx: ProbeContext<T>) => void;
   breakOnMatch?: boolean;
   breakGroup?: string;
-  context?: ProbeContext<T>;
+  context?: T;
 }
 
 export class ProbeRunner {
@@ -107,9 +110,18 @@ export class ProbeRunner {
         `Invalid probe ${probe.name}: initialize must be a function or undefined`
       );
       if (probe.initialize) {
+        const isDefined = Reflect.defineProperty(probe, kProbeOriginalContext, {
+          enumerable: false,
+          value: structuredClone(probe.context),
+          writable: false
+        });
+        if (!isDefined) {
+          throw new Error(`Failed to define original context for probe '${probe.name}'`);
+        }
+
         const context = probe.initialize(this.#getProbeContext(probe));
         if (context) {
-          probe.context = context;
+          probe.context = structuredClone(context);
         }
       }
     }
@@ -193,6 +205,7 @@ export class ProbeRunner {
   finalize(): void {
     for (const probe of this.probes) {
       probe.finalize?.(this.#getProbeContext(probe));
+      probe.context = probe[kProbeOriginalContext];
     }
   }
 }

workspaces/js-x-ray/test/ProbeRunner.spec.ts

@@ -7,6 +7,7 @@ import type { ESTree } from "meriyah";
 
 // Import Internal Dependencies
 import {
+  ProbeContext,
   ProbeRunner
 } from "../src/ProbeRunner.js";
 import { SourceFile } from "../src/SourceFile.js";
@@ -89,6 +90,29 @@ describe("ProbeRunner", () => {
 
       assert.throws(instantiateProbeRunner, Error, "Invalid probe");
     });
+
+    it("should throw if one the provided probe is sealed or frozen", () => {
+      const methods = ["seal", "freeze"];
+      for (const method of methods) {
+        const fakeProbe = Object[method]({
+          name: "frozen-probe",
+          initialize() {
+            return {};
+          },
+          validateNode: mock.fn((_: ESTree.Node) => [true]),
+          main: () => ProbeRunner.Signals.Skip
+        });
+
+        assert.throws(() => {
+          new ProbeRunner(
+            new SourceFile(),
+            [fakeProbe]
+          );
+        }, {
+          message: "Failed to define original context for probe 'frozen-probe'"
+        });
+      }
+    });
   });
 
   describe("walk", () => {
@@ -299,5 +323,45 @@ describe("ProbeRunner", () => {
         expectedContext
       ]);
     });
+
+    it("should deep clone initialization context and clear context when the probe is fully executed", () => {
+      const fakeCtx = {};
+      const fakeProbe: any = {
+        initialize() {
+          return fakeCtx;
+        },
+        validateNode: mock.fn((_: ESTree.Node) => [true]),
+        main(_node: ESTree.Node, ctx: Required<ProbeContext>) {
+          ctx.context.hello = "world";
+
+          return ProbeRunner.Signals.Skip;
+        }
+      };
+
+      const sourceFile = new SourceFile();
+
+      const pr = new ProbeRunner(
+        sourceFile,
+        [fakeProbe]
+      );
+
+      const astNode: ESTree.Literal = {
+        type: "Literal",
+        value: "test"
+      };
+      pr.walk(astNode);
+      assert.deepEqual(fakeProbe.context, {
+        hello: "world"
+      });
+
+      pr.finalize();
+
+      assert.strictEqual(fakeProbe.context, undefined);
+      const { context = null } = fakeProbe.validateNode.mock.calls.at(0)?.arguments[1] ?? {};
+      assert.deepEqual(context, {
+        hello: "world"
+      });
+      assert.notStrictEqual(context, fakeCtx);
+    });
   });
 });