Add prepare() and setToken() for runtime token setup

Users who store their token in a secrets manager (e.g. AWS
Secrets Manager) can't set AIKIDO_TOKEN before the module
loads. prepare() starts instrumentation without a token,
and setToken() connects to the platform once the token is
fetched async. Teams shouldn't have to modify their whole
app structure just to adopt Zen.

Author: hansott

docs/set-token.md

@@ -0,0 +1,89 @@
+# Setting the token at runtime
+
+Zen normally reads the token from the `AIKIDO_TOKEN` environment variable. If you can't set env vars — for example, your token lives in AWS Secrets Manager — you can set it at runtime instead.
+
+## How it works
+
+1. Call `prepare()` at startup. This starts Zen's instrumentation without a token.
+2. Fetch your token async (secrets manager, config service, wherever).
+3. Call `setToken(token)` to connect to the Aikido platform.
+
+Zen detects attacks from step 1, but won't report them until you call `setToken`.
+
+## Example with AWS Secrets Manager
+
+```js
+const Zen = require("@aikidosec/firewall");
+
+// Start instrumentation without a token
+Zen.prepare();
+
+const {
+  SecretsManagerClient,
+  GetSecretValueCommand,
+} = require("@aws-sdk/client-secrets-manager");
+
+async function loadToken() {
+  const client = new SecretsManagerClient();
+  const response = await client.send(
+    new GetSecretValueCommand({ SecretId: "my-secret" })
+  );
+  return response.SecretString;
+}
+
+loadToken().then((token) => {
+  Zen.setToken(token);
+});
+```
+
+## With ESM
+
+Create a setup file for ESM:
+
+```js
+// zen-setup.cjs
+const { prepare } = require("@aikidosec/firewall/instrument");
+
+prepare();
+```
+
+Start your app with:
+
+```sh
+node -r ./zen-setup.cjs app.js
+```
+
+Then call `setToken` in your application code:
+
+```js
+import { setToken } from "@aikidosec/firewall";
+
+const token = await fetchTokenFromSecretsManager();
+setToken(token);
+```
+
+## With Lambda
+
+Call `prepare()` before wrapping your handler:
+
+```js
+const Zen = require("@aikidosec/firewall");
+Zen.prepare();
+
+const zen = require("@aikidosec/firewall/lambda");
+
+module.exports.handler = zen(async (event) => {
+  // Your handler code
+});
+
+// Fetch token outside the handler so it runs once during cold start
+loadToken().then((token) => {
+  Zen.setToken(token);
+});
+```
+
+## Notes
+
+- Call `prepare()` as early as possible, before other packages are loaded.
+- `setToken` only works once. Calling it again is ignored.
+- If `AIKIDO_TOKEN` is already set in the environment, you don't need `prepare()` or `setToken()`. Calling them anyway is fine — they just do nothing.

end2end/tests-new/hono-pg-esm-set-token.test.mjs

@@ -0,0 +1,110 @@
+import { spawn } from "child_process";
+import { resolve } from "path";
+import { test } from "node:test";
+import { equal, fail, ok } from "node:assert";
+import { getRandomPort } from "./utils/get-port.mjs";
+import { timeout } from "./utils/timeout.mjs";
+
+const pathToAppDir = resolve(
+  import.meta.dirname,
+  "../../sample-apps/hono-pg-esm"
+);
+
+const testServerUrl = "http://localhost:5874";
+
+test(
+  "it blocks after setToken is called and sends a heartbeat (ESM)",
+  { timeout: 60000 },
+  async () => {
+    const port = await getRandomPort();
+
+    const response = await fetch(`${testServerUrl}/api/runtime/apps`, {
+      method: "POST",
+    });
+    const body = await response.json();
+    const token = body.token;
+
+    const server = spawn(
+      `node`,
+      ["--require", "./zen-setup.cjs", "./app-set-token.js", port],
+      {
+        cwd: pathToAppDir,
+        env: {
+          ...process.env,
+          AIKIDO_TOKEN: token,
+          AIKIDO_ENDPOINT: testServerUrl,
+          AIKIDO_REALTIME_ENDPOINT: testServerUrl,
+          AIKIDO_DEBUG: "true",
+          AIKIDO_BLOCK: "true",
+        },
+      }
+    );
+
+    try {
+      server.on("error", (err) => {
+        fail(err);
+      });
+
+      let stdout = "";
+      server.stdout.on("data", (data) => {
+        stdout += data.toString();
+      });
+
+      let stderr = "";
+      server.stderr.on("data", (data) => {
+        stderr += data.toString();
+      });
+
+      // Wait for server + setToken (500ms delay in app)
+      await timeout(2000);
+
+      const [sqlInjection, normalAdd] = await Promise.all([
+        fetch(`http://127.0.0.1:${port}/add`, {
+          method: "POST",
+          body: JSON.stringify({
+            name: "Njuska'); DELETE FROM cats_6;-- H",
+          }),
+          headers: { "Content-Type": "application/json" },
+          signal: AbortSignal.timeout(5000),
+        }),
+        fetch(`http://127.0.0.1:${port}/add`, {
+          method: "POST",
+          body: JSON.stringify({ name: "Miau" }),
+          headers: { "Content-Type": "application/json" },
+          signal: AbortSignal.timeout(5000),
+        }),
+      ]);
+
+      equal(sqlInjection.status, 500);
+      equal(normalAdd.status, 200);
+      ok(stdout.includes("Starting agent"), "should log starting agent");
+      ok(
+        stderr.includes("Zen has blocked an SQL injection"),
+        "should log blocked SQL injection"
+      );
+
+      // Wait for heartbeat (agent sends after ~30s)
+      await timeout(31000);
+
+      const eventsResponse = await fetch(
+        `${testServerUrl}/api/runtime/events`,
+        {
+          method: "GET",
+          headers: { Authorization: token },
+          signal: AbortSignal.timeout(5000),
+        }
+      );
+
+      const events = await eventsResponse.json();
+      const startedEvents = events.filter((e) => e.type === "started");
+      equal(startedEvents.length, 1, "should have 1 started event");
+
+      const heartbeatEvents = events.filter((e) => e.type === "heartbeat");
+      equal(heartbeatEvents.length, 1, "should have 1 heartbeat event");
+    } catch (err) {
+      fail(err);
+    } finally {
+      server.kill();
+    }
+  }
+);

