Merge pull request #447 from AikidoSec/support-ip-blocking

Add support for IP address blocking

end2end/server/app.js

@@ -1,11 +1,13 @@
-// This is a insecure mock server for testing purposes
+// This is an insecure mock server for testing purposes
 const express = require("express");
 const config = require("./src/handlers/getConfig");
 const captureEvent = require("./src/handlers/captureEvent");
 const listEvents = require("./src/handlers/listEvents");
 const createApp = require("./src/handlers/createApp");
 const checkToken = require("./src/middleware/checkToken");
 const updateConfig = require("./src/handlers/updateConfig");
+const ipLists = require("./src/handlers/ipLists");
+const updateIPLists = require("./src/handlers/updateIPLists");
 
 const app = express();
 
@@ -19,6 +21,9 @@ app.post("/api/runtime/config", checkToken, updateConfig);
 app.get("/api/runtime/events", checkToken, listEvents);
 app.post("/api/runtime/events", checkToken, captureEvent);
 
+app.get("/api/runtime/firewall/lists", checkToken, ipLists);
+app.post("/api/runtime/firewall/lists", checkToken, updateIPLists);
+
 app.post("/api/runtime/apps", createApp);
 
 app.listen(port, () => {

end2end/server/src/handlers/ipLists.js

@@ -0,0 +1,24 @@
+const { getBlockedIPAddresses } = require("../zen/config");
+
+module.exports = function ipLists(req, res) {
+  if (!req.app) {
+    throw new Error("App is missing");
+  }
+
+  const blockedIps = getBlockedIPAddresses(req.app);
+
+  res.json({
+    success: true,
+    serviceId: req.app.id,
+    blockedIPAddresses:
+      blockedIps.length > 0
+        ? [
+            {
+              source: "geoip",
+              description: "geo restrictions",
+              ips: blockedIps,
+            },
+          ]
+        : [],
+  });
+};

end2end/server/src/handlers/updateIPLists.js

@@ -0,0 +1,32 @@
+const { updateBlockedIPAddresses } = require("../zen/config");
+
+module.exports = function updateIPLists(req, res) {
+  if (!req.app) {
+    throw new Error("App is missing");
+  }
+
+  // Insecure input validation - but this is only a mock server
+  if (
+    !req.body ||
+    typeof req.body !== "object" ||
+    Array.isArray(req.body) ||
+    !Object.keys(req.body).length
+  ) {
+    return res.status(400).json({
+      message: "Request body is missing or invalid",
+    });
+  }
+
+  if (
+    !req.body.blockedIPAddresses ||
+    !Array.isArray(req.body.blockedIPAddresses)
+  ) {
+    return res.status(400).json({
+      message: "blockedIPAddresses is missing or invalid",
+    });
+  }
+
+  updateBlockedIPAddresses(req.app, req.body.blockedIPAddresses);
+
+  res.json({ success: true });
+};

end2end/server/src/zen/config.js

@@ -29,11 +29,43 @@ function updateAppConfig(app, newConfig) {
     getAppConfig(app);
     index = configs.length - 1;
   }
-  configs[index] = { ...configs[index], ...newConfig };
+  configs[index] = {
+    ...configs[index],
+    ...newConfig,
+    lastUpdatedAt: Date.now(),
+  };
   return true;
 }
 
+const blockedIPAddresses = [];
+
+function updateBlockedIPAddresses(app, ips) {
+  let entry = blockedIPAddresses.find((ip) => ip.serviceId === app.serviceId);
+
+  if (entry) {
+    entry.ipAddresses = ips;
+  } else {
+    entry = { serviceId: app.serviceId, ipAddresses: ips };
+    blockedIPAddresses.push(entry);
+  }
+
+  // Bump lastUpdatedAt
+  updateAppConfig(app, {});
+}
+
+function getBlockedIPAddresses(app) {
+  const entry = blockedIPAddresses.find((ip) => ip.serviceId === app.serviceId);
+
+  if (entry) {
+    return entry.ipAddresses;
+  }
+
+  return { serviceId: app.serviceId, ipAddresses: [] };
+}
+
 module.exports = {
   getAppConfig,
   updateAppConfig,
+  updateBlockedIPAddresses,
+  getBlockedIPAddresses,
 };

end2end/tests/hono-xml-geo-blocking.test.js

@@ -0,0 +1,114 @@
+const t = require("tap");
+const { spawn } = require("child_process");
+const { resolve } = require("path");
+const timeout = require("../timeout");
+
+const pathToApp = resolve(__dirname, "../../sample-apps/hono-xml", "app.js");
+const testServerUrl = "http://localhost:5874";
+
+let token;
+t.beforeEach(async () => {
+  const response = await fetch(`${testServerUrl}/api/runtime/apps`, {
+    method: "POST",
+  });
+  const body = await response.json();
+  token = body.token;
+
+  // Apply rate limiting
+  const updateConfigResponse = await fetch(
+    `${testServerUrl}/api/runtime/firewall/lists`,
+    {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: token,
+      },
+      body: JSON.stringify({
+        blockedIPAddresses: ["1.3.2.0/24", "fe80::1234:5678:abcd:ef12/64"],
+      }),
+    }
+  );
+  t.same(updateConfigResponse.status, 200);
+});
+
+t.test("it blocks geo restricted IPs", (t) => {
+  const server = spawn(`node`, ["--preserve-symlinks", pathToApp, "4002"], {
+    env: {
+      ...process.env,
+      AIKIDO_DEBUG: "true",
+      AIKIDO_BLOCKING: "true",
+      AIKIDO_TOKEN: token,
+      AIKIDO_URL: testServerUrl,
+    },
+  });
+
+  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(async () => {
+      const resp1 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Njuska</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "1.3.2.4",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp1.status, 403);
+      t.same(
+        await resp1.text(),
+        "Your IP address is blocked due to geo restrictions. (Your IP: 1.3.2.4)"
+      );
+
+      const resp2 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Harry</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "fe80::1234:5678:abcd:ef12",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp2.status, 403);
+      t.same(
+        await resp2.text(),
+        "Your IP address is blocked due to geo restrictions. (Your IP: fe80::1234:5678:abcd:ef12)"
+      );
+
+      const resp3 = await fetch("http://127.0.0.1:4002/add", {
+        method: "POST",
+        body: "<cat><name>Harry</name></cat>",
+        headers: {
+          "Content-Type": "application/xml",
+          "X-Forwarded-For": "9.8.7.6",
+        },
+        signal: AbortSignal.timeout(5000),
+      });
+      t.same(resp3.status, 200);
+      t.same(await resp3.text(), JSON.stringify({ success: true }));
+    })
+    .catch((error) => {
+      t.fail(error.message);
+    })
+    .finally(() => {
+      server.kill();
+    });
+});

