Allow passing a Router to addExpressMiddleware

Also log warning when middleware is called multiple times during a
request.

Author: hansott

docs/express.md

@@ -59,6 +59,19 @@ Zen.addExpressMiddleware(app);
 app.get(...);
 ```
 
+You can also pass a `Router` instance to `Zen.addExpressMiddleware`:
+
+```js
+const router = express.Router();
+
+// Note: The middleware should run once per request
+Zen.addExpressMiddleware(router);
+
+router.get(...);
+
+app.use(router);
+```
+
 ## Debug mode
 
 If you need to debug the firewall, you can run your express app with the environment variable `AIKIDO_DEBUG` set to `true`:

library/middleware/express.ts

@@ -1,5 +1,5 @@
 /** TS_EXPECT_TYPES_ERROR_OPTIONAL_DEPENDENCY **/
-import type { Express } from "express";
+import type { Express, Router } from "express";
 import { shouldBlockRequest } from "./shouldBlockRequest";
 import { escapeHTML } from "../helpers/escapeHTML";
 
@@ -8,7 +8,7 @@ import { escapeHTML } from "../helpers/escapeHTML";
  * 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 addExpressMiddleware(app: Express) {
+export function addExpressMiddleware(app: Express | Router) {
   app.use((req, res, next) => {
     const result = shouldBlockRequest();
 

library/middleware/shouldBlockRequest.test.ts

@@ -19,19 +19,23 @@ const sampleContext: Context = {
   route: "/posts/:id",
 };
 
-t.test("without context", async (t) => {
-  const logs: string[] = [];
-  wrap(console, "warn", function warn() {
-    return function warn(message: string) {
-      logs.push(message);
-    };
-  });
+let logs: string[] = [];
+wrap(console, "warn", function warn() {
+  return function warn(message: string) {
+    logs.push(message);
+  };
+});
 
+t.beforeEach(() => {
+  logs = [];
+});
+
+t.test("without context", async (t) => {
   const result = shouldBlockRequest();
   shouldBlockRequest();
   t.same(result, { block: false });
   t.same(logs, [
-    "shouldBlockRequest() was called without a context. The request will not be blocked. Make sure to call shouldBlockRequest() within an HTTP request. If you're using serverless functions, make sure to use the handler wrapper provided by Zen. Also ensure you import Zen at the top of your main app file (before any other imports).",
+    "Zen.shouldBlockRequest() was called without a context. The request will not be blocked. Make sure to call shouldBlockRequest() within an HTTP request. If you're using serverless functions, make sure to use the handler wrapper provided by Zen. Also ensure you import Zen at the top of your main app file (before any other imports).",
   ]);
 });
 
@@ -47,3 +51,14 @@ t.test("with agent", async (t) => {
     t.same(shouldBlockRequest(), { block: false });
   });
 });
+
+t.test("multiple calls", async (t) => {
+  createTestAgent();
+  runWithContext(sampleContext, () => {
+    shouldBlockRequest();
+    shouldBlockRequest();
+  });
+  t.same(logs, [
+    "Zen.shouldBlockRequest() was called multiple times. The middleware should be executed once per request.",
+  ]);
+});

library/middleware/shouldBlockRequest.ts

@@ -21,6 +21,11 @@ export function shouldBlockRequest(): Result {
     return { block: false };
   }
 
+  if (context.executedMiddleware) {
+    logWarningAlreadyExecutedMiddleware();
+    return { block: false };
+  }
+
   updateContext(context, "executedMiddleware", true);
   agent.onMiddlewareExecuted();
 
@@ -50,8 +55,23 @@ function logWarningShouldBlockRequestCalledWithoutContext() {
 
   // eslint-disable-next-line no-console
   console.warn(
-    "shouldBlockRequest() was called without a context. The request will not be blocked. Make sure to call shouldBlockRequest() within an HTTP request. If you're using serverless functions, make sure to use the handler wrapper provided by Zen. Also ensure you import Zen at the top of your main app file (before any other imports)."
+    "Zen.shouldBlockRequest() was called without a context. The request will not be blocked. Make sure to call shouldBlockRequest() within an HTTP request. If you're using serverless functions, make sure to use the handler wrapper provided by Zen. Also ensure you import Zen at the top of your main app file (before any other imports)."
   );
 
   loggedWarningShouldBlockRequestCalledWithoutContext = true;
 }
+
+let loggedWarningAlreadyExecutedMiddleware = false;
+
+function logWarningAlreadyExecutedMiddleware() {
+  if (loggedWarningAlreadyExecutedMiddleware) {
+    return;
+  }
+
+  // eslint-disable-next-line no-console
+  console.warn(
+    "Zen.shouldBlockRequest() was called multiple times. The middleware should be executed once per request."
+  );
+
+  loggedWarningAlreadyExecutedMiddleware = true;
+}

library/sources/Express.tests.ts

@@ -718,4 +718,29 @@ export function createExpressTests(expressPackageName: string) {
       );
     }
   );
+
+  t.test("it supports adding middleware to a Router instance", async (t) => {
+    const app = express();
+    const router = express.Router();
+
+    // Add user middleware to set user data
+    router.use((req, res, next) => {
+      setUser({ id: "567" });
+      next();
+    });
+
+    // Add Zen middleware to router instead of app
+    addExpressMiddleware(router);
+
+    router.get("/router-block-user", (req, res) => {
+      res.send({ willNotBeSent: true });
+    });
+
+    app.use(router);
+
+    // Test blocked user is properly handled by the router middleware
+    const blockedResponse = await request(app).get("/router-block-user");
+    t.same(blockedResponse.statusCode, 403);
+    t.same(blockedResponse.text, "You are blocked by Zen.");
+  });
 }