WIP for running unit tests using registerHooks

library/agent/Agent.ts

@@ -25,6 +25,7 @@ import { wrapInstalledPackages } from "./wrapInstalledPackages";
 import { Wrapper } from "./Wrapper";
 import { isAikidoCI } from "../helpers/isAikidoCI";
 import { AttackLogger } from "./AttackLogger";
+import { envToBool } from "../helpers/envToBool";
 
 type WrappedPackage = { version: string | null; supported: boolean };
 
@@ -70,6 +71,10 @@ export class Agent {
     if (typeof this.serverless === "string" && this.serverless.length === 0) {
       throw new Error("Serverless cannot be an empty string");
     }
+
+    if (envToBool(process.env.AIKIDO_TEST_NEW_INSTRUMENTATION)) {
+      this.newInstrumentation = true;
+    }
   }
 
   shouldBlock() {
@@ -567,4 +572,8 @@ export class Agent {
   onMiddlewareExecuted() {
     this.middlewareInstalled = true;
   }
+
+  isUsingNewInstrumentation() {
+    return this.newInstrumentation;
+  }
 }

library/agent/hooks/instrumentation/index.ts

@@ -2,6 +2,8 @@ import { onModuleLoad } from "./loadHook";
 import * as mod from "node:module";
 import type { RegisterHookFunction } from "./types";
 import { patchProcessGetBuiltinModule } from "./processGetBuiltin";
+import { join } from "node:path";
+import { envToBool } from "../../../helpers/envToBool";
 
 let hooksRegistered = false;
 
@@ -25,5 +27,18 @@ export function registerNodeHooks() {
     },
   });
 
+  // DEV rewrite for unit tests
+  if (envToBool(process.env.AIKIDO_TEST_NEW_INSTRUMENTATION)) {
+    (mod.registerHooks as RegisterHookFunction)({
+      resolve(specifier, context, nextResolve) {
+        if (specifier === "@aikidosec/firewall/instrument/internals") {
+          specifier = join(__dirname, "injectedFunctions.ts");
+        }
+
+        return nextResolve(specifier, context);
+      },
+    });
+  }
+
   patchProcessGetBuiltinModule();
 }

library/helpers/envToBool.ts

@@ -3,9 +3,9 @@ const trueValues = ["true", "1", "yes", "y", "on"];
 /**
  * Parses the string value of an environment variable to a boolean.
  */
-export function envToBool(envName: string | undefined): boolean {
-  if (!envName) {
+export function envToBool(env: string | undefined): boolean {
+  if (!env) {
     return false;
   }
-  return trueValues.includes(envName.toLowerCase());
+  return trueValues.includes(env.toLowerCase());
 }

library/helpers/startTestAgent.ts

@@ -23,13 +23,17 @@ export function startTestAgent(opts: {
   const agent = createTestAgent(opts);
   agent.start(opts.wrappers);
 
-  // In order to support multiple versions of the same package, we need to rewrite the package name
-  // e.g. In our sources and sinks, we use the real package name `hooks.addPackage("undici")`
-  // but in the tests we want to `require("undici-v6")` instead of `require("undici")`
-  // The `__internalRewritePackageName` function allows us to do this
-  Object.keys(opts.rewrite).forEach((packageName) => {
-    __internalRewritePackageName(packageName, opts.rewrite[packageName]);
-  });
+  if (!agent.isUsingNewInstrumentation()) {
+    // In order to support multiple versions of the same package, we need to rewrite the package name
+    // e.g. In our sources and sinks, we use the real package name `hooks.addPackage("undici")`
+    // but in the tests we want to `require("undici-v6")` instead of `require("undici")`
+    // The `__internalRewritePackageName` function allows us to do this
+    Object.keys(opts.rewrite).forEach((packageName) => {
+      __internalRewritePackageName(packageName, opts.rewrite[packageName]);
+    });
+  }
+
+  // Todo support this in the new instrumentation
 
   return agent;
 }