library/agent/Agent.test.ts

@@ -1,8 +1,10 @@
 import * as FakeTimers from "@sinonjs/fake-timers";
 import { hostname, platform, release } from "os";
 import * as t from "tap";
+import * as fetch from "../helpers/fetch";
 import { getSemverNodeVersion } from "../helpers/getNodeVersion";
 import { ip } from "../helpers/ipAddress";
+import { wrap } from "../helpers/wrap";
 import { MongoDB } from "../sinks/MongoDB";
 import { Agent } from "./Agent";
 import { ReportingAPIForTesting } from "./api/ReportingAPIForTesting";
@@ -15,6 +17,24 @@ import { LoggerNoop } from "./logger/LoggerNoop";
 import { Wrapper } from "./Wrapper";
 import { Context } from "./Context";
 import { createTestAgent } from "../helpers/createTestAgent";
+import { setTimeout } from "node:timers/promises";
+
+wrap(fetch, "fetch", function mock() {
+  return async function mock() {
+    return {
+      statusCode: 200,
+      body: JSON.stringify({
+        blockedIPAddresses: [
+          {
+            source: "name",
+            description: "Description",
+            ips: ["1.3.2.0/24", "fe80::1234:5678:abcd:ef12/64"],
+          },
+        ],
+      }),
+    };
+  };
+});
 
 t.test("it throws error if serverless is empty string", async () => {
   t.throws(
@@ -40,7 +60,8 @@ t.test("it sends started event", async (t) => {
   });
   agent.start([new MongoDB()]);
 
-  const mongodb = require("mongodb");
+  // Require mongodb to see if agent logs message
+  require("mongodb");
 
   t.match(api.getEvents(), [
     {
@@ -496,11 +517,6 @@ t.test("it sends heartbeat when reached max timings", async () => {
 });
 
 t.test("it logs when failed to report event", async () => {
-  async function waitForCalls() {
-    // API calls are async, wait for them to finish
-    await new Promise((resolve) => setTimeout(resolve, 0));
-  }
-
   const logger = new LoggerForTesting();
   const api = new ReportingAPIThatThrows();
   const agent = createTestAgent({
@@ -510,12 +526,12 @@ t.test("it logs when failed to report event", async () => {
   });
   agent.start([]);
 
-  await waitForCalls();
+  await setTimeout(0);
 
   // @ts-expect-error Private method
   agent.heartbeat();
 
-  await waitForCalls();
+  await setTimeout(0);
 
   agent.onDetectedAttack({
     module: "mongodb",
@@ -545,7 +561,7 @@ t.test("it logs when failed to report event", async () => {
     },
   });
 
-  await waitForCalls();
+  await setTimeout(0);
 
   t.same(logger.getMessages(), [
     "Starting agent...",
@@ -804,7 +820,7 @@ t.test(
     agent.start([]);
 
     // Wait for the event to be sent
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    await setTimeout(0);
 
     t.same(agent.shouldBlock(), true);
   }
@@ -825,7 +841,7 @@ t.test(
     agent.start([]);
 
     // Wait for the event to be sent
-    await new Promise((resolve) => setTimeout(resolve, 0));
+    await setTimeout(0);
 
     t.same(agent.shouldBlock(), false);
   }
@@ -852,7 +868,7 @@ t.test("it enables blocking mode after sending startup event", async () => {
   agent.start([]);
 
   // Wait for the event to be sent
-  await new Promise((resolve) => setTimeout(resolve, 0));
+  await setTimeout(0);
 
   t.same(agent.shouldBlock(), true);
 });
@@ -877,7 +893,7 @@ t.test("it goes into monitoring mode after sending startup event", async () => {
   agent.start([]);
 
   // Wait for the event to be sent
-  await new Promise((resolve) => setTimeout(resolve, 0));
+  await setTimeout(0);
 
   t.same(agent.shouldBlock(), false);
 });
@@ -911,3 +927,37 @@ t.test("it sends middleware installed with heartbeat", async () => {
 
   clock.uninstall();
 });
+
+t.test("it fetches blocked IPs", async () => {
+  const agent = createTestAgent({
+    token: new Token("123"),
+  });
+
+  agent.start([]);
+
+  await setTimeout(0);
+
+  t.same(agent.getConfig().isIPAddressBlocked("1.3.2.4"), {
+    blocked: true,
+    reason: "Description",
+  });
+  t.same(agent.getConfig().isIPAddressBlocked("fe80::1234:5678:abcd:ef12"), {
+    blocked: true,
+    reason: "Description",
+  });
+});
+
+t.test("it does not fetch blocked IPs if serverless", async () => {
+  const agent = createTestAgent({
+    token: new Token("123"),
+    serverless: "gcp",
+  });
+
+  agent.start([]);
+
+  await setTimeout(0);
+
+  t.same(agent.getConfig().isIPAddressBlocked("1.3.2.4"), {
+    blocked: false,
+  });
+});