Add esm support to multiple sinks

AikidoSec · Oct 4, 2024 · 50863a7 · 50863a7
1 parent f18568c
commit 50863a7
Show file tree

Hide file tree

Showing 10 changed files with 135 additions and 105 deletions.
diff --git a/end2end/tests/hono-sqlite3-esm.test.js b/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();
   });
 

diff --git a/library/agent/hooks/wrapImport.ts b/library/agent/hooks/wrapImport.ts
@@ -195,7 +195,7 @@ function patchPackage(
       .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();

diff --git a/library/agent/hooks/wrapRequire.ts b/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();

diff --git a/library/sinks/AwsSDKVersion2.test.ts b/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",

diff --git a/library/sinks/AwsSDKVersion2.ts b/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),

diff --git a/library/sinks/BetterSQLite3.test.ts b/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 {

diff --git a/library/sinks/BetterSQLite3.ts b/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);
             },

diff --git a/library/sinks/ChildProcess.test.ts b/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(

diff --git a/library/sinks/ChildProcess.ts b/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");
+          },
+        });
+      }
     });
   }
 

diff --git a/library/sinks/Fetch.test.ts b/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(
       {