Add WAF rule engine prototype using wirefilter

end2end/server/src/zen/config.ts

@@ -1,5 +1,11 @@
 import type { App } from "./apps.ts";
 
+type WafRule = {
+  id: string;
+  expression: string;
+  action?: string;
+};
+
 type AppConfig = {
   success: boolean;
   serviceId: number;
@@ -10,6 +16,7 @@ type AppConfig = {
   allowedIPAddresses: string[];
   blockNewOutgoingRequests: boolean;
   domains: any[];
+  wafRules?: WafRule[];
   failureRate?: number;
   timeout?: number;
 };

end2end/tests-new/waf-rules.test.mjs

@@ -0,0 +1,136 @@
+import { spawn } from "child_process";
+import { resolve } from "path";
+import { test } from "node:test";
+import { equal } 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-xml"
+);
+
+const testServerUrl = "http://localhost:5874";
+
+test("it blocks requests matching WAF path rule", 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 configResp = await fetch(`${testServerUrl}/api/runtime/config`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: token,
+    },
+    body: JSON.stringify({
+      wafRules: [
+        {
+          id: "block-admin",
+          expression: 'http.request.uri.path contains "/admin"',
+          action: "block",
+        },
+      ],
+    }),
+  });
+  equal(configResp.status, 200);
+
+  const server = spawn(`node`, ["./app.js", port], {
+    cwd: pathToAppDir,
+    env: {
+      ...process.env,
+      AIKIDO_DEBUG: "true",
+      AIKIDO_BLOCKING: "true",
+      AIKIDO_TOKEN: token,
+      AIKIDO_ENDPOINT: testServerUrl,
+      AIKIDO_REALTIME_ENDPOINT: testServerUrl,
+    },
+  });
+
+  try {
+    await timeout(2000);
+
+    const blocked = await fetch(`http://127.0.0.1:${port}/admin`, {
+      headers: { "X-Forwarded-For": "1.2.3.4" },
+      signal: AbortSignal.timeout(5000),
+    });
+    equal(blocked.status, 403);
+    equal(await blocked.text(), "You are blocked by Zen.");
+
+    const allowed = await fetch(`http://127.0.0.1:${port}/`, {
+      headers: { "X-Forwarded-For": "1.2.3.4" },
+      signal: AbortSignal.timeout(5000),
+    });
+    equal(allowed.status, 200);
+  } finally {
+    server.kill();
+  }
+});
+
+test("it blocks requests matching user agent WAF rule", 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 configResp = await fetch(`${testServerUrl}/api/runtime/config`, {
+    method: "POST",
+    headers: {
+      "Content-Type": "application/json",
+      Authorization: token,
+    },
+    body: JSON.stringify({
+      wafRules: [
+        {
+          id: "block-scanners",
+          expression: 'http.user_agent matches "(?i)(sqlmap|nikto)"',
+          action: "block",
+        },
+      ],
+    }),
+  });
+  equal(configResp.status, 200);
+
+  const server = spawn(`node`, ["./app.js", port], {
+    cwd: pathToAppDir,
+    env: {
+      ...process.env,
+      AIKIDO_DEBUG: "true",
+      AIKIDO_BLOCKING: "true",
+      AIKIDO_TOKEN: token,
+      AIKIDO_ENDPOINT: testServerUrl,
+      AIKIDO_REALTIME_ENDPOINT: testServerUrl,
+    },
+  });
+
+  try {
+    await timeout(2000);
+
+    const blocked = await fetch(`http://127.0.0.1:${port}/`, {
+      headers: {
+        "User-Agent": "sqlmap/1.0",
+        "X-Forwarded-For": "1.2.3.4",
+      },
+      signal: AbortSignal.timeout(5000),
+    });
+    equal(blocked.status, 403);
+
+    const allowed = await fetch(`http://127.0.0.1:${port}/`, {
+      headers: {
+        "User-Agent": "Mozilla/5.0",
+        "X-Forwarded-For": "1.2.3.4",
+      },
+      signal: AbortSignal.timeout(5000),
+    });
+    equal(allowed.status, 200);
+  } finally {
+    server.kill();
+  }
+});

library/agent/Agent.ts

@@ -18,6 +18,7 @@ import { Kind } from "./Attack";
 import { Endpoint } from "./Config";
 import { pollForChanges } from "./realtime/pollForChanges";
 import { Context } from "./Context";
+import { setWafRules } from "../waf/waf";
 import { Hostnames } from "./Hostnames";
 import { InspectionStatistics } from "./InspectionStatistics";
 import { Logger } from "./logger/Logger";
@@ -338,6 +339,15 @@ export class Agent {
         );
         this.serviceConfig.updateDomains(response.domains);
       }
+
+      if (response.wafRules && Array.isArray(response.wafRules)) {
+        const wafResult = setWafRules(response.wafRules);
+        if (!wafResult.success) {
+          this.logger.log(
+            `Failed to compile WAF rule ${wafResult.rule_id}: ${wafResult.error}`
+          );
+        }
+      }
     }
   }
 

