feat(tracer): add follow consecutive assignment of return values (#385)

Author: clemgbld

.changeset/beige-insects-taste.md

@@ -0,0 +1,5 @@
+---
+"@nodesecure/tracer": minor
+---
+
+feat(tracer): add follow consecutive assignment of return values

workspaces/tracer/src/VariableTracer.ts

@@ -78,6 +78,7 @@ export class VariableTracer extends EventEmitter {
   #traced = new Map<string, Traced>();
   #variablesRefToGlobal = new Set<string>();
   #neutralCallable = new Set<string>();
+  #assignedReturnValueToTraced = new Map<string, string>();
 
   debug() {
     console.log(this.#traced);
@@ -442,6 +443,9 @@ export class VariableTracer extends EventEmitter {
             type: "ReturnValueAssignment",
             name: id.name
           });
+          if (tracedVariant.followConsecutiveAssignment) {
+            this.#assignedReturnValueToTraced.set(id.name, tracedFullIdentifierName);
+          }
         }
         // const id = Function.prototype.call.call(require, require, "http");
         if (this.#neutralCallable.has(identifierName) || isEvilIdentifierPath(fullIdentifierPath)) {
@@ -489,6 +493,16 @@ export class VariableTracer extends EventEmitter {
           this.#variablesRefToGlobal.add(id.name);
         }
 
+        if (this.#assignedReturnValueToTraced.has(childNode.name)) {
+          const tracedFullIdentifierName = this.#assignedReturnValueToTraced.get(childNode.name)!;
+          const tracedVariant = this.#traced.get(tracedFullIdentifierName)!;
+          tracedVariant.assignmentMemory.push({
+            type: "ReturnValueAssignment",
+            name: id.name
+          });
+          this.#assignedReturnValueToTraced.set(id.name, tracedFullIdentifierName);
+        }
+
         break;
       }
 

workspaces/tracer/test/VariableTracer.spec.ts

@@ -365,3 +365,80 @@ test("should be able to trace the return value of a traced function in an array
     name: "user"
   }]);
 });
+
+test("should be able to follow re-assignment on traced return values", () => {
+  const helpers = createTracer(false);
+  helpers.tracer.trace("os.userInfo", {
+    followConsecutiveAssignment: true,
+    followReturnValueAssignement: true,
+    moduleName: "os"
+  });
+
+  helpers.walkOnCode(`
+    import os from "os";
+
+    const user = [...os.userInfo()];
+    const userBis = user;
+    console.log(userBis);
+  `);
+
+  assert.deepEqual(helpers.tracer.getDataFromIdentifier("os.userInfo")?.assignmentMemory, [{
+    type: "ReturnValueAssignment",
+    name: "user"
+  }, {
+    type: "ReturnValueAssignment",
+    name: "userBis"
+  }]);
+});
+
+test("should be able to follow re-assignment on multiple consecutive traced return values", () => {
+  const helpers = createTracer(false);
+  helpers.tracer.trace("os.userInfo", {
+    followConsecutiveAssignment: true,
+    followReturnValueAssignement: true,
+    moduleName: "os"
+  });
+
+  helpers.walkOnCode(`
+    import os from "os";
+
+    const user = [...os.userInfo()];
+    const userBis = {...user};
+    const userTer = userBis;
+    console.log(userTer);
+  `);
+
+  assert.deepEqual(helpers.tracer.getDataFromIdentifier("os.userInfo")?.assignmentMemory, [{
+    type: "ReturnValueAssignment",
+    name: "user"
+  }, {
+    type: "ReturnValueAssignment",
+    name: "userBis"
+  },
+  {
+    type: "ReturnValueAssignment",
+    name: "userTer"
+  }
+  ]);
+});
+
+test("should not be able to follow re-assignment of a traced return value when followConsecutiveAssignment is not on", () => {
+  const helpers = createTracer(false);
+  helpers.tracer.trace("os.userInfo", {
+    followReturnValueAssignement: true,
+    moduleName: "os"
+  });
+
+  helpers.walkOnCode(`
+    import os from "os";
+
+    const user = [...os.userInfo()];
+    const userBis = user;
+    console.log(userBis);
+  `);
+
+  assert.deepEqual(helpers.tracer.getDataFromIdentifier("os.userInfo")?.assignmentMemory, [{
+    type: "ReturnValueAssignment",
+    name: "user"
+  }]);
+});