Integrate esm logic into agent

Makefile

@@ -72,6 +72,7 @@ install:
 
 .PHONY: build
 build:
+	rm -r library/build
 	cd library && npm run build
 
 .PHONY: watch

benchmarks/api-discovery/benchmark.js

@@ -1,8 +1,10 @@
 /**
  * Runs benchmarks for the api discovery (api schema collection)
  */
-const { Routes } = require("../../build/agent/Routes");
-const { isFeatureEnabled } = require("../../build/helpers/featureFlags");
+const { Routes } = require("../../library/build/agent/Routes");
+const {
+  isFeatureEnabled,
+} = require("../../library/build/helpers/featureFlags");
 const reqBodies = require("./reqBodies");
 
 const MAX_TIME_LIMIT = 0.05; // milliseconds / statement

benchmarks/nosql-injection/withGuard.js

@@ -1,9 +1,9 @@
-require("../../build");
+require("../../library/build");
 
 const measure = require("./measure");
 const getUser = require("./getUser");
 const getClient = require("./getClient");
-const { runWithContext } = require("../../build/agent/Context");
+const { runWithContext } = require("../../library/build/agent/Context");
 
 async function main() {
   const client = await getClient();

benchmarks/shell-injection/benchmark.js

@@ -3,7 +3,7 @@
  */
 const {
   detectShellInjection,
-} = require("../../build/vulnerabilities/shell-injection/detectShellInjection");
+} = require("../../library/build/vulnerabilities/shell-injection/detectShellInjection");
 
 const MAX_TIME_LIMIT = 0.05; // milliseconds / statement
 

benchmarks/sql-injection/benchmark.js

@@ -5,10 +5,10 @@ const fs = require("fs");
 const { join } = require("path");
 const {
   detectSQLInjection,
-} = require("../../build/vulnerabilities/sql-injection/detectSQLInjection");
+} = require("../../library/build/vulnerabilities/sql-injection/detectSQLInjection");
 const {
   SQLDialectMySQL,
-} = require("../../build/vulnerabilities/sql-injection/dialects/SQLDialectMySQL");
+} = require("../../library/build/vulnerabilities/sql-injection/dialects/SQLDialectMySQL");
 
 const MAX_TIME_LIMIT = 0.05; // milliseconds / statement
 

library/agent/Agent.ts

