Prevent double wrapping, increase test coverage

Author: timokoessler

library/agent/hooks/instrumentation/builtinShim.test.ts

@@ -66,6 +66,9 @@ t.test(
 `.replace(/\s+/g, " ")
     );
 
+    const nonExistingModuleShim = generateBuildinShim("abc", "abc", false);
+    t.equal(nonExistingModuleShim, undefined);
+
     t.end();
   }
 );

library/agent/hooks/instrumentation/getExportsForBuiltin.test.ts

@@ -43,5 +43,7 @@ t.test(
         // ...
       ])
     );
+
+    t.match(getExportsForBuiltin("abc"), undefined);
   }
 );

library/agent/hooks/instrumentation/getExportsForBuiltin.ts

@@ -2,7 +2,7 @@ export function getExportsForBuiltin(name: string) {
   const mod = process.getBuiltinModule(name);
 
   if (!mod) {
-    return new Set();
+    return undefined;
   }
 
   return new Set(["default", ...Object.keys(mod)]);

library/agent/hooks/instrumentation/getPackageVersionFromPath.test.ts

@@ -0,0 +1,12 @@
+import * as t from "tap";
+import { getPackageVersionFromPath } from "./getPackageVersionFromPath";
+import { join } from "path";
+
+t.test("it works", async (t) => {
+  t.equal(getPackageVersionFromPath(join(__dirname, "../../..")), "0.0.0");
+  t.equal(
+    getPackageVersionFromPath(join(__dirname, "../../../node_modules/accepts")),
+    "2.0.0"
+  );
+  t.equal(getPackageVersionFromPath("test123"), undefined);
+});

library/agent/hooks/instrumentation/index.ts

@@ -2,11 +2,19 @@ import { onModuleLoad } from "./loadHook";
 import * as mod from "node:module";
 import type { RegisterHookFunction } from "./types";
 
+let hooksRegistered = false;
+
 export function registerNodeHooks() {
   if (!("registerHooks" in mod) || typeof mod.registerHooks !== "function") {
     throw new Error("This Node.js version is not supported");
   }
 
+  // Do not register hooks multiple times
+  if (hooksRegistered) {
+    return;
+  }
+  hooksRegistered = true;
+
   // Hook into the ESM & CJS module loading process
   // Types are required because official Node.js typings are not up-to-date
   (mod.registerHooks as RegisterHookFunction)({

library/agent/hooks/instrumentation/injectedFunctions.ts

@@ -56,6 +56,10 @@ export function __wrapBuiltinExports(id: string, exports: unknown): unknown {
 }
 
 export function __instrumentModifyArgs(id: string, args: unknown[]): unknown[] {
+  if (!Array.isArray(args)) {
+    return [];
+  }
+
   const agent = getInstance();
   if (!agent) {
     return args;

library/agent/hooks/instrumentation/instructions.test.ts

@@ -12,6 +12,7 @@ import { Package } from "../Package";
 import { BuiltinModule } from "../BuiltinModule";
 import {
   __instrumentInspectArgs,
+  __instrumentModifyArgs,
   __wrapBuiltinExports,
 } from "./injectedFunctions";
 import { createTestAgent } from "../../../helpers/createTestAgent";
@@ -119,6 +120,7 @@ t.test("it works", async (t) => {
 t.test("it works using injected functions", async (t) => {
   let pkgInspectArgsCalled = false;
   let builtinOnRequireCalled = false;
+  let pkgModifyArgsCalled = false;
 
   const pkg = new Package("foo");
   pkg.withVersion("^1.0.0").addFileInstrumentation({
@@ -130,6 +132,10 @@ t.test("it works using injected functions", async (t) => {
         inspectArgs: () => {
           pkgInspectArgsCalled = true;
         },
+        modifyArgs: (args) => {
+          pkgModifyArgsCalled = true;
+          return [42];
+        },
       },
     ],
   });
@@ -145,17 +151,96 @@ t.test("it works using injected functions", async (t) => {
 
   t.equal(pkgInspectArgsCalled, false);
   __instrumentInspectArgs("foo.bar.js.bazABCDEF.^1.0.0", []);
+  __instrumentModifyArgs("foo.bar.js.bazABCDEF.^1.0.0", []);
   t.equal(pkgInspectArgsCalled, false);
+  t.equal(pkgModifyArgsCalled, false);
   __instrumentInspectArgs("foo.bar.js.baz.^1.0.0", []);
+  __instrumentModifyArgs("foo.bar.js.baz.^1.0.0", []);
   // No agent yet
   t.equal(pkgInspectArgsCalled, false);
+  t.equal(pkgModifyArgsCalled, false);
+
+  createTestAgent();
+
+  __instrumentInspectArgs("foo.bar.js.bazABCDEF.^1.0.0", []);
+  __instrumentModifyArgs("foo.bar.js.bazABCDEF.^1.0.0", []);
+  t.equal(pkgInspectArgsCalled, false);
+  t.equal(pkgModifyArgsCalled, false);
 
-  const agent = createTestAgent();
   __instrumentInspectArgs("foo.bar.js.baz.^1.0.0", []);
   t.equal(pkgInspectArgsCalled, true);
+  t.same(__instrumentModifyArgs("foo.bar.js.baz.^1.0.0", []), [42]);
+  t.equal(pkgModifyArgsCalled, true);
 
   t.equal(builtinOnRequireCalled, false);
   const wrapped = __wrapBuiltinExports("http", {}) as any;
   t.equal(builtinOnRequireCalled, true);
   t.equal(wrapped.test, 42);
 });
+
+t.test("modifyArgs always returns a array", async (t) => {
+  const pkg = new Package("foo");
+  pkg.withVersion("^1.0.0").addFileInstrumentation({
+    path: "xyz.js",
+    functions: [
+      {
+        nodeType: "MethodDefinition",
+        name: "abc",
+        modifyArgs: (args) => {
+          return args;
+        },
+      },
+      {
+        nodeType: "MethodDefinition",
+        name: "xyz",
+        // @ts-expect-error Testing invalid input
+        modifyArgs: (args) => {
+          return undefined;
+        },
+      },
+    ],
+  });
+
+  setPackagesToInstrument([pkg]);
+  createTestAgent();
+
+  t.same(__instrumentModifyArgs("foo.xyz.js.abc.^1.0.0", [1, 2, 3]), [1, 2, 3]);
+  // @ts-expect-error Testing invalid input
+  t.same(__instrumentModifyArgs("foo.xyz.js.abc.^1.0.0", undefined), []);
+  t.same(
+    __instrumentModifyArgs("foo.xyz.js.doesnotexist", [1, 2, 3]),
+    [1, 2, 3]
+  );
+  t.same(__instrumentModifyArgs("foo.xyz.js.xyz.^1.0.0", [1, 2, 3]), [1, 2, 3]);
+  // @ts-expect-error Testing invalid input
+  t.same(__instrumentModifyArgs("foo.xyz.js.xyz.^1.0.0", undefined), []);
+});
+
+t.test("all injected functions handle errors", async (t) => {
+  const pkg = new Package("foo");
+  pkg.withVersion("^1.0.0").addFileInstrumentation({
+    path: "dist/test.mjs",
+    functions: [
+      {
+        nodeType: "MethodDefinition",
+        name: "abc",
+        inspectArgs: () => {
+          throw new Error("test");
+        },
+        modifyArgs: () => {
+          throw new Error("test");
+        },
+        modifyReturnValue: () => {
+          throw new Error("test");
+        },
+      },
+    ],
+  });
+
+  setPackagesToInstrument([pkg]);
+  createTestAgent();
+
+  __instrumentInspectArgs("foo.dist/test.mjs.abc.^1.0.0", []);
+  __instrumentModifyArgs("foo.dist/test.mjs.abc.^1.0.0", []);
+  t.same(__instrumentModifyArgs("foo.dist/test.mjs.abc.^1.0.0", []), []);
+});

library/agent/hooks/instrumentation/loadHook.test.ts

@@ -3,6 +3,23 @@ import { createTestAgent } from "../../../helpers/createTestAgent";
 import { applyHooks } from "../../applyHooks";
 import { Hooks } from "../Hooks";
 import * as mod from "node:module";
+import { registerNodeHooks } from ".";
+
+t.test(
+  "it throws an error if Node.js version is not supported",
+  {
+    skip:
+      "registerHooks" in mod ? "Only test on older Node.js versions" : false,
+  },
+  async (t) => {
+    createTestAgent();
+    const error = t.throws(() => applyHooks(new Hooks(), true));
+    t.ok(error instanceof Error);
+    if (error instanceof Error) {
+      t.match(error.message, /This Node.js version is not supported/);
+    }
+  }
+);
 
 t.test(
   "it works",
@@ -54,6 +71,9 @@ t.test(
       }
     }
 
+    // Try to register hooks again to test if it still works
+    registerNodeHooks();
+
     try {
       await import("hono");
       t.fail("import should not work");
@@ -99,6 +119,10 @@ t.test(
     t.equal(typeof http.createServer, "function");
     t.equal(typeof http.Server, "function");
 
+    // Require unpatched module
+    const assert = require("assert") as typeof import("assert");
+    t.equal(typeof assert, "function");
+
     process.env.AIKIDO_UNIT_TEST = "false";
   }
 );

library/helpers/createTestAgent.ts

@@ -17,6 +17,7 @@ export function createTestAgent(opts?: {
   token?: Token;
   serverless?: string;
   suppressConsoleLog?: boolean;
+  newInstrumentation?: boolean;
 }) {
   if (opts?.suppressConsoleLog ?? true) {
     wrap(console, "log", function log() {
@@ -29,7 +30,8 @@ export function createTestAgent(opts?: {
     opts?.logger ?? new LoggerNoop(),
     opts?.api ?? new ReportingAPIForTesting(),
     opts?.token, // Defaults to undefined
-    opts?.serverless // Defaults to undefined
+    opts?.serverless, // Defaults to undefined
+    opts?.newInstrumentation ?? false
   );
 
   setInstance(agent);