From 72adc2853dc9390f6affbb7b48e8be4c724a81a7 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Tue, 28 Apr 2026 14:21:53 +0200 Subject: [PATCH 1/2] 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. --- docs/set-token.md | 89 ++++++++++++++ .../tests-new/hono-pg-esm-set-token.test.mjs | 110 ++++++++++++++++++ end2end/tests/express-mysql.set-token.test.js | 105 +++++++++++++++++ library/agent/Agent.ts | 20 +++- library/agent/protect.setToken.test.ts | 74 ++++++++++++ library/agent/protect.ts | 46 +++++++- library/index.ts | 5 + library/instrument/index.ts | 9 +- sample-apps/express-mysql/app-set-token.js | 74 ++++++++++++ sample-apps/hono-pg-esm/app-set-token.js | 50 ++++++++ sample-apps/hono-pg-esm/zen-setup.cjs | 3 + 11 files changed, 581 insertions(+), 4 deletions(-) create mode 100644 docs/set-token.md create mode 100644 end2end/tests-new/hono-pg-esm-set-token.test.mjs create mode 100644 end2end/tests/express-mysql.set-token.test.js create mode 100644 library/agent/protect.setToken.test.ts create mode 100644 sample-apps/express-mysql/app-set-token.js create mode 100644 sample-apps/hono-pg-esm/app-set-token.js create mode 100644 sample-apps/hono-pg-esm/zen-setup.cjs diff --git a/docs/set-token.md b/docs/set-token.md new file mode 100644 index 000000000..b73ce7885 --- /dev/null +++ b/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. diff --git a/end2end/tests-new/hono-pg-esm-set-token.test.mjs b/end2end/tests-new/hono-pg-esm-set-token.test.mjs new file mode 100644 index 000000000..3532e982a --- /dev/null +++ b/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(); + } + } +); diff --git a/end2end/tests/express-mysql.set-token.test.js b/end2end/tests/express-mysql.set-token.test.js new file mode 100644 index 000000000..893ba0037 --- /dev/null +++ b/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(); + }); + } +); diff --git a/library/agent/Agent.ts b/library/agent/Agent.ts index dded6ec43..586831fd3 100644 --- a/library/agent/Agent.ts +++ b/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}` diff --git a/library/agent/protect.setToken.test.ts b/library/agent/protect.setToken.test.ts new file mode 100644 index 000000000..722ea783a --- /dev/null +++ b/library/agent/protect.setToken.test.ts @@ -0,0 +1,74 @@ +/* oxlint-disable no-console */ +import * as t from "tap"; +import { ReportingAPIForTesting } from "./api/ReportingAPIForTesting"; +import { Token } from "./api/Token"; +import { getInstance, setInstance } from "./AgentSingleton"; +import { createTestAgent } from "../helpers/createTestAgent"; +import { setToken } from "./protect"; + +t.beforeEach(() => { + // @ts-expect-error Reset singleton for isolation + setInstance(undefined); +}); + +t.test("it sets the token on an agent without one", async (t) => { + const api = new ReportingAPIForTesting(); + const agent = createTestAgent({ + api, + serverless: "lambda", + }); + agent.start([]); + + t.equal(agent.hasToken(), false); + + setToken("test-token-123"); + + t.equal(agent.hasToken(), true); + + // Give the async onStart a tick to complete + await new Promise((resolve) => setTimeout(resolve, 10)); + + const events = api.getEvents(); + const startedEvents = events.filter((e) => e.type === "started"); + t.equal(startedEvents.length, 1); +}); + +t.test("it ignores setToken if agent already has a token", async (t) => { + const api = new ReportingAPIForTesting(); + createTestAgent({ + api, + token: new Token("existing-token"), + serverless: "lambda", + }); + + setToken("new-token"); + + const agent = getInstance()!; + t.equal(agent.hasToken(), true); +}); + +t.test("it warns if no agent is running", async (t) => { + const warnings: string[] = []; + const originalWarn = console.warn; + console.warn = (msg: string) => warnings.push(msg); + + setToken("test-token"); + + console.warn = originalWarn; + + t.equal(warnings.length, 1); + t.match(warnings[0], /agent is not running/); +}); + +t.test("it warns on empty string", async (t) => { + const warnings: string[] = []; + const originalWarn = console.warn; + console.warn = (msg: string) => warnings.push(msg); + + setToken(""); + + console.warn = originalWarn; + + t.equal(warnings.length, 1); + t.match(warnings[0], /empty string/); +}); diff --git a/library/agent/protect.ts b/library/agent/protect.ts index 692f6a16d..fc8f5d3df 100644 --- a/library/agent/protect.ts +++ b/library/agent/protect.ts @@ -186,7 +186,7 @@ export function protect() { } export function lambda(): (handler: Handler) => Handler { - if (!shouldEnableFirewall()) { + if (!getInstance() && !shouldEnableFirewall()) { return (handler: Handler) => handler; } @@ -199,7 +199,7 @@ export function lambda(): (handler: Handler) => Handler { } export function cloudFunction(): (handler: HttpFunction) => HttpFunction { - if (!shouldEnableFirewall()) { + if (!getInstance() && !shouldEnableFirewall()) { return (handler: HttpFunction) => handler; } @@ -217,3 +217,45 @@ export function protectWithNewInstrumentation() { newInstrumentation: true, }); } + +export function prepare() { + startAgent({ + serverless: undefined, + newInstrumentation: false, + }); +} + +export function prepareWithNewInstrumentation() { + startAgent({ + serverless: undefined, + newInstrumentation: true, + }); +} + +export function setToken(token: string) { + if (token.length === 0) { + // oxlint-disable-next-line no-console + console.warn("AIKIDO: setToken called with an empty string, ignoring."); + return; + } + + const agent = getInstance(); + + if (!agent) { + // oxlint-disable-next-line no-console + console.warn( + "AIKIDO: setToken called but the agent is not running. Call prepare() first." + ); + return; + } + + if (agent.hasToken()) { + // oxlint-disable-next-line no-console + console.warn( + "AIKIDO: setToken called but the agent already has a token, ignoring." + ); + return; + } + + agent.setToken(new Token(token)); +} diff --git a/library/index.ts b/library/index.ts index b1a75baba..474aa0216 100644 --- a/library/index.ts +++ b/library/index.ts @@ -20,6 +20,7 @@ import { withoutIdorProtection } from "./agent/context/withoutIdorProtection"; import { colorText } from "./helpers/colorText"; import { isPreloaded } from "./helpers/isPreloaded"; import { warnIfEntrypointIsModule } from "./helpers/warnIfEntrypointIsModule"; +import { prepare, setToken } from "./agent/protect"; // Prevent logging twice / trying to start agent twice if (!isNewHookSystemUsed()) { @@ -71,6 +72,8 @@ export { setTenantId, enableIdorProtection, withoutIdorProtection, + prepare, + setToken, }; // Required for ESM / TypeScript default export support @@ -90,4 +93,6 @@ export default { setTenantId, enableIdorProtection, withoutIdorProtection, + prepare, + setToken, }; diff --git a/library/instrument/index.ts b/library/instrument/index.ts index 59b2708be..38dcf6bb0 100644 --- a/library/instrument/index.ts +++ b/library/instrument/index.ts @@ -10,11 +10,13 @@ import { isMainThread } from "node:worker_threads"; import { isESM } from "../helpers/isESM"; import { isPreloaded } from "../helpers/isPreloaded"; import { colorText } from "../helpers/colorText"; +import { getInstance } from "../agent/AgentSingleton"; setIsNewHookSystemUsed(true); const isSupported = isFirewallSupported(); -const shouldEnable = shouldEnableFirewall(); +const alreadyRunning = !!getInstance(); +const shouldEnable = alreadyRunning || shouldEnableFirewall(); const notAlreadyImported = checkIndexImportGuard(); function start() { @@ -52,3 +54,8 @@ function start() { } start(); + +export { + prepareWithNewInstrumentation as prepare, + setToken, +} from "../agent/protect"; diff --git a/sample-apps/express-mysql/app-set-token.js b/sample-apps/express-mysql/app-set-token.js new file mode 100644 index 000000000..ab892b6d3 --- /dev/null +++ b/sample-apps/express-mysql/app-set-token.js @@ -0,0 +1,74 @@ +const Zen = require("@aikidosec/firewall"); + +Zen.prepare(); + +const Cats = require("./Cats"); +const express = require("express"); +const asyncHandler = require("express-async-handler"); +const mysql = require("mysql"); + +async function createConnection() { + const connection = await mysql.createConnection({ + host: "localhost", + user: "root", + password: "mypassword", + database: "catsdb", + port: 27015, + multipleStatements: true, + }); + + await connection.query(` + CREATE TABLE IF NOT EXISTS cats ( + petname varchar(255) + ); + `); + + return connection; +} + +async function main(port) { + const db = await createConnection(); + const cats = new Cats(db); + + const app = express(); + + app.get( + "/", + asyncHandler(async (req, res) => { + if (req.query["petname"]) { + await cats.add(req.query["petname"]); + } + + res.send(`All cats: ${(await cats.getAll()).join(", ")}`); + }) + ); + + // Set the token after startup (simulates fetching from a secrets manager) + setTimeout(() => { + Zen.setToken(process.env.AIKIDO_TOKEN); + }, 500); + + return new Promise((resolve, reject) => { + try { + app.listen(port, () => { + console.log(`Listening on port ${port}`); + resolve(); + }); + } catch (err) { + reject(err); + } + }); +} + +function getPort() { + const port = parseInt(process.argv[2], 10) || 4000; + + if (isNaN(port)) { + console.error("Invalid port"); + process.exit(1); + } + + return port; +} + +main(getPort()); diff --git a/sample-apps/hono-pg-esm/app-set-token.js b/sample-apps/hono-pg-esm/app-set-token.js new file mode 100644 index 000000000..1b6dd5e7c --- /dev/null +++ b/sample-apps/hono-pg-esm/app-set-token.js @@ -0,0 +1,50 @@ +import { Hono } from "hono"; +import { serve } from "@hono/node-server"; +import { createConnection } from "./db.js"; +import { setToken } from "@aikidosec/firewall"; + +const app = new Hono(); +const db = await createConnection(); + +app.get("/", async (c) => { + return c.text("Hello, World!"); +}); + +app.post("/add", async (c) => { + const json = await c.req.json(); + const name = json.name; + if (!name) { + return c.status(400).text("Name is required"); + } + + await db.query( + `INSERT INTO cats_6 (petname, user_id) VALUES ('${name}', 1);` + ); + return c.text("OK"); +}); + +function getPort() { + const port = parseInt(process.argv[2], 10) || 4000; + + if (isNaN(port)) { + console.error("Invalid port"); + process.exit(1); + } + + return port; +} + +// Set the token after startup (simulates fetching from a secrets manager) +setTimeout(() => { + setToken(process.env.AIKIDO_TOKEN); +}, 500); + +serve( + { + fetch: app.fetch, + port: getPort(), + }, + (info) => { + console.log(`Server is running on http://${info.address}:${info.port}`); + } +); diff --git a/sample-apps/hono-pg-esm/zen-setup.cjs b/sample-apps/hono-pg-esm/zen-setup.cjs new file mode 100644 index 000000000..9d48e4f7b --- /dev/null +++ b/sample-apps/hono-pg-esm/zen-setup.cjs @@ -0,0 +1,3 @@ +const { prepare } = require("@aikidosec/firewall/instrument"); + +prepare(); From a0fc9235170d1f8b04187d2833b834aabe2cec84 Mon Sep 17 00:00:00 2001 From: Hans Ott Date: Thu, 30 Apr 2026 13:10:21 +0200 Subject: [PATCH 2/2] Refactor: Use AIKIDO_INSTRUMENT=true instead of prepare() --- .../tests-new/hono-pg-esm-set-token.test.mjs | 3 +- end2end/tests/express-mysql.set-token.test.js | 3 +- .../tests/functions-framework-sqlite3.test.js | 4 +- end2end/tests/lambda-mongodb.test.js | 4 +- library/agent/Agent.ts | 23 ++++--- library/agent/protect.setToken.test.ts | 68 +++++++++++++++++-- library/agent/protect.ts | 16 ++--- library/cloud-function/index.ts | 12 +++- library/cloud-function/instrument.ts | 9 +++ library/helpers/shouldEnableFirewall.test.ts | 9 +++ library/helpers/shouldEnableFirewall.ts | 8 ++- library/helpers/shouldInstrument.ts | 5 ++ library/index.ts | 4 +- library/instrument/index.ts | 63 ++--------------- library/instrument/start.ts | 53 +++++++++++++++ library/lambda/index.ts | 12 +++- library/lambda/instrument.ts | 6 ++ library/package.json | 2 + library/sources/FunctionsFramework.test.ts | 21 ++++++ library/sources/FunctionsFramework.ts | 2 +- library/sources/Lambda.test.ts | 24 +++++++ library/sources/Lambda.ts | 2 +- library/start/index.ts | 9 +++ sample-apps/express-mysql/app-set-token.js | 4 +- sample-apps/hono-pg-esm/app-set-token.js | 2 +- sample-apps/hono-pg-esm/zen-setup.cjs | 4 +- sample-apps/test-exports/test.js | 2 + sample-apps/test-exports/test.ts | 2 + 28 files changed, 270 insertions(+), 106 deletions(-) create mode 100644 library/cloud-function/instrument.ts create mode 100644 library/helpers/shouldInstrument.ts create mode 100644 library/instrument/start.ts create mode 100644 library/lambda/instrument.ts create mode 100644 library/start/index.ts diff --git a/end2end/tests-new/hono-pg-esm-set-token.test.mjs b/end2end/tests-new/hono-pg-esm-set-token.test.mjs index 3532e982a..83592263b 100644 --- a/end2end/tests-new/hono-pg-esm-set-token.test.mjs +++ b/end2end/tests-new/hono-pg-esm-set-token.test.mjs @@ -31,7 +31,8 @@ test( cwd: pathToAppDir, env: { ...process.env, - AIKIDO_TOKEN: token, + AIKIDO_INSTRUMENT: "true", + TEST_AIKIDO_TOKEN: token, AIKIDO_ENDPOINT: testServerUrl, AIKIDO_REALTIME_ENDPOINT: testServerUrl, AIKIDO_DEBUG: "true", diff --git a/end2end/tests/express-mysql.set-token.test.js b/end2end/tests/express-mysql.set-token.test.js index 893ba0037..69a87d2d9 100644 --- a/end2end/tests/express-mysql.set-token.test.js +++ b/end2end/tests/express-mysql.set-token.test.js @@ -27,7 +27,8 @@ t.test( const server = spawn(`node`, [pathToApp, "4020"], { env: { ...process.env, - AIKIDO_TOKEN: token, + AIKIDO_INSTRUMENT: "true", + TEST_AIKIDO_TOKEN: token, AIKIDO_ENDPOINT: testServerUrl, AIKIDO_REALTIME_ENDPOINT: testServerUrl, AIKIDO_DEBUG: "true", diff --git a/end2end/tests/functions-framework-sqlite3.test.js b/end2end/tests/functions-framework-sqlite3.test.js index 7a10669ba..e8ec435ec 100644 --- a/end2end/tests/functions-framework-sqlite3.test.js +++ b/end2end/tests/functions-framework-sqlite3.test.js @@ -115,7 +115,7 @@ t.test("it does not block in dry mode", (t) => { ); t.notMatch( stdout, - /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG/ + /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG, AIKIDO_INSTRUMENT/ ); }) .catch((error) => { @@ -171,7 +171,7 @@ t.test("it does not enable Zen when no environment variables are set", (t) => { t.equal(normalAdd.status, 200); t.match( stdout, - /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG/ + /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG, AIKIDO_INSTRUMENT/ ); }) .catch((error) => { diff --git a/end2end/tests/lambda-mongodb.test.js b/end2end/tests/lambda-mongodb.test.js index e8e0c25b1..c4d7c387c 100644 --- a/end2end/tests/lambda-mongodb.test.js +++ b/end2end/tests/lambda-mongodb.test.js @@ -72,7 +72,7 @@ t.test( t.notMatch( stdout, - /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG/ + /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG, AIKIDO_INSTRUMENT/ ); t.same(getJsonFromLogs(stdout.toString()), { @@ -102,7 +102,7 @@ t.test("it does not enable if no environment variable is set", async (t) => { t.match( stdout, - /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG/ + /AIKIDO: Zen is disabled\. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG, AIKIDO_INSTRUMENT/ ); t.same(getJsonFromLogs(stdout.toString()), { diff --git a/library/agent/Agent.ts b/library/agent/Agent.ts index 586831fd3..e0824d4e9 100644 --- a/library/agent/Agent.ts +++ b/library/agent/Agent.ts @@ -545,24 +545,27 @@ export class Agent { return; } - this.onStart() - .then(() => { - this.startHeartbeats(); - this.startPollingForConfigChanges(); - }) - .catch((err) => { - console.error(`Aikido: Failed to start agent: ${err.message}`); - }); + if (this.token) { + this.startReporting(); + } } - hasToken() { + hasToken(): boolean { return this.token !== undefined; } - setToken(token: Token) { + setToken(token: Token): void { this.token = token; this.logger.log("Token set, enabling reporting."); + if (this.serverless) { + return; + } + + this.startReporting(); + } + + private startReporting(): void { this.onStart() .then(() => { this.startHeartbeats(); diff --git a/library/agent/protect.setToken.test.ts b/library/agent/protect.setToken.test.ts index 722ea783a..4b54c8b72 100644 --- a/library/agent/protect.setToken.test.ts +++ b/library/agent/protect.setToken.test.ts @@ -4,23 +4,51 @@ import { ReportingAPIForTesting } from "./api/ReportingAPIForTesting"; import { Token } from "./api/Token"; import { getInstance, setInstance } from "./AgentSingleton"; import { createTestAgent } from "../helpers/createTestAgent"; -import { setToken } from "./protect"; +import { cloudFunction, lambda, setToken } from "./protect"; t.beforeEach(() => { // @ts-expect-error Reset singleton for isolation setInstance(undefined); + delete process.env.AIKIDO_BLOCK; + delete process.env.AIKIDO_BLOCKING; + delete process.env.AIKIDO_DEBUG; + delete process.env.AIKIDO_DISABLE; + delete process.env.AIKIDO_INSTRUMENT; + delete process.env.AIKIDO_TOKEN; }); -t.test("it sets the token on an agent without one", async (t) => { +t.test( + "it sets the token on a serverless agent without starting reporting", + async (t) => { + const api = new ReportingAPIForTesting(); + const agent = createTestAgent({ + api, + serverless: "lambda", + }); + agent.start([]); + + t.equal(agent.hasToken(), false); + + setToken("test-token-123"); + + t.equal(agent.hasToken(), true); + + // Give any accidental async reporting a tick to run + await new Promise((resolve) => setTimeout(resolve, 10)); + + const events = api.getEvents(); + const startedEvents = events.filter((e) => e.type === "started"); + t.equal(startedEvents.length, 0); + } +); + +t.test("it starts reporting for a non-serverless agent", async (t) => { const api = new ReportingAPIForTesting(); const agent = createTestAgent({ api, - serverless: "lambda", }); agent.start([]); - t.equal(agent.hasToken(), false); - setToken("test-token-123"); t.equal(agent.hasToken(), true); @@ -57,7 +85,7 @@ t.test("it warns if no agent is running", async (t) => { console.warn = originalWarn; t.equal(warnings.length, 1); - t.match(warnings[0], /agent is not running/); + t.match(warnings[0], /AIKIDO_INSTRUMENT/); }); t.test("it warns on empty string", async (t) => { @@ -72,3 +100,31 @@ t.test("it warns on empty string", async (t) => { t.equal(warnings.length, 1); t.match(warnings[0], /empty string/); }); + +t.test( + "AIKIDO_INSTRUMENT starts the lambda wrapper without a token", + async (t) => { + process.env.AIKIDO_INSTRUMENT = "true"; + + lambda(); + + const agent = getInstance(); + t.ok(agent); + t.equal(agent?.hasToken(), false); + t.equal(agent?.isServerless(), true); + } +); + +t.test( + "AIKIDO_INSTRUMENT starts the cloud function wrapper without a token", + async (t) => { + process.env.AIKIDO_INSTRUMENT = "true"; + + cloudFunction(); + + const agent = getInstance(); + t.ok(agent); + t.equal(agent?.hasToken(), false); + t.equal(agent?.isServerless(), true); + } +); diff --git a/library/agent/protect.ts b/library/agent/protect.ts index fc8f5d3df..215e62156 100644 --- a/library/agent/protect.ts +++ b/library/agent/protect.ts @@ -218,22 +218,22 @@ export function protectWithNewInstrumentation() { }); } -export function prepare() { +export function lambdaWithNewInstrumentation() { startAgent({ - serverless: undefined, - newInstrumentation: false, + serverless: "lambda", + newInstrumentation: true, }); } -export function prepareWithNewInstrumentation() { +export function cloudFunctionWithNewInstrumentation() { startAgent({ - serverless: undefined, + serverless: "gcp", newInstrumentation: true, }); } -export function setToken(token: string) { - if (token.length === 0) { +export function setToken(token: string | undefined) { + if (typeof token !== "string" || token.trim().length === 0) { // oxlint-disable-next-line no-console console.warn("AIKIDO: setToken called with an empty string, ignoring."); return; @@ -244,7 +244,7 @@ export function setToken(token: string) { if (!agent) { // oxlint-disable-next-line no-console console.warn( - "AIKIDO: setToken called but the agent is not running. Call prepare() first." + "AIKIDO: setToken called but the agent is not running. Set AIKIDO_INSTRUMENT=true before importing Zen." ); return; } diff --git a/library/cloud-function/index.ts b/library/cloud-function/index.ts index e6aee067c..61e6686c5 100644 --- a/library/cloud-function/index.ts +++ b/library/cloud-function/index.ts @@ -1,3 +1,11 @@ -import { cloudFunction } from "../agent/protect"; +import type { HttpFunction } from "@google-cloud/functions-framework"; +import { cloudFunction, setToken } from "../agent/protect"; -export = cloudFunction(); +type CloudFunctionWrapper = ((handler: HttpFunction) => HttpFunction) & { + setToken: typeof setToken; +}; + +const wrapper = cloudFunction() as CloudFunctionWrapper; +wrapper.setToken = setToken; + +export = wrapper; diff --git a/library/cloud-function/instrument.ts b/library/cloud-function/instrument.ts new file mode 100644 index 000000000..14377aded --- /dev/null +++ b/library/cloud-function/instrument.ts @@ -0,0 +1,9 @@ +import { + cloudFunctionWithNewInstrumentation, + setToken, +} from "../agent/protect"; +import { startWithNewInstrumentation } from "../instrument/start"; + +startWithNewInstrumentation(cloudFunctionWithNewInstrumentation); + +export { setToken }; diff --git a/library/helpers/shouldEnableFirewall.test.ts b/library/helpers/shouldEnableFirewall.test.ts index d39923080..db5b1f672 100644 --- a/library/helpers/shouldEnableFirewall.test.ts +++ b/library/helpers/shouldEnableFirewall.test.ts @@ -30,6 +30,15 @@ t.test("works with AIKIDO_TOKEN", async () => { t.same(shouldEnableFirewall(), false); }); +t.test("works with AIKIDO_INSTRUMENT", async () => { + process.env.AIKIDO_INSTRUMENT = "1"; + t.same(shouldEnableFirewall(), true); + process.env.AIKIDO_INSTRUMENT = "true"; + t.same(shouldEnableFirewall(), true); + process.env.AIKIDO_INSTRUMENT = ""; + t.same(shouldEnableFirewall(), false); +}); + t.test("it works if multiple are set", async () => { process.env.AIKIDO_DEBUG = "1"; process.env.AIKIDO_BLOCK = "1"; diff --git a/library/helpers/shouldEnableFirewall.ts b/library/helpers/shouldEnableFirewall.ts index 532322a74..e4d3c06c1 100644 --- a/library/helpers/shouldEnableFirewall.ts +++ b/library/helpers/shouldEnableFirewall.ts @@ -2,6 +2,7 @@ import { envToBool } from "./envToBool"; import { isAikidoCI } from "./isAikidoCI"; import { isDebugging } from "./isDebugging"; import { shouldBlock } from "./shouldBlock"; +import { shouldInstrument } from "./shouldInstrument"; /** * Only enable firewall if at least one of the following environment variables is set to a valid value: @@ -9,12 +10,17 @@ import { shouldBlock } from "./shouldBlock"; * - AIKIDO_BLOCK * - AIKIDO_TOKEN * - AIKIDO_DEBUG + * - AIKIDO_INSTRUMENT */ export default function shouldEnableFirewall() { if (envToBool(process.env.AIKIDO_DISABLE)) { return false; } + if (shouldInstrument()) { + return true; + } + if (shouldBlock()) { return true; } @@ -30,7 +36,7 @@ export default function shouldEnableFirewall() { if (!isAikidoCI()) { // oxlint-disable-next-line no-console console.log( - "AIKIDO: Zen is disabled. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG." + "AIKIDO: Zen is disabled. Configure one of the following environment variables to enable it: AIKIDO_BLOCK, AIKIDO_TOKEN, AIKIDO_DEBUG, AIKIDO_INSTRUMENT." ); } return false; diff --git a/library/helpers/shouldInstrument.ts b/library/helpers/shouldInstrument.ts new file mode 100644 index 000000000..513c04ee0 --- /dev/null +++ b/library/helpers/shouldInstrument.ts @@ -0,0 +1,5 @@ +import { envToBool } from "./envToBool"; + +export function shouldInstrument() { + return envToBool(process.env.AIKIDO_INSTRUMENT); +} diff --git a/library/index.ts b/library/index.ts index 474aa0216..1e21a40ed 100644 --- a/library/index.ts +++ b/library/index.ts @@ -20,7 +20,7 @@ import { withoutIdorProtection } from "./agent/context/withoutIdorProtection"; import { colorText } from "./helpers/colorText"; import { isPreloaded } from "./helpers/isPreloaded"; import { warnIfEntrypointIsModule } from "./helpers/warnIfEntrypointIsModule"; -import { prepare, setToken } from "./agent/protect"; +import { setToken } from "./agent/protect"; // Prevent logging twice / trying to start agent twice if (!isNewHookSystemUsed()) { @@ -72,7 +72,6 @@ export { setTenantId, enableIdorProtection, withoutIdorProtection, - prepare, setToken, }; @@ -93,6 +92,5 @@ export default { setTenantId, enableIdorProtection, withoutIdorProtection, - prepare, setToken, }; diff --git a/library/instrument/index.ts b/library/instrument/index.ts index 38dcf6bb0..03f01a772 100644 --- a/library/instrument/index.ts +++ b/library/instrument/index.ts @@ -1,61 +1,6 @@ -/* oxlint-disable no-console */ +import { protectWithNewInstrumentation, setToken } from "../agent/protect"; +import { startWithNewInstrumentation } from "./start"; -import * as mod from "node:module"; -import shouldEnableFirewall from "../helpers/shouldEnableFirewall"; -import isFirewallSupported from "../helpers/isFirewallSupported"; -import { protectWithNewInstrumentation } from "../agent/protect"; -import { setIsNewHookSystemUsed } from "../agent/isNewHookSystemUsed"; -import { checkIndexImportGuard } from "../helpers/indexImportGuard"; -import { isMainThread } from "node:worker_threads"; -import { isESM } from "../helpers/isESM"; -import { isPreloaded } from "../helpers/isPreloaded"; -import { colorText } from "../helpers/colorText"; -import { getInstance } from "../agent/AgentSingleton"; +startWithNewInstrumentation(protectWithNewInstrumentation); -setIsNewHookSystemUsed(true); - -const isSupported = isFirewallSupported(); -const alreadyRunning = !!getInstance(); -const shouldEnable = alreadyRunning || shouldEnableFirewall(); -const notAlreadyImported = checkIndexImportGuard(); - -function start() { - if (!isSupported || !shouldEnable || !notAlreadyImported) { - return; - } - - if (!("registerHooks" in mod) || typeof mod.registerHooks !== "function") { - console.error( - colorText( - "red", - "AIKIDO: Error: Zen requires that your Node.js version supports the `module.registerHooks` API. Please upgrade to a newer version of Node.js. See our ESM documentation for setup instructions (https://github.com/AikidoSec/firewall-node/blob/main/docs/esm.md)." - ) - ); - return; - } - - if (!isMainThread) { - console.warn( - "AIKIDO: Zen does not instrument worker threads. Zen will only be active in the main thread." - ); - return; - } - - if (isESM() === true && !isPreloaded()) { - console.error( - colorText( - "red", - "AIKIDO: Error: Your application seems to be running in ESM mode without preloading the library. Please use --require to preload the library. See our ESM documentation for setup instructions (https://github.com/AikidoSec/firewall-node/blob/main/docs/esm.md)." - ) - ); - } - - protectWithNewInstrumentation(); -} - -start(); - -export { - prepareWithNewInstrumentation as prepare, - setToken, -} from "../agent/protect"; +export { setToken }; diff --git a/library/instrument/start.ts b/library/instrument/start.ts new file mode 100644 index 000000000..96b2b03be --- /dev/null +++ b/library/instrument/start.ts @@ -0,0 +1,53 @@ +/* oxlint-disable no-console */ + +import * as mod from "node:module"; +import { isMainThread } from "node:worker_threads"; +import { getInstance } from "../agent/AgentSingleton"; +import { setIsNewHookSystemUsed } from "../agent/isNewHookSystemUsed"; +import { colorText } from "../helpers/colorText"; +import { checkIndexImportGuard } from "../helpers/indexImportGuard"; +import isFirewallSupported from "../helpers/isFirewallSupported"; +import { isESM } from "../helpers/isESM"; +import { isPreloaded } from "../helpers/isPreloaded"; +import shouldEnableFirewall from "../helpers/shouldEnableFirewall"; + +export function startWithNewInstrumentation(startAgent: () => void): void { + setIsNewHookSystemUsed(true); + + const isSupported = isFirewallSupported(); + const alreadyRunning = !!getInstance(); + const shouldEnable = alreadyRunning || shouldEnableFirewall(); + const notAlreadyImported = checkIndexImportGuard(); + + if (!isSupported || !shouldEnable || !notAlreadyImported) { + return; + } + + if (!("registerHooks" in mod) || typeof mod.registerHooks !== "function") { + console.error( + colorText( + "red", + "AIKIDO: Error: Zen requires that your Node.js version supports the `module.registerHooks` API. Please upgrade to a newer version of Node.js. See our ESM documentation for setup instructions (https://github.com/AikidoSec/firewall-node/blob/main/docs/esm.md)." + ) + ); + return; + } + + if (!isMainThread) { + console.warn( + "AIKIDO: Zen does not instrument worker threads. Zen will only be active in the main thread." + ); + return; + } + + if (isESM() === true && !isPreloaded()) { + console.error( + colorText( + "red", + "AIKIDO: Error: Your application seems to be running in ESM mode without preloading the library. Please use --require to preload the library. See our ESM documentation for setup instructions (https://github.com/AikidoSec/firewall-node/blob/main/docs/esm.md)." + ) + ); + } + + startAgent(); +} diff --git a/library/lambda/index.ts b/library/lambda/index.ts index 4f891372a..8d6c2747b 100644 --- a/library/lambda/index.ts +++ b/library/lambda/index.ts @@ -1,3 +1,11 @@ -import { lambda } from "../agent/protect"; +import type { Handler } from "aws-lambda"; +import { lambda, setToken } from "../agent/protect"; -export = lambda(); +type LambdaWrapper = ((handler: Handler) => Handler) & { + setToken: typeof setToken; +}; + +const wrapper = lambda() as LambdaWrapper; +wrapper.setToken = setToken; + +export = wrapper; diff --git a/library/lambda/instrument.ts b/library/lambda/instrument.ts new file mode 100644 index 000000000..6b5706bb2 --- /dev/null +++ b/library/lambda/instrument.ts @@ -0,0 +1,6 @@ +import { lambdaWithNewInstrumentation, setToken } from "../agent/protect"; +import { startWithNewInstrumentation } from "../instrument/start"; + +startWithNewInstrumentation(lambdaWithNewInstrumentation); + +export { setToken }; diff --git a/library/package.json b/library/package.json index d74e555ac..f22ebed50 100644 --- a/library/package.json +++ b/library/package.json @@ -13,8 +13,10 @@ "./instrument/internals": "./instrument/internals.js", "./context": "./context/index.js", "./lambda": "./lambda/index.js", + "./lambda/instrument": "./lambda/instrument.js", "./nopp": "./nopp/index.js", "./cloud-function": "./cloud-function/index.js", + "./cloud-function/instrument": "./cloud-function/instrument.js", "./bundler": "./bundler/index.js" }, "homepage": "https://aikido.dev/zen", diff --git a/library/sources/FunctionsFramework.test.ts b/library/sources/FunctionsFramework.test.ts index 35af5caa1..9eafd547c 100644 --- a/library/sources/FunctionsFramework.test.ts +++ b/library/sources/FunctionsFramework.test.ts @@ -203,6 +203,27 @@ t.test("it flushes stats first invoke", async (t) => { t.same(api.getEvents().length, 2); }); +t.test("it sends startup on first request after token is set", async (t) => { + const api = new ReportingAPIForTesting(); + const agent = createTestAgent({ + api, + serverless: "gcp", + }); + agent.start([]); + + api.clear(); + + const app = getExpressApp(); + + await request(app).get("/"); + t.same(api.getEvents(), []); + + agent.setToken(new Token("123")); + + await request(app).get("/"); + t.match(api.getEvents(), [{ type: "started" }, { type: "heartbeat" }]); +}); + t.test("it hooks into functions framework", async () => { const agent = createTestAgent({ serverless: "gcp", diff --git a/library/sources/FunctionsFramework.ts b/library/sources/FunctionsFramework.ts index 12862e9af..9708ecf61 100644 --- a/library/sources/FunctionsFramework.ts +++ b/library/sources/FunctionsFramework.ts @@ -45,7 +45,7 @@ export function createCloudFunctionWrapper(fn: HttpFunction): HttpFunction { return async (req, res) => { // Send startup event on first invocation - if (agent && !startupEventSent) { + if (agent && !startupEventSent && agent.hasToken()) { startupEventSent = true; try { await agent.onStart(getTimeoutInMS()); diff --git a/library/sources/Lambda.test.ts b/library/sources/Lambda.test.ts index ce0fa49c5..b4f3e126e 100644 --- a/library/sources/Lambda.test.ts +++ b/library/sources/Lambda.test.ts @@ -208,6 +208,30 @@ t.test("it passes through unknown types of events", async () => { t.same(result, undefined); }); +t.test("it sends startup on first invocation after token is set", async () => { + const testing = new ReportingAPIForTesting(); + const agent = createTestAgent({ + block: false, + serverless: "lambda", + api: testing, + }); + agent.start([]); + + const handler = createLambdaWrapper(async (event, context) => { + return getContext(); + }); + + testing.clear(); + + await handler(gatewayEvent, lambdaContext, () => {}); + t.same(testing.getEvents(), []); + + agent.setToken(new Token("token")); + + await handler(gatewayEvent, lambdaContext, () => {}); + t.match(testing.getEvents(), [{ type: "started" }, { type: "heartbeat" }]); +}); + t.test("it sends heartbeat after first and every 10 minutes", async () => { const clock = FakeTimers.install(); diff --git a/library/sources/Lambda.ts b/library/sources/Lambda.ts index 9fe4bc443..0b3cccb6c 100644 --- a/library/sources/Lambda.ts +++ b/library/sources/Lambda.ts @@ -154,7 +154,7 @@ export function createLambdaWrapper(handler: Handler): Handler { return async (event, context) => { // Send startup event on first invocation - if (agent && !startupEventSent) { + if (agent && !startupEventSent && agent.hasToken()) { startupEventSent = true; try { await agent.onStart(getTimeoutInMS()); diff --git a/library/start/index.ts b/library/start/index.ts new file mode 100644 index 000000000..aacd72727 --- /dev/null +++ b/library/start/index.ts @@ -0,0 +1,9 @@ +import isFirewallSupported from "../helpers/isFirewallSupported"; +import shouldEnableFirewall from "../helpers/shouldEnableFirewall"; +import { protect, setToken } from "../agent/protect"; + +if (isFirewallSupported() && shouldEnableFirewall({ enabledByDefault: true })) { + protect(); +} + +export { setToken }; diff --git a/sample-apps/express-mysql/app-set-token.js b/sample-apps/express-mysql/app-set-token.js index ab892b6d3..4ae3f3887 100644 --- a/sample-apps/express-mysql/app-set-token.js +++ b/sample-apps/express-mysql/app-set-token.js @@ -1,7 +1,5 @@ const Zen = require("@aikidosec/firewall"); -Zen.prepare(); - const Cats = require("./Cats"); const express = require("express"); const asyncHandler = require("express-async-handler"); @@ -45,7 +43,7 @@ async function main(port) { // Set the token after startup (simulates fetching from a secrets manager) setTimeout(() => { - Zen.setToken(process.env.AIKIDO_TOKEN); + Zen.setToken(process.env.TEST_AIKIDO_TOKEN); }, 500); return new Promise((resolve, reject) => { diff --git a/sample-apps/hono-pg-esm/app-set-token.js b/sample-apps/hono-pg-esm/app-set-token.js index 1b6dd5e7c..5fbe064a3 100644 --- a/sample-apps/hono-pg-esm/app-set-token.js +++ b/sample-apps/hono-pg-esm/app-set-token.js @@ -36,7 +36,7 @@ function getPort() { // Set the token after startup (simulates fetching from a secrets manager) setTimeout(() => { - setToken(process.env.AIKIDO_TOKEN); + setToken(process.env.TEST_AIKIDO_TOKEN); }, 500); serve( diff --git a/sample-apps/hono-pg-esm/zen-setup.cjs b/sample-apps/hono-pg-esm/zen-setup.cjs index 9d48e4f7b..879ef2660 100644 --- a/sample-apps/hono-pg-esm/zen-setup.cjs +++ b/sample-apps/hono-pg-esm/zen-setup.cjs @@ -1,3 +1 @@ -const { prepare } = require("@aikidosec/firewall/instrument"); - -prepare(); +require("@aikidosec/firewall/instrument"); diff --git a/sample-apps/test-exports/test.js b/sample-apps/test-exports/test.js index 7bca8c6f2..67bef44af 100644 --- a/sample-apps/test-exports/test.js +++ b/sample-apps/test-exports/test.js @@ -10,5 +10,7 @@ assert.ok(typeof Zen.addExpressMiddleware === "function"); assert.ok(typeof Zen.setUser === "function"); assert.ok(typeof context.setUser === "function"); assert.ok(typeof lambda === "function"); +assert.ok(typeof lambda.setToken === "function"); assert.ok(typeof cloudFunction === "function"); +assert.ok(typeof cloudFunction.setToken === "function"); assert.ok(Array.isArray(bundler.externals())); diff --git a/sample-apps/test-exports/test.ts b/sample-apps/test-exports/test.ts index 43015472d..6570b90d7 100644 --- a/sample-apps/test-exports/test.ts +++ b/sample-apps/test-exports/test.ts @@ -11,5 +11,7 @@ assert.ok(typeof Zen.addExpressMiddleware === "function"); assert.ok(typeof Zen.setUser === "function"); assert.ok(typeof context.setUser === "function"); assert.ok(typeof lambda === "function"); +assert.ok(typeof lambda.setToken === "function"); assert.ok(typeof cloudFunction === "function"); +assert.ok(typeof cloudFunction.setToken === "function"); assert.ok(Array.isArray(externals()));