@@ -53,7 +53,8 @@ export class Agent {
     private readonly logger: Logger,
     private readonly api: ReportingAPI,
     private readonly token: Token | undefined,
-    private readonly serverless: string | undefined
+    private readonly serverless: string | undefined,
+    private readonly isESM: boolean | undefined = false
   ) {
     if (typeof this.serverless === "string" && this.serverless.length === 0) {
       throw new Error("Serverless cannot be an empty string");
@@ -415,7 +416,7 @@ export class Agent {
       setInstance(this);
     }
 
-    wrapInstalledPackages(wrappers);
+    wrapInstalledPackages(wrappers, this.isESM || false);
 
     // Send startup event and wait for config
     // Then start heartbeats and polling for config changes

library/agent/applyHooks.ts

@@ -5,6 +5,11 @@
   wrapRequire,
 } from "./hooks/wrapRequire";
 import { wrapExport } from "./hooks/wrapExport";
+import {
+  setImportBuiltinModulesToPatch,
+  setImportPackagesToPatch,
+  wrapImport,
+} from "./hooks/wrapImport";
 
 /**
  * Hooks allows you to register packages and then wrap specific methods on
@@ -13,10 +18,17 @@
  * This method wraps the require function and sets up the hooks.
  * Globals are wrapped directly.
  */
-export function applyHooks(hooks: Hooks) {
-  setPackagesToPatch(hooks.getPackages());
-  setBuiltinModulesToPatch(hooks.getBuiltInModules());
-  wrapRequire();
+export function applyHooks(hooks: Hooks, isESM: boolean) {
+  if (!isESM) {
+    // Todo check if we need to wrap require too in ESM mode under certain conditions
+    setPackagesToPatch(hooks.getPackages());
+    setBuiltinModulesToPatch(hooks.getBuiltInModules());
+    wrapRequire();
+  } else {
+    setImportPackagesToPatch(hooks.getPackages());
+    setImportBuiltinModulesToPatch(hooks.getBuiltInModules());
+    wrapImport();
+  }
 
   hooks.getGlobals().forEach((g) => {
     const name = g.getName();

library/agent/hooks/isBuiltinModule.ts

@@ -6,7 +6,7 @@ const moduleList = mod.builtinModules;
 /**
  * Returns true if the module is a builtin module, otherwise false.
  */
-export function isBuiltinModule(moduleName: string) {
+export function isBuiltinModule(moduleName: string): boolean {
   // Added in Node.js v18.6.0, v16.17.0
   if (typeof mod.isBuiltin === "function") {
     return mod.isBuiltin(moduleName);

library/agent/hooks/wrapImport.ts

@@ -0,0 +1,66 @@
+/* eslint-disable max-lines-per-function */
+import { register } from "module";
+import { pathToFileURL } from "url";
+import { Hook } from "import-in-the-middle";
+import { BuiltinModule } from "./BuiltinModule";
+import { Package } from "./Package";
+
+let isImportHookRegistered = false;
+
+let packages: Package[] = [];
+let builtinModules: BuiltinModule[] = [];
+let pkgCache = new Map<string, unknown>();
+let builtinCache = new Map<string, unknown>();
+
+/**
+ * Intercept esm package imports.
+ * This function makes sure that the import function is only wrapped once.
+ */
+export function wrapImport() {
+  if (isImportHookRegistered) {
+    return;
+  }
+
+  // Prevent registering the import hook multiple times
+  isImportHookRegistered = true;
+
+  // Register import-in-the-middle hook
+  register("import-in-the-middle/hook.mjs", pathToFileURL(__filename));
+
+  const allPackageNames = [packages, builtinModules]
+    .flat()
+    .map((p) => p.getName());
+
+  new Hook(
+    allPackageNames,
+    {
+      internals: true,
+    },
+    (exports, name, baseDir) => {
+      // Todo ...
+      //console.log(name);
+      //console.log(baseDir);
+      //console.log(process.pid);
+    }
+  );
+}
+
+/**
+ * Update the list of external packages that should be patched.
+ */
+export function setImportPackagesToPatch(packagesToPatch: Package[]) {
+  packages = packagesToPatch;
+  // Reset cache
+  pkgCache = new Map();
+}
+
+/**
+ * Update the list of builtin modules that should be patched.
+ */
+export function setImportBuiltinModulesToPatch(
+  builtinModulesToPatch: BuiltinModule[]
+) {
+  builtinModules = builtinModulesToPatch;
+  // Reset cache
+  builtinCache = new Map();
+}

library/agent/protect.ts

@@ -93,7 +93,13 @@
     : undefined;
 }
 
-function getAgent({ serverless }: { serverless: string | undefined }) {
+function getAgent({
+  serverless,
+  isESM,
+}: {
+  serverless: string | undefined;
+  isESM?: boolean;
+}) {
   const current = getInstance();
 
   if (current) {
@@ -105,7 +111,8 @@
     getLogger(),
     getAPI(),
     getTokenFromEnv(),
-    serverless
+    serverless,
+    isESM
   );
 
   setInstance(agent);
@@ -142,9 +149,10 @@
   ];
 }
 
-export function protect() {
+export function protect(isESM = false) {
   const agent = getAgent({
     serverless: undefined,
+    isESM,
   });
 
   agent.start(getWrappers());

library/agent/wrapInstalledPackages.ts

@@ -2,11 +2,11 @@ import { applyHooks } from "./applyHooks";
 import { Hooks } from "./hooks/Hooks";
 import { Wrapper } from "./Wrapper";
 
-export function wrapInstalledPackages(wrappers: Wrapper[]) {
+export function wrapInstalledPackages(wrappers: Wrapper[], isESM: boolean) {
   const hooks = new Hooks();
   wrappers.forEach((wrapper) => {
     wrapper.wrap(hooks);
   });
 
-  return applyHooks(hooks);
+  return applyHooks(hooks, isESM);
 }

library/esm/index.ts

@@ -1,17 +1,26 @@
-import { Hook, createAddHookMessageChannel } from "import-in-the-middle";
-import { register } from "module";
-import { pathToFileURL } from "url";
+import isFirewallSupported from "../helpers/isFirewallSupported";
+import {
+  getMajorNodeVersion,
+  getMinorNodeVersion,
+} from "../helpers/getNodeVersion";
 
-// Todo init agent?
+// Was added in v20.6.0 and v18.19.0
+function isESMSupported() {
+  const nodeMajor = getMajorNodeVersion();
+  const nodeMinor = getMinorNodeVersion();
+  return (
+    nodeMajor >= 22 ||
+    (nodeMajor === 20 && nodeMinor >= 6) ||
+    (nodeMajor === 18 && nodeMinor >= 19)
+  );
+}
 
-const { registerOptions, waitForAllMessagesAcknowledged } =
-  createAddHookMessageChannel();
-
-register(
-  "import-in-the-middle/hook.mjs",
-  pathToFileURL(__filename),
-  // @ts-ignore test
-  registerOptions
-);
-
-// Todo
+if (isFirewallSupported()) {
+  if (isESMSupported()) {
+    require("../agent/protect").protect(true);
+  } else {
+    console.error(
+      "Error: Aikido Firewall requires Node.js v20.6.0 / v18.19.0 or higher to support ESM."
+    );
+  }
+}

library/helpers/getAgentVersion.ts

@@ -1,7 +1,7 @@
 import { resolve } from "path";
 
 export function getAgentVersion(): string {
-  const json = require(resolve(__dirname, "../package.json"));
+  const json = require(resolve(__dirname, "../../package.json"));
 
   /* c8 ignore start */
   if (!json.version) {