Add esm support to multiple sinks

end2end/tests/hono-sqlite3-esm.test.js

@@ -33,11 +33,13 @@ t.test("it blocks in blocking mode", (t) => {
 
   let stdout = "";
   server.stdout.on("data", (data) => {
+    console.log(data.toString());
     stdout += data.toString();
   });
 
   let stderr = "";
   server.stderr.on("data", (data) => {
+    console.log(data.toString());
     stderr += data.toString();
   });
 

library/agent/hooks/wrapImport.ts

@@ -68,10 +68,10 @@
 
     patchPackage(exports, name, baseDir);
   } catch (error) {
     if (error instanceof Error) {
       getInstance()?.onFailedToWrapModule(name, error);
     }
   }
 }
 
 function patchBuiltinModule(exports: any, name: string) {
@@ -96,8 +96,8 @@
 
 function getRelativeImportedFilePath(packageName: string, importName: string) {
   if (!importName.startsWith(packageName)) {
     return importName;
   }
   return importName.substring(packageName.length + 1);
 }
 
@@ -108,13 +108,13 @@
 ) {
   // Ignore .json files
   if (importName.endsWith(".json")) {
     return;
   }
 
   // Required for getting package version
   if (!baseDir) {
     throw new Error("Can not patch package without baseDir");
   }
 
   // Base dir and importName include the package name
   // The importName also includes the path to the file inside the imported package
@@ -122,9 +122,9 @@
   // Do not use pathInfo.base because its wrong, use baseDir instead!
   const pathInfoForName = getModuleInfoFromPath(`node_modules/${importName}`);
   if (!pathInfoForName) {
     // Can happen if the package is not inside a node_modules folder, like the dev build of our library itself
     return;
   }
   const packageName = pathInfoForName.name;
 
   const pathInfo: ModulePathInfo = {
@@ -141,8 +141,8 @@
 
   // We don't want to patch this package because we do not have any hooks for it
   if (!versionedPackages.length) {
     return;
   }
 
   // Get the original require function (if it was wrapped)
   const requireFunc = getOrignalRequire() || require;
@@ -153,10 +153,10 @@
   // Get the version of the installed package
   const installedPkgVersion = packageJson.version;
   if (!installedPkgVersion) {
     throw new Error(
       `Could not get installed package version for ${packageName} on import`
     );
   }
 
   // Check if the installed package version is supported (get all matching versioned packages)
   const matchingVersionedPackages = versionedPackages.filter((pkg) =>
@@ -165,17 +165,17 @@
 
   const agent = getInstance();
   if (agent) {
     // Report to the agent that the package was wrapped or not if it's version is not supported
     agent.onPackageWrapped(packageName, {
       version: installedPkgVersion,
       supported: !!matchingVersionedPackages.length,
     });
   }
 
   if (!matchingVersionedPackages.length) {
     // We don't want to patch this package version
     return;
   }
 
   const fullFilePath = join(pathInfo.base, pathInfo.path);
 
@@ -195,11 +195,11 @@
       .map((pkg) => pkg.getRequireInterceptors())
       .flat();
   } else {
-    // If its not the main file, we want to check if the want to patch the required file
+    // If its not the main file, we check if we want to patch the required file
     interceptors = matchingVersionedPackages
       .map((pkg) => pkg.getRequireFileInterceptor(pathInfo.path) || [])
       .flat();
   }
 
   executeInterceptors(interceptors, exports, undefined, undefined, {
     name: pathInfo.name,

library/agent/hooks/wrapRequire.ts

@@ -235,7 +235,7 @@ function patchPackage(this: mod, id: string, originalExports: unknown) {
       .map((pkg) => pkg.getRequireInterceptors())
       .flat();
   } else {
-    // If its not the main file, we want to check if the want to patch the required file
+    // If its not the main file, we check if we want to patch the required file
     interceptors = matchingVersionedPackages
       .map((pkg) => pkg.getRequireFileInterceptor(pathInfo.path) || [])
       .flat();

library/sinks/AwsSDKVersion2.test.ts

@@ -4,6 +4,7 @@ import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { Context, runWithContext } from "../agent/Context";
 import { LoggerForTesting } from "../agent/logger/LoggerForTesting";
 import { AwsSDKVersion2 } from "./AwsSDKVersion2";
+import { isCJS } from "../helpers/isCJS";
 
 // Suppress upgrade to SDK v3 notice
 require("aws-sdk/lib/maintenance_mode_message").suppress = true;
@@ -32,12 +33,13 @@ t.test("it works", async (t) => {
     logger,
     new ReportingAPIForTesting(),
     undefined,
-    undefined
+    undefined,
+    !isCJS()
   );
 
   agent.start([new AwsSDKVersion2()]);
 
-  const AWS = require("aws-sdk");
+  const AWS = isCJS() ? require("aws-sdk") : (await import("aws-sdk")).default;
 
   const s3 = new AWS.S3({
     region: "us-east-1",

library/sinks/AwsSDKVersion2.ts

@@ -73,7 +73,9 @@ export class AwsSDKVersion2 implements Wrapper {
       .addPackage("aws-sdk")
       .withVersion("^2.0.0")
       .onRequire((exports, pkgInfo) => {
-        wrapNewInstance(exports, "S3", pkgInfo, (instance) => {
+        const base = pkgInfo.isESMImport ? exports.default : exports;
+
+        wrapNewInstance(base, "S3", pkgInfo, (instance) => {
           for (const operation of operationsWithKey) {
             wrapExport(instance, operation, pkgInfo, {
               inspectArgs: (args) => this.inspectS3Operation(args, operation),

library/sinks/BetterSQLite3.test.ts

@@ -4,6 +4,7 @@ import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { runWithContext, type Context } from "../agent/Context";
 import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { BetterSQLite3 } from "./BetterSQLite3";
+import { isCJS } from "../helpers/isCJS";
 
 const dangerousContext: Context = {
   remoteAddress: "::1",
@@ -54,11 +55,14 @@ t.test("it detects SQL injections", async (t) => {
     new LoggerNoop(),
     new ReportingAPIForTesting(),
     undefined,
-    "lambda"
+    "lambda",
+    !isCJS()
   );
   agent.start([new BetterSQLite3()]);
 
-  const betterSqlite3 = require("better-sqlite3");
+  const betterSqlite3 = isCJS()
+    ? require("better-sqlite3")
+    : (await import("better-sqlite3")).default;
   const db = new betterSqlite3(":memory:");
 
   try {

library/sinks/BetterSQLite3.ts

@@ -69,15 +69,17 @@ export class BetterSQLite3 implements Wrapper {
       .addPackage("better-sqlite3")
       .withVersion("^11.0.0 || ^10.0.0 || ^9.0.0 || ^8.0.0")
       .onRequire((exports, pkgInfo) => {
+        const base = pkgInfo.isESMImport ? exports.default : exports;
+
         for (const func of sqlFunctions) {
-          wrapExport(exports.prototype, func, pkgInfo, {
+          wrapExport(base.prototype, func, pkgInfo, {
             inspectArgs: (args) => {
               return this.inspectQuery(`better-sqlite3.${func}`, args);
             },
           });
         }
         for (const func of fsPathFunctions) {
-          wrapExport(exports.prototype, func, pkgInfo, {
+          wrapExport(base.prototype, func, pkgInfo, {
             inspectArgs: (args) => {
               return this.inspectPath(`better-sqlite3.${func}`, args);
             },

library/sinks/ChildProcess.test.ts

@@ -4,7 +4,8 @@ import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { Context, runWithContext } from "../agent/Context";
 import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { ChildProcess } from "./ChildProcess";
-import { execFile, execFileSync, fork } from "child_process";
+import { execFile, execFileSync } from "child_process";
+import { isCJS } from "../helpers/isCJS";
 
 const unsafeContext: Context = {
   remoteAddress: "::1",
@@ -36,12 +37,15 @@ t.test("it works", async (t) => {
     new LoggerNoop(),
     new ReportingAPIForTesting(),
     undefined,
-    "lambda"
+    "lambda",
+    !isCJS()
   );
 
   agent.start([new ChildProcess()]);
 
-  const { exec, execSync, spawn, spawnSync, fork } = require("child_process");
+  const { exec, execSync, spawn, spawnSync, fork } = isCJS()
+    ? require("child_process")
+    : await import("child_process");
 
   const runCommandsWithInvalidArgs = () => {
     throws(

library/sinks/ChildProcess.ts

@@ -19,41 +19,47 @@ const PATH_PREFIXES = [
 export class ChildProcess implements Wrapper {
   wrap(hooks: Hooks) {
     hooks.addBuiltinModule("child_process").onRequire((exports, pkgInfo) => {
-      wrapExport(exports, "exec", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectExec(args, "exec");
-        },
-      });
-      wrapExport(exports, "execSync", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectExec(args, "execSync");
-        },
-      });
-      wrapExport(exports, "spawn", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectSpawn(args, "spawn");
-        },
-      });
-      wrapExport(exports, "spawnSync", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectSpawn(args, "spawnSync");
-        },
-      });
-      wrapExport(exports, "execFile", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectExecFile(args, "execFile");
-        },
-      });
-      wrapExport(exports, "execFileSync", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectExecFile(args, "execFileSync");
-        },
-      });
-      wrapExport(exports, "fork", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectFork(args, "fork");
-        },
-      });
+      const toWrap = pkgInfo.isESMImport
+        ? [exports, exports.default]
+        : [exports];
+
+      for (const exportObj of toWrap) {
+        wrapExport(exportObj, "exec", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectExec(args, "exec");
+          },
+        });
+        wrapExport(exportObj, "execSync", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectExec(args, "execSync");
+          },
+        });
+        wrapExport(exportObj, "spawn", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectSpawn(args, "spawn");
+          },
+        });
+        wrapExport(exportObj, "spawnSync", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectSpawn(args, "spawnSync");
+          },
+        });
+        wrapExport(exportObj, "execFile", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectExecFile(args, "execFile");
+          },
+        });
+        wrapExport(exportObj, "execFileSync", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectExecFile(args, "execFileSync");
+          },
+        });
+        wrapExport(exportObj, "fork", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectFork(args, "fork");
+          },
+        });
+      }
     });
   }
 

library/sinks/Fetch.test.ts

@@ -8,35 +8,39 @@ import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { wrap } from "../helpers/wrap";
 import { Fetch } from "./Fetch";
 import * as dns from "dns";
+import { isCJS } from "../helpers/isCJS";
 
 const calls: Record<string, number> = {};
-wrap(dns, "lookup", function lookup(original) {
-  return function lookup() {
-    const hostname = arguments[0];
 
-    if (!calls[hostname]) {
-      calls[hostname] = 0;
-    }
+if (isCJS()) {
+  wrap(dns, "lookup", function lookup(original) {
+    return function lookup() {
+      const hostname = arguments[0];
 
-    calls[hostname]++;
+      if (!calls[hostname]) {
+        calls[hostname] = 0;
+      }
 
-    if (hostname === "thisdomainpointstointernalip.com") {
-      return original.apply(this, [
-        "localhost",
-        ...Array.from(arguments).slice(1),
-      ]);
-    }
+      calls[hostname]++;
 
-    if (hostname === "example,prefix.thisdomainpointstointernalip.com") {
-      return original.apply(this, [
-        "localhost",
-        ...Array.from(arguments).slice(1),
-      ]);
-    }
+      if (hostname === "thisdomainpointstointernalip.com") {
+        return original.apply(this, [
+          "localhost",
+          ...Array.from(arguments).slice(1),
+        ]);
+      }
 
-    original.apply(this, arguments);
-  };
-});
+      if (hostname === "example,prefix.thisdomainpointstointernalip.com") {
+        return original.apply(this, [
+          "localhost",
+          ...Array.from(arguments).slice(1),
+        ]);
+      }
+
+      original.apply(this, arguments);
+    };
+  });
+}
 
 const context: Context = {
   remoteAddress: "::1",
@@ -76,7 +80,8 @@ t.test(
       new LoggerNoop(),
       api,
       new Token("123"),
-      undefined
+      undefined,
+      !isCJS()
     );
     agent.start([new Fetch()]);
 
@@ -154,47 +159,50 @@ t.test(
       }
     });
 
-    await runWithContext(
-      {
-        ...context,
-        ...{
-          body: {
-            image2: [
-              "http://example",
-              "prefix.thisdomainpointstointernalip.com",
-            ],
-            image: "http://thisdomainpointstointernalip.com/path",
+    // Because we need to wrap dns here in the test
+    if (isCJS()) {
+      await runWithContext(
+        {
+          ...context,
+          ...{
+            body: {
+              image2: [
+                "http://example",
+                "prefix.thisdomainpointstointernalip.com",
+              ],
+              image: "http://thisdomainpointstointernalip.com/path",
+            },
           },
         },
-      },
-      async () => {
-        const error = await t.rejects(() =>
-          fetch("http://thisdomainpointstointernalip.com")
-        );
-        if (error instanceof Error) {
-          t.same(
-            // @ts-expect-error Type is not defined
-            error.cause.message,
-            "Zen has blocked a server-side request forgery: fetch(...) originating from body.image"
+        async () => {
+          const error = await t.rejects(() =>
+            fetch("http://thisdomainpointstointernalip.com")
           );
-        }
+          if (error instanceof Error) {
+            t.same(
+              // @ts-expect-error Type is not defined
+              error.cause.message,
+              "Zen has blocked a server-side request forgery: fetch(...) originating from body.image"
+            );
+          }
 
-        const error2 = await t.rejects(() =>
-          fetch(["http://example", "prefix.thisdomainpointstointernalip.com"])
-        );
-        if (error2 instanceof Error) {
-          t.same(
-            // @ts-expect-error Type is not defined
-            error2.cause.message,
-            "Zen has blocked a server-side request forgery: fetch(...) originating from body.image2"
+          const error2 = await t.rejects(() =>
+            fetch(["http://example", "prefix.thisdomainpointstointernalip.com"])
           );
-        }
+          if (error2 instanceof Error) {
+            t.same(
+              // @ts-expect-error Type is not defined
+              error2.cause.message,
+              "Zen has blocked a server-side request forgery: fetch(...) originating from body.image2"
+            );
+          }
 
-        // Ensure the lookup is only called once per hostname
-        // Otherwise, it could be vulnerable to TOCTOU
-        t.same(calls["thisdomainpointstointernalip.com"], 1);
-      }
-    );
+          // Ensure the lookup is only called once per hostname
+          // Otherwise, it could be vulnerable to TOCTOU
+          t.same(calls["thisdomainpointstointernalip.com"], 1);
+        }
+      );
+    }
 
     await runWithContext(
       {