library/agent/Config.ts

@@ -1,4 +1,5 @@
 import type { IPMatcher } from "../helpers/ip-matcher/IPMatcher";
+import type { WafRule } from "../waf/WafRule";
 
 export type EndpointConfig = {
   method: string;
@@ -31,4 +32,5 @@ export type Config = {
   block?: boolean;
   blockNewOutgoingRequests?: boolean;
   domains?: Domain[];
+  wafRules?: WafRule[];
 };

library/middleware/express.ts

@@ -27,6 +27,11 @@ export function addExpressMiddleware(app: Express | Router) {
         res.status(403).type("text").send("You are blocked by Zen.");
         return;
       }
+
+      if (result.type === "waf") {
+        res.status(403).type("text").send("You are blocked by Zen.");
+        return;
+      }
     }
 
     next();

library/middleware/hono.ts

@@ -29,6 +29,10 @@ export function addHonoMiddleware<
       if (result.type === "blocked") {
         return c.text("You are blocked by Zen.", 403);
       }
+
+      if (result.type === "waf") {
+        return c.text("You are blocked by Zen.", 403);
+      }
     }
 
     await next();

library/middleware/shouldBlockRequest.ts

@@ -3,11 +3,12 @@
 import { getInstance } from "../agent/AgentSingleton";
 import { getContext, updateContext } from "../agent/Context";
 import { shouldRateLimitRequest } from "../ratelimiting/shouldRateLimitRequest";
+import { evaluateWafRules } from "../waf/waf";
 
 type Result = {
   block: boolean;
-  type?: "ratelimited" | "blocked";
-  trigger?: "ip" | "user" | "group";
+  type?: "ratelimited" | "blocked" | "waf";
+  trigger?: "ip" | "user" | "group" | "waf";
   ip?: string;
 };
 
@@ -39,6 +40,17 @@ export function shouldBlockRequest(): Result {
     return { block: true, type: "blocked", trigger: "user" };
   }
 
+  // WAF rule evaluation
+  const wafResult = evaluateWafRules(context);
+  if (wafResult.matched && wafResult.action === "block") {
+    return {
+      block: true,
+      type: "waf",
+      trigger: "waf",
+      ip: context.remoteAddress,
+    };
+  }
+
   const rateLimitResult = shouldRateLimitRequest(context, agent);
   if (rateLimitResult.block) {
     // Mark the request as rate limited in the context

library/waf/WafRule.ts

@@ -0,0 +1,15 @@
+export type WafRuleAction = "block";
+
+export type WafRule = {
+  id: string;
+  expression: string;
+  action?: WafRuleAction;
+};
+
+export type WafSetRulesResult =
+  | { success: true }
+  | { success: false; error: string; rule_id?: string };
+
+export type WafEvaluateResult =
+  | { matched: false; error?: string }
+  | { matched: true; rule_id: string; action: string };

library/waf/waf.ts

@@ -0,0 +1,59 @@
+import {
+  wasm_waf_set_rules,
+  wasm_waf_evaluate,
+} from "../internals/zen_internals";
+import type { Context } from "../agent/Context";
+import type { WafRule, WafSetRulesResult, WafEvaluateResult } from "./WafRule";
+
+export function setWafRules(rules: WafRule[]): WafSetRulesResult {
+  return wasm_waf_set_rules(JSON.stringify(rules));
+}
+
+export function evaluateWafRules(context: Context): WafEvaluateResult {
+  if (!context.remoteAddress) {
+    return { matched: false, error: "Missing remote address" };
+  }
+
+  const url = context.url || "/";
+  let path = url;
+  let query = "";
+  const queryIndex = url.indexOf("?");
+  if (queryIndex !== -1) {
+    path = url.substring(0, queryIndex);
+    query = url.substring(queryIndex);
+  }
+
+  const host =
+    (typeof context.headers?.host === "string"
+      ? context.headers.host
+      : undefined) || "";
+
+  const requestData = {
+    host: host,
+    method: context.method || "GET",
+    path: path,
+    query: query,
+    uri: url,
+    full_uri: `${host}${url}`,
+    user_agent:
+      typeof context.headers?.["user-agent"] === "string"
+        ? context.headers["user-agent"]
+        : undefined,
+    cookie:
+      typeof context.headers?.cookie === "string"
+        ? context.headers.cookie
+        : undefined,
+    referer:
+      typeof context.headers?.referer === "string"
+        ? context.headers.referer
+        : undefined,
+    x_forwarded_for:
+      typeof context.headers?.["x-forwarded-for"] === "string"
+        ? context.headers["x-forwarded-for"]
+        : undefined,
+    body: typeof context.body === "string" ? context.body : undefined,
+    ip_src: context.remoteAddress,
+  };
+
+  return wasm_waf_evaluate(JSON.stringify(requestData));
+}