Merge pull request #362 from AikidoSec/koa

Add support for koa

Author: hansott

Makefile

@@ -72,6 +72,10 @@ nestjs-sentry:
 fastify-mysql2:
 	cd sample-apps/fastify-mysql2 && AIKIDO_DEBUG=true AIKIDO_BLOCKING=true node --preserve-symlinks app.js
 
+.PHONY: koa-sqlite3
+koa-sqlite3:
+	cd sample-apps/koa-sqlite3 && AIKIDO_DEBUG=true AIKIDO_BLOCKING=true node --preserve-symlinks app.js
+
 .PHONY: install
 install:
 	mkdir -p build

README.md

@@ -44,6 +44,7 @@ Zen for Node.js 16+ is compatible with:
 * ✅ [micro](docs/micro.md) 10.x
 * ✅ [Next.js](docs/next.md) 12.x, 13.x and 14.x
 * ✅ [Fastify](docs/fastify.md) 4.x and 5.x
+* ✅ [Koa](docs/koa.md) 2.x
 
 ### Database drivers
 
@@ -89,6 +90,11 @@ See list above for supported database drivers.
 
 * ✅ [`ShellJS`](https://www.npmjs.com/package/shelljs) 0.8.x, 0.7.x
 
+### Routers
+
+* ✅ [`@koa/router`](https://www.npmjs.com/package/@koa/router) 13.x, 12.x, 11.x and 10.x
+
+
 ## Installation
 
 We recommend testing Zen locally or on staging before deploying to production.

docs/koa.md

@@ -0,0 +1,86 @@
+# Koa
+
+At the very beginning of your app.js file, add the following line:
+
+```js
+require('@aikidosec/firewall'); // <-- Include this before any other code or imports
+
+const Koa = require("koa");
+
+const app = Koa();
+
+app.use(...);
+
+// ...
+```
+
+or ESM import style:
+
+```js
+import "@aikidosec/firewall";
+
+// ...
+```
+
+Zen also supports `@koa/router` or `koa-router`.
+
+## Blocking mode
+
+By default, the firewall will run in non-blocking mode. When it detects an attack, the attack will be reported to Aikido if the environment variable `AIKIDO_TOKEN` is set and continue executing the call.
+
+You can enable blocking mode by setting the environment variable `AIKIDO_BLOCK` to `true`:
+
+```sh
+AIKIDO_BLOCK=true node app.js
+```
+
+It's recommended to enable this on your staging environment for a considerable amount of time before enabling it on your production environment (e.g. one week).
+
+## Rate limiting and user blocking
+
+If you want to add the rate limiting feature to your app, modify your code like this:
+
+```js
+const Zen = require("@aikidosec/firewall");
+
+const app = Koa();
+
+// Optional, if you want to use user based rate limiting or block specific users
+app.use(async (ctx, next) => {
+  // Get the user from your authentication middleware
+  // or wherever you store the user
+  Zen.setUser({
+    id: "123",
+    name: "John Doe", // Optional
+  });
+
+  await next();
+});
+
+// Place this middleware after your authentication middleware
+// As early as possible in the middleware chain
+Zen.addKoaMiddleware(app);
+
+app.get(...);
+```
+
+If you are using `@koa/router` or `koa-router`, please make sure to place the `.use(router.routes())` middleware after the Zen middleware:
+
+## Debug mode
+
+If you need to debug the firewall, you can run your express app with the environment variable `AIKIDO_DEBUG` set to `true`:
+
+```sh
+AIKIDO_DEBUG=true node app.js
+```
+
+This will output debug information to the console (e.g. if the agent failed to start, no token was found, unsupported packages, ...).
+
+## Preventing prototype pollution
+
+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 [koa sample app](../sample-apps/koa-sqlite3).

end2end/tests/koa-sqlite3.test.js

@@ -0,0 +1,120 @@
+const t = require("tap");
+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const timeout = require("../timeout");
+
+const pathToApp = resolve(__dirname, "../../sample-apps/koa-sqlite3", "app.js");
+
+t.test("it blocks in blocking mode", (t) => {
+  const server = spawn(`node`, ["--preserve-symlinks", pathToApp, "4002"], {
+    env: { ...process.env, 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 the server to start
+  timeout(2000)
+    .then(() => {
+      return Promise.all([
+        fetch("http://127.0.0.1:4002/add", {
+          method: "POST",
+          body: JSON.stringify({ name: "Test'), ('Test2');--" }),
+          headers: {
+            "Content-Type": "application/json",
+          },
+          signal: AbortSignal.timeout(5000),
+        }),
+        fetch("http://127.0.0.1:4002/add", {
+          method: "POST",
+          body: JSON.stringify({ name: "Miau" }),
+          headers: {
+            "Content-Type": "application/json",
+          },
+          signal: AbortSignal.timeout(5000),
+        }),
+      ]);
+    })
+    .then(([sqlInjection, normalAdd]) => {
+      t.equal(sqlInjection.status, 500);
+      t.equal(normalAdd.status, 200);
+      t.match(stdout, /Starting agent/);
+      t.match(stderr, /Zen has blocked an SQL injection/);
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});
+
+t.test("it does not block in dry mode", (t) => {
+  const server = spawn(`node`, ["--preserve-symlinks", pathToApp, "4003"], {
+    env: { ...process.env, AIKIDO_DEBUG: "true" },
+  });
+
+  server.on("close", () => {
+    t.end();
+  });
+
+  let stdout = "";
+  server.stdout.on("data", (data) => {
+    stdout += data.toString();
+  });
+
+  let stderr = "";
+  server.stderr.on("data", (data) => {
+    stderr += data.toString();
+  });
+
+  // Wait for the server to start
+  timeout(2000)
+    .then(() =>
+      Promise.all([
+        fetch("http://127.0.0.1:4003/add", {
+          method: "POST",
+          body: JSON.stringify({ name: "Test'), ('Test2');--" }),
+          headers: {
+            "Content-Type": "application/json",
+          },
+          signal: AbortSignal.timeout(5000),
+        }),
+        fetch("http://127.0.0.1:4003/add", {
+          method: "POST",
+          body: JSON.stringify({ name: "Miau" }),
+          headers: {
+            "Content-Type": "application/json",
+          },
+          signal: AbortSignal.timeout(5000),
+        }),
+      ])
+    )
+    .then(([sqlInjection, normalAdd]) => {
+      t.equal(sqlInjection.status, 200);
+      t.equal(normalAdd.status, 200);
+      t.match(stdout, /Starting agent/);
+      t.notMatch(stderr, /Zen has blocked an SQL injection/);
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});

library/agent/protect.ts

@@ -45,6 +45,8 @@ import { isDebugging } from "../helpers/isDebugging";
 import { shouldBlock } from "../helpers/shouldBlock";
 import { Postgresjs } from "../sinks/Postgresjs";
 import { Fastify } from "../sources/Fastify";
+import { Koa } from "../sources/Koa";
+import { KoaRouter } from "../sources/KoaRouter";
 
 function getLogger(): Logger {
   if (isDebugging()) {
@@ -132,6 +134,8 @@ function getWrappers() {
     new BetterSQLite3(),
     new Postgresjs(),
     new Fastify(),
+    new Koa(),
+    new KoaRouter(),
   ];
 }
 

library/index.ts

@@ -7,6 +7,7 @@ import { addExpressMiddleware } from "./middleware/express";
 import { addHonoMiddleware } from "./middleware/hono";
 import { addHapiMiddleware } from "./middleware/hapi";
 import { addFastifyHook } from "./middleware/fastify";
+import { addKoaMiddleware } from "./middleware/koa";
 
 const supported = isFirewallSupported();
 const shouldEnable = shouldEnableFirewall();
@@ -22,4 +23,5 @@ export {
   addHonoMiddleware,
   addHapiMiddleware,
   addFastifyHook,
+  addKoaMiddleware,
 };

library/middleware/koa.ts

@@ -0,0 +1,37 @@
+import { shouldBlockRequest } from "./shouldBlockRequest";
+import { escapeHTML } from "../helpers/escapeHTML";
+import type * as Application from "koa";
+
+/**
+ * Calling this function will setup rate limiting and user blocking for the provided Express app.
+ * Attacks will still be blocked by Zen if you do not call this function.
+ * Execute this function as early as possible in your Express app, but after the middleware that sets the user.
+ */
+export function addKoaMiddleware(app: Application): void {
+  app.use(async (ctx, next) => {
+    const result = shouldBlockRequest();
+
+    if (result.block) {
+      if (result.type === "ratelimited") {
+        let message = "You are rate limited by Zen.";
+        if (result.trigger === "ip" && result.ip) {
+          message += ` (Your IP: ${escapeHTML(result.ip)})`;
+        }
+
+        ctx.type = "text/plain";
+        ctx.body = message;
+        ctx.status = 429;
+        return;
+      }
+
+      if (result.type === "blocked") {
+        ctx.type = "text/plain";
+        ctx.body = "You are blocked by Zen.";
+        ctx.status = 403;
+        return;
+      }
+    }
+
+    await next();
+  });
+}