Send heartbeat on shutdown

docs/express.md

@@ -90,3 +90,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 
 That's it! Your app is now protected by Zen.  
 If you want to see a full example, check our [express sample app](../sample-apps/express-mongodb).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/fastify.md

@@ -77,3 +77,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 
 That's it! Your app is now protected by Zen.  
 If you want to see a full example, check our [fastify sample app](../sample-apps/fastify-mysql2).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/graceful-shutdown.md

@@ -0,0 +1,20 @@
+# Graceful Shutdown
+
+If you already registered a shutdown handler, e.g. using `process.on('SIGTERM', ...)`, you can call `await Zen.shutdown()` to ensure the latest stats are sent before the process exits.
+See the example below for how to implement this, if you don't already have a shutdown handler in place.
+Please note that this might not work correctly on Windows, due to [differences in how signals like SIGTERM and SIGINT are handled](https://nodejs.org/api/process.html#signal-events).
+
+```js
+const Zen = require("@aikidosec/firewall");
+
+async function shutdownHandler() {
+  // Perform any cleanup or final tasks here
+  await Zen.shutdown();
+
+  process.exit(process.exitCode || 0);
+}
+
+process.on("SIGTERM", shutdownHandler);
+process.on("SIGINT", shutdownHandler);
+// Handle other signals as needed
+```

docs/hapi.md

@@ -81,3 +81,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 
 That's it! Your app is now protected by Zen.  
 If you want to see a full example, check our [hapi sample app](../sample-apps/hapi-postgres).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/hono.md

@@ -79,3 +79,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 
 That's it! Your app is now protected by Zen.  
 If you want to see a full example, check our [hono sample app](../sample-apps/hono-mongodb).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/koa.md

@@ -84,3 +84,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 
 That's it! Your app is now protected by Zen.  
 If you want to see a full example, check our [koa sample app](../sample-apps/koa-sqlite3).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/micro.md

@@ -86,3 +86,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 That's it! Your app is now protected by Zen.
 
 If you want to see a full example, check our [micro sample app](../sample-apps/micro).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/nestjs.md

@@ -92,3 +92,7 @@ Read [Protect against prototype pollution](./prototype-pollution.md) to learn ho
 
 That's it! Your app is now protected by Zen.  
 If you want to see a full example, check our [NestJS sample app](../sample-apps/nestjs-sentry).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/next.md

@@ -71,3 +71,7 @@ AIKIDO_DEBUG=true node -r @aikidosec/firewall .next/standalone/server.js
 ```
 
 This will output debug information to the console (e.g. if the agent failed to start, no token was found, unsupported packages, ...).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/pubsub.md

@@ -48,3 +48,7 @@ This will output debug information to the console (e.g. if the agent failed to s
 Zen can also protect your application against prototype pollution attacks.
 
 Read [Protect against prototype pollution](./prototype-pollution.md) to learn how to set it up.
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

docs/restify.md

@@ -85,4 +85,8 @@ Zen can also protect your application against prototype pollution attacks.
 Read [Protect against prototype pollution](./prototype-pollution.md) to learn how to set it up.
 
 That's it! Your app is now protected by Zen.  
-If you want to see a full example, check our [Restify sample app](../sample-apps/restify-postgres). 
+If you want to see a full example, check our [Restify sample app](../sample-apps/restify-postgres).
+
+## Graceful shutdown
+
+It is recommended to add a shutdown handler to your app to ensure that no statistics are lost when the app is stopped. You can find more information [here](./graceful-shutdown.md).

end2end/tests/hono-sqlite3.test.js

@@ -14,10 +14,6 @@ t.test("it blocks in blocking mode", (t) => {
     env: { ...process.env, AIKIDO_DEBUG: "true", AIKIDO_BLOCKING: "true" },
   });
 
-  server.on("close", () => {
-    t.end();
-  });
-
   server.on("error", (err) => {
     t.fail(err.message);
   });
@@ -32,6 +28,11 @@ t.test("it blocks in blocking mode", (t) => {
     stderr += data.toString();
   });
 
+  server.on("close", () => {
+    t.match(stdout, /AIKIDO: Shutting down agent.../);
+    t.end();
+  });
+
   // Wait for the server to start
   timeout(2000)
     .then(() => {

library/agent/Agent.test.ts

@@ -19,6 +19,7 @@ import { Context } from "./Context";
 import { createTestAgent } from "../helpers/createTestAgent";
 import { setTimeout } from "node:timers/promises";
 import type { Response } from "./api/fetchBlockedLists";
+import { shutdown } from "./shutdown";
 
 let shouldOnlyAllowSomeIPAddresses = false;
 
@@ -1230,3 +1231,42 @@ t.test("it includes agent's own package in heartbeat", async () => {
 
   clock.uninstall();
 });
+
+t.test("it sends heartbeat when onShutdown handler is called", async () => {
+  const clock = FakeTimers.install();
+
+  const logger = new LoggerNoop();
+  const api = new ReportingAPIForTesting();
+  const agent = createTestAgent({
+    api,
+    logger,
+    token: new Token("123"),
+    suppressConsoleLog: false,
+  });
+  agent.start([]);
+
+  // After 5 seconds, nothing should happen
+  clock.tick(1000 * 5);
+
+  t.match(api.getEvents(), [
+    {
+      type: "started",
+    },
+  ]);
+
+  clock.tick(1 * 90 * 1000);
+  await clock.nextAsync();
+
+  await shutdown();
+
+  t.match(api.getEvents(), [
+    {
+      type: "started",
+    },
+    {
+      type: "heartbeat",
+    },
+  ]);
+
+  clock.uninstall();
+});

library/agent/Agent.ts

@@ -609,4 +609,15 @@ export class Agent {
         return this.sendHeartbeatEveryMS;
     }
   }
+
+  public async shutdown() {
+    this.logger.log("Shutting down agent...");
+    if (
+      performance.now() > 30000 &&
+      performance.now() - this.lastHeartbeat > 3
+    ) {
+      // Only send heartbeat if we are running for more than 30 seconds and haven't sent a heartbeat in the last 3 seconds
+      await this.flushStats(1000);
+    }
+  }
 }

library/agent/shutdown.ts

@@ -0,0 +1,5 @@
+import { getInstance } from "./AgentSingleton";
+
+export async function shutdown(): Promise<void> {
+  await getInstance()?.shutdown();
+}

library/index.ts

@@ -13,6 +13,7 @@ import { addRestifyMiddleware } from "./middleware/restify";
 import { isESM } from "./helpers/isESM";
 import { checkIndexImportGuard } from "./helpers/indexImportGuard";
 import { setRateLimitGroup } from "./ratelimiting/group";
+import { shutdown } from "./agent/shutdown";
 
 const supported = isFirewallSupported();
 const shouldEnable = shouldEnableFirewall();
@@ -39,6 +40,7 @@ export {
   addKoaMiddleware,
   addRestifyMiddleware,
   setRateLimitGroup,
+  shutdown,
 };
 
 // Required for ESM / TypeScript default export support
@@ -54,4 +56,5 @@ export default {
   addKoaMiddleware,
   addRestifyMiddleware,
   setRateLimitGroup,
+  shutdown,
 };

sample-apps/hono-sqlite3/app.js

@@ -115,3 +115,9 @@ main().then((app) => {
     console.log(`Server is running on port ${port}`);
   });
 });
+
+process.on("SIGTERM", async () => {
+  console.log("SIGTERM received, shutting down gracefully...");
+  await Zen.shutdown();
+  process.exit(process.exitCode || 0);
+});