Test more sinks with ESM

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

Hide file tree

Showing 6 changed files with 190 additions and 133 deletions.
diff --git a/library/sinks/FileSystem.test.ts b/library/sinks/FileSystem.test.ts
@@ -4,6 +4,7 @@ import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { Context, runWithContext } from "../agent/Context";
 import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { FileSystem } from "./FileSystem";
+import { isCJS } from "../helpers/isCJS";
 
 const unsafeContext: Context = {
   remoteAddress: "::1",
@@ -52,7 +53,8 @@ t.test("it works", async (t) => {
     new LoggerNoop(),
     new ReportingAPIForTesting(),
     undefined,
-    "lambda"
+    "lambda",
+    !isCJS()
   );
 
   agent.start([new FileSystem()]);
@@ -64,8 +66,10 @@ t.test("it works", async (t) => {
     realpath,
     promises: fsDotPromise,
     realpathSync,
-  } = require("fs");
-  const { writeFile: writeFilePromise } = require("fs/promises");
+  } = isCJS() ? require("fs") : await import("fs");
+  const { writeFile: writeFilePromise } = isCJS()
+    ? require("fs/promises")
+    : await import("fs/promises");
 
   t.ok(typeof realpath.native === "function");
   t.ok(typeof realpathSync.native === "function");

diff --git a/library/sinks/FileSystem.ts b/library/sinks/FileSystem.ts
@@ -85,48 +85,60 @@ export class FileSystem implements Wrapper {
   wrap(hooks: Hooks) {
     // Wrap fs
     hooks.addBuiltinModule("fs").onRequire((exports, pkgInfo) => {
-      Object.keys(functions).forEach((name) => {
-        const { pathsArgs, sync, promise } = functions[name];
+      const toWrap = pkgInfo.isESMImport
+        ? [exports, exports.default]
+        : [exports];
 
-        wrapExport(exports, name, pkgInfo, {
-          inspectArgs: (args) => {
-            return this.inspectPath(args, name, pathsArgs);
-          },
-        });
+      for (const exportObj of toWrap) {
+        Object.keys(functions).forEach((name) => {
+          const { pathsArgs, sync, promise } = functions[name];
 
-        if (sync) {
-          wrapExport(exports, `${name}Sync`, pkgInfo, {
+          wrapExport(exportObj, name, pkgInfo, {
             inspectArgs: (args) => {
-              return this.inspectPath(args, `${name}Sync`, pathsArgs);
+              return this.inspectPath(args, name, pathsArgs);
             },
           });
-        }
-      });
 
-      // Wrap realpath.native
-      wrapExport(exports.realpath, "native", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectPath(args, "realpath.native", 1);
-        },
-      });
-      wrapExport(exports.realpathSync, "native", pkgInfo, {
-        inspectArgs: (args) => {
-          return this.inspectPath(args, "realpathSync.native", 1);
-        },
-      });
+          if (sync) {
+            wrapExport(exportObj, `${name}Sync`, pkgInfo, {
+              inspectArgs: (args) => {
+                return this.inspectPath(args, `${name}Sync`, pathsArgs);
+              },
+            });
+          }
+        });
+
+        // Wrap realpath.native
+        wrapExport(exportObj.realpath, "native", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectPath(args, "realpath.native", 1);
+          },
+        });
+        wrapExport(exportObj.realpathSync, "native", pkgInfo, {
+          inspectArgs: (args) => {
+            return this.inspectPath(args, "realpathSync.native", 1);
+          },
+        });
+      }
     });
 
     // Wrap fs/promises
     hooks.addBuiltinModule("fs/promises").onRequire((exports, pkgInfo) => {
-      Object.keys(functions).forEach((name) => {
-        const { pathsArgs, sync, promise } = functions[name];
+      const toWrap = pkgInfo.isESMImport
+        ? [exports, exports.default]
+        : [exports];
 
-        if (promise) {
-          wrapExport(exports, name, pkgInfo, {
-            inspectArgs: (args) => this.inspectPath(args, name, pathsArgs),
-          });
-        }
-      });
+      for (const exportObj of toWrap) {
+        Object.keys(functions).forEach((name) => {
+          const { pathsArgs, promise } = functions[name];
+
+          if (promise) {
+            wrapExport(exportObj, name, pkgInfo, {
+              inspectArgs: (args) => this.inspectPath(args, name, pathsArgs),
+            });
+          }
+        });
+      }
     });
   }
 }