end2end/tests/express-mysql.set-token.test.js

@@ -0,0 +1,105 @@
+const t = require("tap");
+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const timeout = require("../timeout");
+
+const pathToApp = resolve(
+  __dirname,
+  "../../sample-apps/express-mysql",
+  "app-set-token.js"
+);
+
+const testServerUrl = "http://localhost:5874";
+
+let token;
+t.beforeEach(async () => {
+  const response = await fetch(`${testServerUrl}/api/runtime/apps`, {
+    method: "POST",
+  });
+  const body = await response.json();
+  token = body.token;
+});
+
+t.test(
+  "it blocks after setToken is called and sends a heartbeat",
+  { timeout: 60000 },
+  (t) => {
+    const server = spawn(`node`, [pathToApp, "4020"], {
+      env: {
+        ...process.env,
+        AIKIDO_TOKEN: token,
+        AIKIDO_ENDPOINT: testServerUrl,
+        AIKIDO_REALTIME_ENDPOINT: testServerUrl,
+        AIKIDO_DEBUG: "true",
+        AIKIDO_BLOCKING: "true",
+      },
+    });
+
+    server.on("close", () => {
+      t.end();
+    });
+
+    server.on("error", (err) => {
+      t.fail(err.message);
+    });
+
+    let stdout = "";
+    server.stdout.on("data", (data) => {
+      stdout += data.toString();
+    });
+
+    let stderr = "";
+    server.stderr.on("data", (data) => {
+      stderr += data.toString();
+    });
+
+    // Wait for server + setToken (500ms delay in app)
+    timeout(2000)
+      .then(() => {
+        return Promise.all([
+          fetch(
+            `http://localhost:4020/?petname=${encodeURIComponent("Njuska'); DELETE FROM cats;-- H")}`,
+            {
+              signal: AbortSignal.timeout(5000),
+            }
+          ),
+          fetch("http://localhost:4020/?petname=Njuska", {
+            signal: AbortSignal.timeout(5000),
+          }),
+        ]);
+      })
+      .then(([sqlInjection, normalSearch]) => {
+        t.equal(sqlInjection.status, 500);
+        t.equal(normalSearch.status, 200);
+        t.match(stdout, /Starting agent/);
+        t.match(stderr, /Zen has blocked an SQL injection/);
+      })
+      .then(() => {
+        // Wait for heartbeat (agent sends after ~30s)
+        return timeout(31000);
+      })
+      .then(() => {
+        return fetch(`${testServerUrl}/api/runtime/events`, {
+          method: "GET",
+          headers: {
+            Authorization: token,
+          },
+          signal: AbortSignal.timeout(5000),
+        });
+      })
+      .then((response) => response.json())
+      .then((events) => {
+        const startedEvents = events.filter((e) => e.type === "started");
+        t.equal(startedEvents.length, 1);
+
+        const heartbeatEvents = events.filter((e) => e.type === "heartbeat");
+        t.equal(heartbeatEvents.length, 1);
+      })
+      .catch((error) => {
+        t.fail(error.message);
+      })
+      .finally(() => {
+        server.kill();
+      });
+  }
+);

library/agent/Agent.ts

@@ -72,7 +72,7 @@ export class Agent {
     private block: boolean,
     private readonly logger: Logger,
     private readonly api: ReportingAPI,
-    private readonly token: Token | undefined,
+    private token: Token | undefined,
     private readonly serverless: string | undefined,
     private readonly newInstrumentation: boolean = false,
     private readonly fetchListsAPI: FetchListsAPI
@@ -555,6 +555,24 @@ export class Agent {
       });
   }
 
+  hasToken() {
+    return this.token !== undefined;
+  }
+
+  setToken(token: Token) {
+    this.token = token;
+    this.logger.log("Token set, enabling reporting.");
+
+    this.onStart()
+      .then(() => {
+        this.startHeartbeats();
+        this.startPollingForConfigChanges();
+      })
+      .catch((err) => {
+        console.error(`Aikido: Failed to start agent: ${err.message}`);
+      });
+  }
+
   onFailedToWrapMethod(module: string, name: string, error: Error) {
     this.logger.log(
       `Failed to wrap method ${name} in module ${module}: ${error.message}`