Proof of concept for serverless async check

Author: hansott

library/agent/Agent.ts

@@ -546,4 +546,8 @@ export class Agent {
   onMiddlewareExecuted() {
     this.middlewareInstalled = true;
   }
+
+  getToken() {
+    return this.token;
+  }
 }

library/index.ts

@@ -3,7 +3,10 @@ import isFirewallSupported from "./helpers/isFirewallSupported";
 import shouldEnableFirewall from "./helpers/shouldEnableFirewall";
 import { setUser } from "./agent/context/user";
 import { shouldBlockRequest } from "./middleware/shouldBlockRequest";
-import { addExpressMiddleware } from "./middleware/express";
+import {
+  addExpressMiddleware,
+  addExpressMiddlewareAsync,
+} from "./middleware/express";
 import { addHonoMiddleware } from "./middleware/hono";
 import { addHapiMiddleware } from "./middleware/hapi";
 import { addFastifyHook } from "./middleware/fastify";
@@ -20,6 +23,7 @@ export {
   setUser,
   shouldBlockRequest,
   addExpressMiddleware,
+  addExpressMiddlewareAsync,
   addHonoMiddleware,
   addHapiMiddleware,
   addFastifyHook,
@@ -32,6 +36,7 @@ export default {
   setUser,
   shouldBlockRequest,
   addExpressMiddleware,
+  addExpressMiddlewareAsync,
   addHonoMiddleware,
   addHapiMiddleware,
   addFastifyHook,

library/middleware/express.ts

@@ -1,6 +1,9 @@
 /** TS_EXPECT_TYPES_ERROR_OPTIONAL_DEPENDENCY **/
 import type { Express } from "express";
-import { shouldBlockRequest } from "./shouldBlockRequest";
+import {
+  shouldBlockRequest,
+  shouldBlockRequestAsync,
+} from "./shouldBlockRequest";
 import { escapeHTML } from "../helpers/escapeHTML";
 
 /**
@@ -30,3 +33,31 @@ export function addExpressMiddleware(app: Express) {
     next();
   });
 }
+
+export function addExpressMiddlewareAsync(app: Express) {
+  app.use((req, res, next) => {
+    shouldBlockRequestAsync()
+      .then((result) => {
+        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)})`;
+            }
+
+            return res.status(429).type("text").send(message);
+          }
+
+          if (result.type === "blocked") {
+            return res.status(403).type("text").send("You are blocked by Zen.");
+          }
+        }
+
+        next();
+      })
+      .catch((error) => {
+        console.error(error);
+        next();
+      });
+  });
+}

library/middleware/shouldBlockRequest.ts

@@ -9,6 +9,38 @@ type Result = {
   ip?: string;
 };
 
+const SERVERLESS_URL = "http://localhost:5132";
+
+export async function shouldBlockRequestAsync(): Promise<Result> {
+  const context = getContext();
+  if (!context) {
+    return { block: false };
+  }
+
+  const agent = getInstance();
+  if (!agent) {
+    return { block: false };
+  }
+
+  const response = await fetch(`${SERVERLESS_URL}/check-request`, {
+    method: "POST",
+    headers: {
+      // TODO: Remove "!", token is not guaranteed to be present
+      Authorization: agent.getToken()!.asString(),
+      "Content-Type": "application/json",
+    },
+    body: JSON.stringify({
+      method: context.method,
+      headers: context.headers,
+      url: context.url,
+      route: context.route,
+      clientIp: context.remoteAddress,
+    }),
+  });
+
+  return await response.json();
+}
+
 export function shouldBlockRequest(): Result {
   const context = getContext();
   if (!context) {

sample-apps/express-mysql/app.js

@@ -1,5 +1,5 @@
 require("dotenv").config();
-require("@aikidosec/firewall");
+const Zen = require("@aikidosec/firewall");
 const Sentry = require("@sentry/node");
 
 Sentry.init({
@@ -59,6 +59,8 @@ async function main(port) {
 
   const app = express();
 
+  Zen.addExpressMiddlewareAsync(app);
+
   app.use(Sentry.Handlers.requestHandler());
   app.use(morgan("tiny"));