Run some unit tests with new instrumentation

Author: timokoessler

.github/workflows/unit-test.yml

@@ -64,7 +64,19 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        node-version: [16.x, 18.x, 20.x, 22.x, 24.x]
+        include:
+          - node-version: 16.x
+            new-instrumentation: false
+          - node-version: 18.x
+            new-instrumentation: false
+          - node-version: 20.x
+            new-instrumentation: false
+          - node-version: 22.x
+            new-instrumentation: false
+          - node-version: 24.x
+            new-instrumentation: false
+          - node-version: 24.x
+            new-instrumentation: true
     timeout-minutes: 10
     steps:
       - uses: actions/checkout@v4
@@ -93,6 +105,9 @@ jobs:
       - run: npm run install-lib-only
       - run: npm run build
       - run: npm run test:ci
+        if: ${{ matrix.new-instrumentation == 'false' }}
+      - run: npm run test:ci:new
+        if: ${{ matrix.new-instrumentation == 'true' }}
       - name: "Upload coverage"
         uses: codecov/codecov-action@0565863a31f2c772f9f0395002a31e3f06189574 # v5
         with:

library/agent/Agent.test.ts

@@ -19,6 +19,7 @@ import { Context } from "./Context";
 import { createTestAgent } from "../helpers/createTestAgent";
 import { setTimeout } from "node:timers/promises";
 import type { Response } from "./api/fetchBlockedLists";
+import { isNewInstrumentationUnitTest } from "../helpers/isNewInstrumentationUnitTest";
 
 let shouldOnlyAllowSomeIPAddresses = false;
 
@@ -123,11 +124,18 @@ t.test("it sends started event", async (t) => {
     },
   ]);
 
-  t.same(logger.getMessages(), [
-    "Starting agent v0.0.0...",
-    "Found token, reporting enabled!",
-    "mongodb@6.16.0 is supported!",
-  ]);
+  if (!isNewInstrumentationUnitTest()) {
+    t.same(logger.getMessages(), [
+      "Starting agent v0.0.0...",
+      "Found token, reporting enabled!",
+      "mongodb@6.16.0 is supported!",
+    ]);
+  } else {
+    t.same(logger.getMessages(), [
+      "Starting agent v0.0.0...",
+      "Found token, reporting enabled!",
+    ]);
+  }
 });
 
 t.test("it throws error if already started", async () => {

library/agent/Agent.ts

@@ -25,9 +25,9 @@ import { wrapInstalledPackages } from "./wrapInstalledPackages";
 import { Wrapper } from "./Wrapper";
 import { isAikidoCI } from "../helpers/isAikidoCI";
 import { AttackLogger } from "./AttackLogger";
-import { envToBool } from "../helpers/envToBool";
 import { Packages } from "./Packages";
 import { AIStatistics } from "./AIStatistics";
+import { isNewInstrumentationUnitTest } from "../helpers/isNewInstrumentationUnitTest";
 
 type WrappedPackage = { version: string | null; supported: boolean };
 
@@ -76,7 +76,7 @@ export class Agent {
       throw new Error("Serverless cannot be an empty string");
     }
 
-    if (envToBool(process.env.AIKIDO_TEST_NEW_INSTRUMENTATION)) {
+    if (isNewInstrumentationUnitTest()) {
       this.newInstrumentation = true;
     }
   }

library/agent/hooks/instrumentation/codeTransformation.ts

@@ -2,8 +2,8 @@ import type { PackageFileInstrumentationInstructionJSON } from "./types";
 // eslint-disable-next-line camelcase
 import { wasm_transform_code_str } from "./wasm/node_code_instrumentation";
 import { getSourceType, PackageLoadFormat } from "./getSourceType";
-import { envToBool } from "../../../helpers/envToBool";
 import { join } from "path";
+import { isNewInstrumentationUnitTest } from "../../../helpers/isNewInstrumentationUnitTest";
 
 export function transformCode(
   pkgName: string,
@@ -26,7 +26,7 @@ export function transformCode(
   }
 
   // Rewrite import path for unit tests if environment variable is set to true
-  if (envToBool(process.env.AIKIDO_TEST_NEW_INSTRUMENTATION)) {
+  if (isNewInstrumentationUnitTest()) {
     return result.replace(
       "@aikidosec/firewall/instrument/internals",
       join(__dirname, "injectedFunctions.ts")

library/helpers/isNewInstrumentationUnitTest.ts

@@ -0,0 +1,5 @@
+import { envToBool } from "./envToBool";
+
+export function isNewInstrumentationUnitTest(): boolean {
+  return envToBool(process.env.AIKIDO_TEST_NEW_INSTRUMENTATION);
+}

library/index.ts

@@ -24,8 +24,8 @@ if (!isNewHookSystemUsed()) {
       console.warn(
         "AIKIDO: Your application seems to be running in ESM mode. You need to use the new hook system to enable AIKIDO. Please refer to the documentation for more information."
       );
-      require("./agent/protect").protect();
     }
+    require("./agent/protect").protect();
   }
 }
 

library/package.json

@@ -136,6 +136,7 @@
     "test": "node ../scripts/run-tap.js",
     "test:ci": "CI=true node ../scripts/run-tap.js",
     "test:new": "node ../scripts/run-tap.js --test-new-instrumentation",
+    "test:ci:new": "node ../scripts/run-tap.js --ci --test-new-instrumentation",
     "build": "tsc -p tsconfig.build.json",
     "build:watch": "tsc --watch -p tsconfig.build.json",
     "lint": "npm run lint-eslint && npm run lint-tsc && npm run check-format",

package.json

@@ -14,6 +14,7 @@
     "test": "cd library && npm run test",
     "test:ci": "cd library && npm run test:ci",
     "test:new": "cd library && npm run test:new",
+    "test:ci:new": "cd library && npm run test:ci:new",
     "lint": "cd library && npm run lint",
     "end2end": "cd end2end && npm run test",
     "format": "prettier --write .",

scripts/run-tap.js

@@ -2,7 +2,6 @@ const { execSync } = require("child_process");
 
 const version = process.versions.node.split(".");
 const major = parseInt(version[0], 10);
-const minor = parseInt(version[1], 10);
 
 let args = "--allow-incomplete-coverage";
 
@@ -17,14 +16,30 @@ if (process.env.CI) {
 
 if (process.argv.includes("--test-new-instrumentation")) {
   process.env.AIKIDO_TEST_NEW_INSTRUMENTATION = "true";
+
+  if (major < 22) {
+    console.error(
+      "Error:: --test-new-instrumentation is not supported on Node.js versions below 22."
+    );
+    process.exit(1);
+  }
+
+  const excludedTestFilesForNewInstrumentation = [
+    "**/sinks/**",
+    "**/sources/**",
+  ];
+
+  for (const exclude of excludedTestFilesForNewInstrumentation) {
+    args += ` --exclude='${exclude}'`;
+  }
 }
 
-execSync(`tap ${args}`, {
+execSync(`tap run ${args}`, {
   stdio: "inherit",
   env: {
     ...process.env,
     AIKIDO_CI: "true",
-    // In v23 some sub-dependencies are calling require on a esm module triggering an experimental warning
+    // In v24 some sub-dependencies are calling require on a esm module triggering an experimental warning
     NODE_OPTIONS: major === 24 ? "--disable-warning=ExperimentalWarning" : "",
   },
 });