feat(ProbeRunner): implement context (#395)

fraxken · web-flow · commit fad019fd25c2 · 2025-07-16T19:53:58.000+02:00
Update workspaces/js-x-ray/src/probes/isSyncIO.ts

Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;

chore: add changeset
diff --git a/.changeset/quiet-files-melt.md b/.changeset/quiet-files-melt.md
@@ -0,0 +1,5 @@
+---
+"@nodesecure/js-x-ray": minor
+---
+
+Implement context for Probe and ProbeRunner
diff --git a/workspaces/js-x-ray/docs/AstAnalyser.md b/workspaces/js-x-ray/docs/AstAnalyser.md
@@ -143,11 +143,11 @@ Below a basic probe that detect a string assignation to `danger`:
 export const customProbes = [
   {
     name: "customProbeUnsafeDanger",
-    validateNode: (node, sourceFile) => [
+    validateNode: (node) => [
       node.type === "VariableDeclaration" && node.declarations[0].init.value === "danger"
     ],
-    main: (node, options) => {
-      const { sourceFile, data: calleeName } = options;
+    main: (node, ctx) => {
+      const { sourceFile, data: calleeName } = ctx;
       if (node.declarations[0].init.value === "danger") {
         sourceFile.warnings.push({
           kind: "unsafe-danger",
diff --git a/workspaces/js-x-ray/src/ProbeRunner.ts b/workspaces/js-x-ray/src/ProbeRunner.ts
@@ -24,24 +24,29 @@ import type { SourceFile } from "./SourceFile.js";
 import type { OptionalWarningName } from "./warnings.js";
 
 export type ProbeReturn = void | null | symbol;
-export type ProbeInitializeCallback = (sourceFile: SourceFile) => void;
-export type ProbeFinalizeCallback = (sourceFile: SourceFile) => void;
-export type ProbeMainCallback = (
-  node: any,
-  options: { sourceFile: SourceFile; data?: any; }
-) => ProbeReturn;
-export type ProbeTeardownCallback = (options: { sourceFile: SourceFile; }) => void;
-export type ProbeValidationCallback = (node: ESTree.Node, sourceFile: SourceFile) => [boolean, any?];
-
-export interface Probe {
+export type ProbeContextDef = Record<string, any>;
+export type ProbeContext<T extends ProbeContextDef = ProbeContextDef> = {
+  sourceFile: SourceFile;
+  context?: T;
+};
+
+export type ProbeValidationCallback<T extends ProbeContextDef = ProbeContextDef> = (
+  node: ESTree.Node, ctx: ProbeContext<T>
+) => [boolean, any?];
+
+export interface Probe<T extends ProbeContextDef = ProbeContextDef> {
   name: string;
-  initialize?: ProbeInitializeCallback;
-  finalize?: ProbeFinalizeCallback;
-  validateNode: ProbeValidationCallback | ProbeValidationCallback[];
-  main: ProbeMainCallback;
-  teardown?: ProbeTeardownCallback;
+  initialize?: (ctx: ProbeContext<T>) => void | ProbeContext;
+  finalize?: (ctx: ProbeContext<T>) => void;
+  validateNode: ProbeValidationCallback<T> | ProbeValidationCallback<T>[];
+  main: (
+    node: any,
+    ctx: ProbeContext<T> & { data?: any; }
+  ) => ProbeReturn;
+  teardown?: (ctx: ProbeContext<T>) => void;
   breakOnMatch?: boolean;
   breakGroup?: string;
+  context?: ProbeContext<T>;
 }
 
 export const ProbeSignals = Object.freeze({
@@ -97,28 +102,42 @@ export class ProbeRunner {
         `Invalid probe ${probe.name}: initialize must be a function or undefined`
       );
       if (probe.initialize) {
-        probe.initialize(sourceFile);
+        const context = probe.initialize(this.#getProbeContext(probe));
+        if (context) {
+          probe.context = context;
+        }
       }
     }
 
     this.probes = probes;
   }
 
+  #getProbeContext(
+    probe: Probe
+  ): ProbeContext {
+    return {
+      sourceFile: this.sourceFile,
+      context: probe.context
+    };
+  }
+
   #runProbe(
     probe: Probe,
     node: ESTree.Node
   ): ProbeReturn {
     const validationFns = Array.isArray(probe.validateNode) ?
       probe.validateNode : [probe.validateNode];
+    const ctx = this.#getProbeContext(probe);
+
     for (const validateNode of validationFns) {
       const [isMatching, data = null] = validateNode(
         node,
-        this.sourceFile
+        ctx
       );
 
       if (isMatching) {
         return probe.main(node, {
-          sourceFile: this.sourceFile,
+          ...ctx,
           data
         });
       }
@@ -158,9 +177,7 @@ export class ProbeRunner {
         }
       }
       finally {
-        if (probe.teardown) {
-          probe.teardown({ sourceFile: this.sourceFile });
-        }
+        probe.teardown?.(this.#getProbeContext(probe));
       }
     }
 
@@ -169,9 +186,7 @@ export class ProbeRunner {
 
   finalize(): void {
     for (const probe of this.probes) {
-      if (probe.finalize) {
-        probe.finalize(this.sourceFile);
-      }
+      probe.finalize?.(this.#getProbeContext(probe));
     }
   }
 }
diff --git a/workspaces/js-x-ray/src/probes/isFetch.ts b/workspaces/js-x-ray/src/probes/isFetch.ts
@@ -3,12 +3,13 @@ import { getCallExpressionIdentifier } from "@nodesecure/estree-ast-utils";
 import type { ESTree } from "meriyah";
 
 // Import Internal Dependencies
-import { SourceFile } from "../SourceFile.js";
+import type { ProbeContext } from "../ProbeRunner.js";
 
 function validateNode(
   node: ESTree.Node,
-  { tracer }: SourceFile
+  ctx: ProbeContext
 ): [boolean, any?] {
+  const { tracer } = ctx.sourceFile;
   const id = getCallExpressionIdentifier(node);
 
   if (id === null) {
@@ -20,13 +21,17 @@ function validateNode(
   return [data !== null && data.identifierOrMemberExpr === "fetch"];
 }
 
-function initialize(sourceFile: SourceFile) {
+function initialize(
+  ctx: ProbeContext
+) {
+  const { sourceFile } = ctx;
+
   sourceFile.tracer.trace("fetch", { followConsecutiveAssignment: true });
 }
 
 function main(
   _node: ESTree.Node,
-  { sourceFile }: { sourceFile: SourceFile; }
+  { sourceFile }: ProbeContext
 ) {
   sourceFile.flags.add("fetch");
 }
diff --git a/workspaces/js-x-ray/src/probes/isRequire/isRequire.ts b/workspaces/js-x-ray/src/probes/isRequire/isRequire.ts
@@ -10,16 +10,16 @@ import {
 import type { ESTree } from "meriyah";
 
 // Import Internal Dependencies
-import { ProbeSignals } from "../../ProbeRunner.js";
-import { SourceFile } from "../../SourceFile.js";
+import { ProbeSignals, type ProbeContext } from "../../ProbeRunner.js";
 import { isLiteral } from "../../types/estree.js";
 import { RequireCallExpressionWalker } from "./RequireCallExpressionWalker.js";
 import { generateWarning } from "../../warnings.js";
 
 function validateNodeRequire(
   node: ESTree.Node,
-  { tracer }: SourceFile
+  ctx: ProbeContext
 ): [boolean, any?] {
+  const { tracer } = ctx.sourceFile;
   const id = getCallExpressionIdentifier(node, {
     resolveCallExpression: false
   });
@@ -63,14 +63,14 @@ function validateNodeEvalRequire(
 }
 
 function teardown(
-  { sourceFile }: { sourceFile: SourceFile; }
+  ctx: ProbeContext
 ) {
-  sourceFile.dependencyAutoWarning = false;
+  ctx.sourceFile.dependencyAutoWarning = false;
 }
 
 function main(
   node: ESTree.CallExpression,
-  options: { sourceFile: SourceFile; data?: string; }
+  options: ProbeContext & { data?: string; }
 ) {
   const { sourceFile, data: calleeName } = options;
   const { tracer } = sourceFile;
diff --git a/workspaces/js-x-ray/src/probes/isSerializeEnv.ts b/workspaces/js-x-ray/src/probes/isSerializeEnv.ts
@@ -6,9 +6,8 @@ import {
 import type { ESTree } from "meriyah";
 
 // Import Internal Dependencies
-import { SourceFile } from "../SourceFile.js";
 import { generateWarning } from "../warnings.js";
-import { ProbeSignals } from "../ProbeRunner.js";
+import { ProbeSignals, type ProbeContext } from "../ProbeRunner.js";
 
 /**
  * @description Detect serialization of process.env which could indicate environment variable exfiltration
@@ -20,8 +19,10 @@ import { ProbeSignals } from "../ProbeRunner.js";
  */
 function validateNode(
   node: ESTree.Node,
-  { tracer }: SourceFile
+  ctx: ProbeContext
 ): [boolean, any?] {
+  const { tracer } = ctx.sourceFile;
+
   const id = getCallExpressionIdentifier(node);
 
   if (id === null) {
@@ -58,9 +59,9 @@ function validateNode(
 
 function main(
   node: ESTree.Node,
-  options: { sourceFile: SourceFile; }
+  ctx: ProbeContext
 ) {
-  const { sourceFile } = options;
+  const { sourceFile } = ctx;
 
   const warning = generateWarning("serialize-environment", {
     value: "JSON.stringify(process.env)",
@@ -72,8 +73,10 @@ function main(
 }
 
 function initialize(
-  { tracer }: SourceFile
+  ctx: ProbeContext
 ) {
+  const { tracer } = ctx.sourceFile;
+
   tracer
     .trace("process.env", {
       followConsecutiveAssignment: true
diff --git a/workspaces/js-x-ray/src/probes/isSyncIO.ts b/workspaces/js-x-ray/src/probes/isSyncIO.ts
@@ -3,7 +3,7 @@ import { getCallExpressionIdentifier } from "@nodesecure/estree-ast-utils";
 import type { ESTree } from "meriyah";
 
 // Import Internal Dependencies
-import { SourceFile } from "../SourceFile.js";
+import type { ProbeContext } from "../ProbeRunner.js";
 import { generateWarning } from "../warnings.js";
 
 // CONSTANTS
@@ -40,8 +40,9 @@ const kSyncIOIdentifierOrMemberExps = [
 
 function validateNode(
   node: ESTree.Node,
-  { tracer }: SourceFile
+  ctx: ProbeContext
 ): [boolean, any?] {
+  const { tracer } = ctx.sourceFile;
   const id = getCallExpressionIdentifier(
     node,
     {
@@ -64,12 +65,12 @@ function validateNode(
 }
 
 function initialize(
-  sourceFile: SourceFile
+  ctx: ProbeContext
 ) {
   kSyncIOIdentifierOrMemberExps.forEach((identifierOrMemberExp) => {
     const moduleName = identifierOrMemberExp.split(".")[0];
 
-    return sourceFile.tracer.trace(identifierOrMemberExp, {
+    ctx.sourceFile.tracer.trace(identifierOrMemberExp, {
       followConsecutiveAssignment: true,
       moduleName
     });
@@ -78,13 +79,13 @@ function initialize(
 
 function main(
   node: ESTree.CallExpression,
-  { sourceFile }: { sourceFile: SourceFile; }
+  ctx: ProbeContext
 ) {
   const warning = generateWarning("synchronous-io", {
     value: node.callee.name,
     location: node.loc
   });
-  sourceFile.warnings.push(warning);
+  ctx.sourceFile.warnings.push(warning);
 }
 
 export default {
diff --git a/workspaces/js-x-ray/src/probes/isWeakCrypto.ts b/workspaces/js-x-ray/src/probes/isWeakCrypto.ts
@@ -3,7 +3,7 @@ import { getCallExpressionIdentifier } from "@nodesecure/estree-ast-utils";
 import type { ESTree } from "meriyah";
 
 // Import Internal Dependencies
-import { SourceFile } from "../SourceFile.js";
+import type { ProbeContext } from "../ProbeRunner.js";
 import { generateWarning } from "../warnings.js";
 import {
   isLiteral
@@ -20,9 +20,9 @@ const kWeakAlgorithms = new Set([
 
 function validateNode(
   node: ESTree.Node,
-  sourceFile: SourceFile
+  ctx: ProbeContext
 ): [boolean, any?] {
-  const { tracer } = sourceFile;
+  const { tracer } = ctx.sourceFile;
 
   const id = getCallExpressionIdentifier(node);
   if (id === null || !tracer.importedModules.has("crypto")) {
@@ -37,18 +37,21 @@ function validateNode(
 }
 
 function initialize(
-  sourceFile: SourceFile
+  ctx: ProbeContext
 ) {
-  sourceFile.tracer.trace("crypto.createHash", {
+  const { tracer } = ctx.sourceFile;
+
+  tracer.trace("crypto.createHash", {
     followConsecutiveAssignment: true,
     moduleName: "crypto"
   });
 }
 
 function main(
   node: ESTree.CallExpression,
-  { sourceFile }: { sourceFile: SourceFile; }
+  ctx: ProbeContext
 ) {
+  const { sourceFile } = ctx;
   const arg = node.arguments.at(0);
 
   if (isLiteral(arg) && kWeakAlgorithms.has(arg.value)) {
diff --git a/workspaces/js-x-ray/test/ProbeRunner.spec.ts b/workspaces/js-x-ray/test/ProbeRunner.spec.ts

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +---
 +"@nodesecure/js-x-ray": minor
 +---
++
 +Implement context for Probe and ProbeRunner