kickstart custom probe as param

Author: tchapacan

src/AstAnalyser.js

@@ -16,8 +16,9 @@ export class AstAnalyser {
    * @constructor
    * @param { SourceParser } [parser]
    */
-  constructor(parser = new JsSourceParser()) {
+  constructor(parser = new JsSourceParser(), customProbes = []) {
     this.parser = parser;
+    this.customProbes = customProbes;
   }
 
   analyse(str, options = Object.create(null)) {
@@ -31,7 +32,7 @@ export class AstAnalyser {
       isEcmaScriptModule: Boolean(module)
     });
 
-    const source = new SourceFile(str);
+    const source = new SourceFile(str, this.customProbes);
 
     // we walk each AST Nodes, this is a purely synchronous I/O
     walk(body, {

src/SourceFile.js

@@ -20,14 +20,20 @@ export class SourceFile {
   encodedLiterals = new Map();
   warnings = [];
 
-  constructor(sourceCodeString) {
+  constructor(sourceCodeString, customProbes = []) {
     this.tracer = new VariableTracer()
       .enableDefaultTracing()
       .trace("crypto.createHash", {
         followConsecutiveAssignment: true, moduleName: "crypto"
       });
 
-    this.probesRunner = new ProbeRunner(this);
+    if (Array.isArray(customProbes) && customProbes.length > 0) {
+      this.probesRunner = new ProbeRunner(this, customProbes);
+    }
+    else {
+      this.probesRunner = new ProbeRunner(this);
+    }
+
     if (trojan.verify(sourceCodeString)) {
       this.addWarning("obfuscated-code", "trojan-source");
     }

test/issues/221-inject-custom-probes.spec.js

@@ -0,0 +1,40 @@
+// Import Node.js Dependencies
+import { test } from "node:test";
+import assert from "node:assert";
+
+// Import Internal Dependencies
+import { JsSourceParser } from "../../src/JsSourceParser.js";
+import { AstAnalyser } from "../../src/AstAnalyser.js";
+import { ProbeSignals } from "../../src/ProbeRunner.js";
+
+/**
+ * @see https://github.com/NodeSecure/js-x-ray/issues/221
+ */
+// CONSTANTS
+const kIncriminedCodeSample = "const danger = 'danger';";
+const kWarningUnsafeDanger = "unsafe-danger";
+
+test("should detect a custom probe alert unsafe-danger", () => {
+  const customProbes = [
+    {
+      name: "customProbeUnsafeDanger",
+      validateNode: (node, sourceFile) => [true]
+      ,
+      main: (node, options) => {
+        const { sourceFile, data: calleeName } = options;
+        if (node.declarations[0].init.value === "danger") {
+          sourceFile.addWarning("unsafe-danger", calleeName, node.loc);
+
+          return ProbeSignals.Skip;
+        }
+
+        return null;
+      }
+    }
+  ];
+
+  const analyser = new AstAnalyser(new JsSourceParser(), customProbes);
+  const result = analyser.analyse(kIncriminedCodeSample);
+
+  assert.equal(result.warnings[0].kind, kWarningUnsafeDanger);
+});

test/utils/index.js

@@ -38,7 +38,8 @@ export function getSastAnalysis(
       return this.sourceFile.dependencies;
     },
     execute(body) {
-      const probeRunner = new ProbeRunner(this.sourceFile, [probe]);
+      const probes = Array.isArray(probe) ? probe : [probe];
+      const probeRunner = new ProbeRunner(this.sourceFile, probes);
       const self = this;
 
       walk(body, {