diff --git a/library/sinks/HTTPRequest.test.ts b/library/sinks/HTTPRequest.test.ts
@@ -8,31 +8,35 @@ import { Context, runWithContext } from "../agent/Context";
 import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { wrap } from "../helpers/wrap";
 import { HTTPRequest } from "./HTTPRequest";
+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;
-    }
-
-    calls[hostname]++;
-
-    if (
-      hostname === "thisdomainpointstointernalip.com" ||
-      hostname === "thisdomainpointstointernalip2.com"
-    ) {
-      return original.apply(this, [
-        "localhost",
-        ...Array.from(arguments).slice(1),
-      ]);
-    }
-
-    original.apply(this, arguments);
-  };
-});
+if (isCJS()) {
+  wrap(dns, "lookup", function lookup(original) {
+    return function lookup() {
+      const hostname = arguments[0];
+
+      if (!calls[hostname]) {
+        calls[hostname] = 0;
+      }
+
+      calls[hostname]++;
+
+      if (
+        hostname === "thisdomainpointstointernalip.com" ||
+        hostname === "thisdomainpointstointernalip2.com"
+      ) {
+        return original.apply(this, [
+          "localhost",
+          ...Array.from(arguments).slice(1),
+        ]);
+      }
+
+      original.apply(this, arguments);
+    };
+  });
+}
 
 const context: Context = {
   remoteAddress: "::1",
@@ -49,21 +53,30 @@ const context: Context = {
   route: "/posts/:id",
 };
 
-t.test("it works", (t) => {
-  const agent = new Agent(
+let http: typeof import("http");
+let https: typeof import("https");
+let agent: Agent;
+
+// Separate test for the setup is required because the setup is async
+// The second test can not be async because we can not defer the t.end() call in the end
+t.test("setup", async (t) => {
+  agent = new Agent(
     true,
     new LoggerNoop(),
     new ReportingAPIForTesting(),
     new Token("123"),
-    undefined
+    undefined,
+    !isCJS()
   );
   agent.start([new HTTPRequest()]);
 
   t.same(agent.getHostnames().asArray(), []);
 
-  const http = require("http");
-  const https = require("https");
+  http = isCJS() ? require("http") : (await import("http")).default;
+  https = isCJS() ? require("https") : (await import("https")).default;
+});
 
+t.test("it works", (t) => {
   runWithContext(context, () => {
     const aikido = http.request("http://aikido.dev");
     aikido.end();
@@ -133,47 +146,55 @@ t.test("it works", (t) => {
   t.same(agent.getHostnames().asArray(), []);
   agent.getHostnames().clear();
 
-  runWithContext(
-    { ...context, ...{ body: { image: "thisdomainpointstointernalip.com" } } },
-    () => {
-      https
-        .request("https://thisdomainpointstointernalip.com")
-        .on("error", (error) => {
-          t.match(
-            error.message,
-            "Zen has blocked a server-side request forgery: https.request(...) originating from body.image"
-          );
-
-          // Ensure the lookup is only called once per hostname
-          // Otherwise, it could be vulnerable to TOCTOU
-          t.same(calls["thisdomainpointstointernalip.com"], 1);
-        })
-        .on("finish", () => {
-          t.fail("should not finish");
-        })
-        .end();
-    }
-  );
+  if (isCJS()) {
+    runWithContext(
+      {
+        ...context,
+        ...{ body: { image: "thisdomainpointstointernalip.com" } },
+      },
+      () => {
+        https
+          .request("https://thisdomainpointstointernalip.com")
+          .on("error", (error) => {
+            t.match(
+              error.message,
+              "Zen has blocked a server-side request forgery: https.request(...) originating from body.image"
+            );
+
+            // Ensure the lookup is only called once per hostname
+            // Otherwise, it could be vulnerable to TOCTOU
+            t.same(calls["thisdomainpointstointernalip.com"], 1);
+          })
+          .on("finish", () => {
+            t.fail("should not finish");
+          })
+          .end();
+      }
+    );
 
-  runWithContext(
-    { ...context, ...{ body: { image: "thisdomainpointstointernalip2.com" } } },
-    () => {
-      https
-        .request("https://thisdomainpointstointernalip2.com", (res) => {
-          t.fail("should not respond");
-        })
-        .on("error", (error) => {
-          t.match(
-            error.message,
-            "Zen has blocked a server-side request forgery: https.request(...) originating from body.image"
-          );
-        })
-        .on("finish", () => {
-          t.fail("should not finish");
-        })
-        .end();
-    }
-  );
+    runWithContext(
+      {
+        ...context,
+        ...{ body: { image: "thisdomainpointstointernalip2.com" } },
+      },
+      () => {
+        https
+          .request("https://thisdomainpointstointernalip2.com", (res) => {
+            t.fail("should not respond");
+          })
+          .on("error", (error) => {
+            t.match(
+              error.message,
+              "Zen has blocked a server-side request forgery: https.request(...) originating from body.image"
+            );
+          })
+          .on("finish", () => {
+            t.fail("should not finish");
+          })
+          .end();
+      }
+    );
+  }
 
   runWithContext(context, () => {
     // With lookup function specified

diff --git a/library/sinks/HTTPRequest.ts b/library/sinks/HTTPRequest.ts
@@ -164,16 +164,22 @@ export class HTTPRequest implements Wrapper {
 
     for (const module of modules) {
       hooks.addBuiltinModule(module).onRequire((exports, pkgInfo) => {
-        for (const method of methods) {
-          wrapExport(exports, method, pkgInfo, {
-            // Whenever a request is made, we'll check the hostname whether it's a private IP
-            inspectArgs: (args, agent) =>
-              this.inspectHttpRequest(args, agent, module),
-            // Whenever a request is made, we'll modify the options to pass a custom lookup function
-            // that will inspect resolved IP address (and thus preventing TOCTOU attacks)
-            modifyArgs: (args, agent) =>
-              this.monitorDNSLookups(args, agent, module),
-          });
+        const toWrap = pkgInfo.isESMImport
+          ? [exports, exports.default]
+          : [exports];
+
+        for (const exportObj of toWrap) {
+          for (const method of methods) {
+            wrapExport(exportObj, method, pkgInfo, {
+              // Whenever a request is made, we'll check the hostname whether it's a private IP
+              inspectArgs: (args, agent) =>
+                this.inspectHttpRequest(args, agent, module),
+              // Whenever a request is made, we'll modify the options to pass a custom lookup function
+              // that will inspect resolved IP address (and thus preventing TOCTOU attacks)
+              modifyArgs: (args, agent) =>
+                this.monitorDNSLookups(args, agent, module),
+            });
+          }
         }
       });
     }

diff --git a/library/sinks/MongoDB.test.ts b/library/sinks/MongoDB.test.ts
@@ -4,6 +4,7 @@ import { ReportingAPIForTesting } from "../agent/api/ReportingAPIForTesting";
 import { Context, runWithContext } from "../agent/Context";
 import { LoggerNoop } from "../agent/logger/LoggerNoop";
 import { MongoDB } from "./MongoDB";
+import { isCJS } from "../helpers/isCJS";
 
 const unsafeContext: Context = {
   remoteAddress: "::1",
@@ -41,11 +42,14 @@ t.test("it inspects method calls and blocks if needed", async (t) => {
     new LoggerNoop(),
     new ReportingAPIForTesting(),
     undefined,
-    "lambda"
+    "lambda",
+    !isCJS()
   );
   agent.start([new MongoDB()]);
 
-  const { MongoClient } = require("mongodb");
+  const { MongoClient } = isCJS()
+    ? require("mongodb")
+    : await import("mongodb");
   const client = new MongoClient("mongodb://root:[email protected]:27017");
   await client.connect();