esbuild plugin (POC)

library/agent/Agent.ts

@@ -71,7 +71,8 @@ export class Agent {
     private readonly api: ReportingAPI,
     private readonly token: Token | undefined,
     private readonly serverless: string | undefined,
-    private readonly newInstrumentation: boolean = false
+    private readonly newInstrumentation: boolean = false,
+    private readonly isBundlingProcess: boolean = false
   ) {
     if (typeof this.serverless === "string" && this.serverless.length === 0) {
       throw new Error("Serverless cannot be an empty string");
@@ -480,27 +481,33 @@ export class Agent {
 
     this.logger.log(`Starting agent v${getAgentVersion()}...`);
 
-    if (!this.block) {
-      this.logger.log("Dry mode enabled, no requests will be blocked!");
-    }
+    if (!this.isBundlingProcess) {
+      if (!this.block) {
+        this.logger.log("Dry mode enabled, no requests will be blocked!");
+      }
 
-    if (this.token) {
-      this.logger.log("Found token, reporting enabled!");
-    } else {
-      this.logger.log("No token provided, disabling reporting.");
+      if (this.token) {
+        this.logger.log("Found token, reporting enabled!");
+      } else {
+        this.logger.log("No token provided, disabling reporting.");
 
-      if (!this.block && !isAikidoCI()) {
-        console.log(
-          "AIKIDO: Running in monitoring only mode without reporting to Aikido Cloud. Set AIKIDO_BLOCK=true to enable blocking."
-        );
+        if (!this.block && !isAikidoCI()) {
+          console.log(
+            "AIKIDO: Running in monitoring only mode without reporting to Aikido Cloud. Set AIKIDO_BLOCK=true to enable blocking."
+          );
+        }
       }
     }
 
     // When our library is required, we are not intercepting `require` calls yet
     // We need to add our library to the list of packages manually
     this.onPackageRequired("@aikidosec/firewall", getAgentVersion());
 
-    wrapInstalledPackages(wrappers, this.newInstrumentation, this.serverless);
+    wrapInstalledPackages(wrappers, this.newInstrumentation, !this.serverless);
+
+    if (this.isBundlingProcess) {
+      return;
+    }
 
     // Send startup event and wait for config
     // Then start heartbeats and polling for config changes

library/agent/hooks/instrumentation/loadHook.ts

@@ -56,7 +56,7 @@ export function onModuleLoad(
   }
 }
 
-function patchPackage(
+export function patchPackage(
   path: string,
   previousLoadResult: ReturnType<LoadFunction>
 ) {

library/agent/protect.ts

@@ -96,9 +96,11 @@ function getTokenFromEnv(): Token | undefined {
 function startAgent({
   serverless,
   newInstrumentation,
+  isBundlingProcess,
 }: {
   serverless: string | undefined;
   newInstrumentation: boolean;
+  isBundlingProcess: boolean;
 }) {
   const current = getInstance();
 
@@ -112,7 +114,8 @@ function startAgent({
     getAPI(),
     getTokenFromEnv(),
     serverless,
-    newInstrumentation
+    newInstrumentation,
+    isBundlingProcess
   );
 
   setInstance(agent);
@@ -169,13 +172,15 @@ export function protect() {
   startAgent({
     serverless: undefined,
     newInstrumentation: false,
+    isBundlingProcess: false,
   });
 }
 
 export function lambda(): (handler: Handler) => Handler {
   startAgent({
     serverless: "lambda",
     newInstrumentation: false,
+    isBundlingProcess: false,
   });
 
   return createLambdaWrapper;
@@ -185,6 +190,7 @@ export function cloudFunction(): (handler: HttpFunction) => HttpFunction {
   startAgent({
     serverless: "gcp",
     newInstrumentation: false,
+    isBundlingProcess: false,
   });
 
   return createCloudFunctionWrapper;
@@ -194,5 +200,14 @@ export function protectWithNewInstrumentation() {
   startAgent({
     serverless: undefined,
     newInstrumentation: true,
+    isBundlingProcess: false,
+  });
+}
+
+export function protectDuringBundling() {
+  startAgent({
+    serverless: undefined,
+    newInstrumentation: true,
+    isBundlingProcess: true,
   });
 }

library/agent/wrapInstalledPackages.ts

@@ -6,14 +6,17 @@ import { Wrapper } from "./Wrapper";
 export function wrapInstalledPackages(
   wrappers: Wrapper[],
   newInstrumentation: boolean,
-  serverless: string | undefined
+  printWarningForAlreadyRequiredPkgs: boolean
 ) {
   const hooks = new Hooks();
   wrappers.forEach((wrapper) => {
     wrapper.wrap(hooks);
   });
 
-  if (!serverless && isAnyPkgAlreadyRequired(hooks.getPackages())) {
+  if (
+    printWarningForAlreadyRequiredPkgs &&
+    isAnyPkgAlreadyRequired(hooks.getPackages())
+  ) {
     // eslint-disable-next-line no-console
     console.warn(
       "AIKIDO: Some packages can't be protected because they were imported before Zen was initialized. Please make sure to import Zen as the first module in your application."

library/bundler/esbuild.ts

@@ -0,0 +1,87 @@
+import type { OnLoadArgs, OnLoadResult, Plugin, PluginBuild } from "esbuild";
+import { patchPackage } from "../agent/hooks/instrumentation/loadHook";
+import { copyFile, readFile } from "fs/promises";
+import { protectDuringBundling } from "../agent/protect";
+import { basename, dirname, join } from "path";
+
+export const zenEsbuildPlugin: Plugin = {
+  name: "@aikidosec/firewall",
+  async setup(build) {
+    const isBundlingPackages =
+      build.initialOptions.bundle &&
+      build.initialOptions.packages !== "external";
+
+    if (!isBundlingPackages) {
+      return;
+    }
+    const outDir = getOutDir(build);
+
+    await copyWasmFiles(outDir);
+
+    protectDuringBundling();
+
+    build.onLoad({ filter: /\.(js|mjs|cjs|jsx)$/, namespace: "file" }, onLoad);
+  },
+};
+
+async function onLoad(args: OnLoadArgs): Promise<OnLoadResult> {
+  try {
+    const fileContent = await readFile(args.path, "utf-8");
+    const result = patchPackage(args.path, {
+      source: fileContent,
+      format: "unambiguous",
+      shortCircuit: false,
+    });
+
+    if (typeof result.source !== "string") {
+      result.source = new TextDecoder("utf-8").decode(result.source);
+    }
+
+    return {
+      contents: result.source,
+    };
+  } catch (error) {
+    return {
+      errors: [
+        {
+          text: `Failed to patch package at ${args.path}: ${String(error)}`,
+        },
+      ],
+    };
+  }
+}
+
+function getOutDir(build: PluginBuild): string {
+  const outdir = build.initialOptions.outdir;
+  const outfile = build.initialOptions.outfile;
+
+  // Pick destination dir (support both outfile & outdir)
+  const outDirResolved = outdir ?? (outfile ? dirname(outfile) : null);
+  if (!outDirResolved) {
+    throw new Error(
+      "Zen esbuild plugin requires either 'outdir' or 'outfile' to be set in the build options."
+    );
+  }
+
+  return outDirResolved;
+}
+
+async function copyWasmFiles(outDir: string): Promise<void> {
+  const zenLibDir = dirname(require.resolve("@aikidosec/firewall"));
+
+  const wasmFiles = [
+    join(zenLibDir, "internals", "zen_internals_bg.wasm"),
+    join(
+      zenLibDir,
+      "agent",
+      "hooks",
+      "instrumentation",
+      "wasm",
+      "node_code_instrumentation_bg.wasm"
+    ),
+  ];
+
+  await Promise.all(
+    wasmFiles.map((file) => copyFile(file, join(outDir, basename(file))))
+  );
+}