Implement ip allowlist (e.g. geo based)

timokoessler · timokoessler · commit b3a658d8cd6e · 2025-01-27T17:20:31.000+01:00
diff --git a/library/agent/Agent.test.ts b/library/agent/Agent.test.ts
@@ -19,6 +19,8 @@ import { Context } from "./Context";
 import { createTestAgent } from "../helpers/createTestAgent";
 import { setTimeout } from "node:timers/promises";
 
+let shouldOnlyAllowSomeIPAddresses = false;
+
 wrap(fetch, "fetch", function mock() {
   return async function mock() {
     return {
@@ -32,6 +34,15 @@ wrap(fetch, "fetch", function mock() {
           },
         ],
         blockedUserAgents: "AI2Bot|Bytespider",
+        allowedIPAddresses: shouldOnlyAllowSomeIPAddresses
+          ? [
+              {
+                source: "name",
+                description: "Description",
+                ips: ["4.3.2.1"],
+              },
+            ]
+          : [],
       }),
     };
   };
@@ -1099,6 +1110,8 @@ t.test("it does not fetch blocked IPs if serverless", async () => {
     blocked: false,
   });
 
+  t.same(agent.getConfig().shouldOnlyAllowSomeIPAddresses(), false);
+
   t.same(
     agent
       .getConfig()
@@ -1110,3 +1123,28 @@ t.test("it does not fetch blocked IPs if serverless", async () => {
     }
   );
 });
+
+t.test("it only allows some IP addresses", async () => {
+  shouldOnlyAllowSomeIPAddresses = true;
+  const agent = createTestAgent({
+    token: new Token("123"),
+    suppressConsoleLog: false,
+  });
+
+  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.same(agent.getConfig().shouldOnlyAllowSomeIPAddresses(), true);
+  t.same(agent.getConfig().isOnlyAllowedIPAddress("1.2.3.4"), false);
+  t.same(agent.getConfig().isOnlyAllowedIPAddress("4.3.2.1"), true);
+});
diff --git a/library/agent/Agent.ts b/library/agent/Agent.ts
@@ -41,7 +41,15 @@ export class Agent {
   private timeoutInMS = 10000;
   private hostnames = new Hostnames(200);
   private users = new Users(1000);
-  private serviceConfig = new ServiceConfig([], Date.now(), [], [], true, []);
+  private serviceConfig = new ServiceConfig(
+    [],
+    Date.now(),
+    [],
+    [],
+    true,
+    [],
+    []
+  );
   private routes: Routes = new Routes(200);
   private rateLimiter: RateLimiter = new RateLimiter(5000, 120 * 60 * 1000);
   private statistics = new InspectionStatistics({
@@ -352,7 +360,7 @@ export class Agent {
     this.interval.unref();
   }
 
-  private async updateBlockedLists() {
+  async updateBlockedLists() {
     if (!this.token) {
       return;
     }
@@ -363,11 +371,11 @@ export class Agent {
     }
 
     try {
-      const { blockedIPAddresses, blockedUserAgents } = await fetchBlockedLists(
-        this.token
-      );
+      const { blockedIPAddresses, blockedUserAgents, allowedIPAddresses } =
+        await fetchBlockedLists(this.token);
       this.serviceConfig.updateBlockedIPAddresses(blockedIPAddresses);
       this.serviceConfig.updateBlockedUserAgents(blockedUserAgents);
+      this.serviceConfig.updateOnlyAllowedIPAddresses(allowedIPAddresses);
     } catch (error: any) {
       console.error(`Aikido: Failed to update blocked lists: ${error.message}`);
     }
diff --git a/library/agent/ServiceConfig.test.ts b/library/agent/ServiceConfig.test.ts
@@ -2,7 +2,7 @@ import * as t from "tap";
 import { ServiceConfig } from "./ServiceConfig";
 
 t.test("it returns false if empty rules", async () => {
-  const config = new ServiceConfig([], 0, [], [], false, []);
+  const config = new ServiceConfig([], 0, [], [], false, [], []);
   t.same(config.getLastUpdatedAt(), 0);
   t.same(config.isUserBlocked("id"), false);
   t.same(config.isAllowedIP("1.2.3.4"), false);
@@ -54,6 +54,7 @@ t.test("it works", async () => {
     ["123"],
     [],
     false,
+    [],
     []
   );
 
@@ -81,25 +82,33 @@ t.test("it works", async () => {
 });
 
 t.test("it checks if IP is allowed", async () => {
-  const config = new ServiceConfig([], 0, [], ["1.2.3.4"], false, []);
+  const config = new ServiceConfig([], 0, [], ["1.2.3.4"], false, [], []);
   t.same(config.isAllowedIP("1.2.3.4"), true);
   t.same(config.isAllowedIP("1.2.3.5"), false);
 });
 
 t.test("ip blocking works", async () => {
-  const config = new ServiceConfig([], 0, [], [], false, [
-    {
-      source: "geoip",
-      description: "description",
-      ips: [
-        "1.2.3.4",
-        "192.168.2.1/24",
-        "fd00:1234:5678:9abc::1",
-        "fd00:3234:5678:9abc::1/64",
-        "5.6.7.8/32",
-      ],
-    },
-  ]);
+  const config = new ServiceConfig(
+    [],
+    0,
+    [],
+    [],
+    false,
+    [
+      {
+        source: "geoip",
+        description: "description",
+        ips: [
+          "1.2.3.4",
+          "192.168.2.1/24",
+          "fd00:1234:5678:9abc::1",
+          "fd00:3234:5678:9abc::1/64",
+          "5.6.7.8/32",
+        ],
+      },
+    ],
+    []
+  );
   t.same(config.isIPAddressBlocked("1.2.3.4"), {
     blocked: true,
     reason: "description",
@@ -132,14 +141,42 @@ t.test("ip blocking works", async () => {
 });
 
 t.test("it blocks bots", async () => {
-  const config = new ServiceConfig([], 0, [], [], true, []);
+  const config = new ServiceConfig([], 0, [], [], true, [], []);
   config.updateBlockedUserAgents("googlebot|bingbot");
 
   t.same(config.isUserAgentBlocked("googlebot"), { blocked: true });
   t.same(config.isUserAgentBlocked("123 bingbot abc"), { blocked: true });
   t.same(config.isUserAgentBlocked("bing"), { blocked: false });
 
+  t.same(config.shouldOnlyAllowSomeIPAddresses(), false);
+
   config.updateBlockedUserAgents("");
 
   t.same(config.isUserAgentBlocked("googlebot"), { blocked: false });
 });
+
+t.test("restricting access to some ips", async () => {
+  const config = new ServiceConfig(
+    [],
+    0,
+    [],
+    [],
+    true,
+    [],
+    [
+      {
+        source: "geoip",
+        description: "description",
+        ips: ["1.2.3.4"],
+      },
+    ]
+  );
+
+  t.same(config.shouldOnlyAllowSomeIPAddresses(), true);
+
+  t.same(config.isOnlyAllowedIPAddress("1.2.3.4"), true);
+  t.same(config.isOnlyAllowedIPAddress("4.3.2.1"), false);
+
+  config.updateOnlyAllowedIPAddresses([]);
+  t.same(config.isOnlyAllowedIPAddress("1.2.3.4"), false);
+});
diff --git a/library/agent/ServiceConfig.ts b/library/agent/ServiceConfig.ts
@@ -1,7 +1,7 @@
 import { IPMatcher } from "../helpers/ip-matcher/IPMatcher";
 import { LimitedContext, matchEndpoints } from "../helpers/matchEndpoints";
 import { Endpoint } from "./Config";
-import { Blocklist as BlocklistType } from "./api/fetchBlockedLists";
+import { IPList as IPList } from "./api/fetchBlockedLists";
 
 export class ServiceConfig {
   private blockedUserIds: Map<string, string> = new Map();
@@ -11,19 +11,22 @@ export class ServiceConfig {
   private blockedIPAddresses: { blocklist: IPMatcher; description: string }[] =
     [];
   private blockedUserAgentRegex: RegExp | undefined;
+  private onlyAllowedIPAddresses: IPMatcher | undefined;
 
   constructor(
     endpoints: Endpoint[],
     private lastUpdatedAt: number,
     blockedUserIds: string[],
     allowedIPAddresses: string[],
     private receivedAnyStats: boolean,
-    blockedIPAddresses: BlocklistType[]
+    blockedIPAddresses: IPList[],
+    onlyAllowedIPAddresses: IPList[]
   ) {
     this.setBlockedUserIds(blockedUserIds);
     this.setAllowedIPAddresses(allowedIPAddresses);
     this.setEndpoints(endpoints);
     this.setBlockedIPAddresses(blockedIPAddresses);
+    this.setOnlyAllowedIPAddresses(onlyAllowedIPAddresses);
   }
 
   private setEndpoints(endpoints: Endpoint[]) {
@@ -96,7 +99,7 @@ export class ServiceConfig {
     return { blocked: false };
   }
 
-  private setBlockedIPAddresses(blockedIPAddresses: BlocklistType[]) {
+  private setBlockedIPAddresses(blockedIPAddresses: IPList[]) {
     this.blockedIPAddresses = [];
 
     for (const source of blockedIPAddresses) {
@@ -107,7 +110,7 @@ export class ServiceConfig {
     }
   }
 
-  updateBlockedIPAddresses(blockedIPAddresses: BlocklistType[]) {
+  updateBlockedIPAddresses(blockedIPAddresses: IPList[]) {
     this.setBlockedIPAddresses(blockedIPAddresses);
   }
 
@@ -126,6 +129,32 @@ export class ServiceConfig {
     return { blocked: false };
   }
 
+  private setOnlyAllowedIPAddresses(ipAddresses: IPList[]) {
+    this.onlyAllowedIPAddresses = undefined;
+
+    if (ipAddresses.length === 0) {
+      return;
+    }
+
+    const ips = ipAddresses.map((source) => source.ips).flat();
+
+    this.onlyAllowedIPAddresses = new IPMatcher(ips);
+  }
+
+  updateOnlyAllowedIPAddresses(ipAddresses: IPList[]) {
+    this.setOnlyAllowedIPAddresses(ipAddresses);
+  }
+
+  shouldOnlyAllowSomeIPAddresses() {
+    return this.onlyAllowedIPAddresses !== undefined;
+  }
+
+  isOnlyAllowedIPAddress(ip: string) {
+    return this.onlyAllowedIPAddresses
+      ? this.onlyAllowedIPAddresses.has(ip)
+      : false;
+  }
+
   updateConfig(
     endpoints: Endpoint[],
     lastUpdatedAt: number,
diff --git a/library/agent/api/fetchBlockedLists.ts b/library/agent/api/fetchBlockedLists.ts
@@ -2,14 +2,15 @@ import { fetch } from "../../helpers/fetch";
 import { getAPIURL } from "../getAPIURL";
 import { Token } from "./Token";
 
-export type Blocklist = {
+export type IPList = {
   source: string;
   description: string;
   ips: string[];
 };
 
 export async function fetchBlockedLists(token: Token): Promise<{
-  blockedIPAddresses: Blocklist[];
+  blockedIPAddresses: IPList[];
+  allowedIPAddresses: IPList[];
   blockedUserAgents: string;
 }> {
   const baseUrl = getAPIURL();
@@ -34,7 +35,8 @@ export async function fetchBlockedLists(token: Token): Promise<{
   }
 
   const result: {
-    blockedIPAddresses: Blocklist[];
+    blockedIPAddresses: IPList[];
+    allowedIPAddresses: IPList[];
     blockedUserAgents: string;
   } = JSON.parse(body);
 
@@ -43,6 +45,10 @@ export async function fetchBlockedLists(token: Token): Promise<{
       result && Array.isArray(result.blockedIPAddresses)
         ? result.blockedIPAddresses
         : [],
+    allowedIPAddresses:
+      result && Array.isArray(result.allowedIPAddresses)
+        ? result.allowedIPAddresses
+        : [],
     // Blocked user agents are stored as a string pattern for usage in a regex (e.g. "Googlebot|Bingbot")
     blockedUserAgents:
       result && typeof result.blockedUserAgents === "string"
diff --git a/library/ratelimiting/getRateLimitedEndpoint.test.ts b/library/ratelimiting/getRateLimitedEndpoint.test.ts
@@ -18,7 +18,10 @@ const context: Context = {
 
 t.test("it returns undefined if no endpoints", async () => {
   t.same(
-    getRateLimitedEndpoint(context, new ServiceConfig([], 0, [], [], true, [])),
+    getRateLimitedEndpoint(
+      context,
+      new ServiceConfig([], 0, [], [], true, [], [])
+    ),
     undefined
   );
 });
@@ -45,6 +48,7 @@ t.test("it returns undefined if no matching endpoints", async () => {
         [],
         [],
         false,
+        [],
         []
       )
     ),
@@ -74,6 +78,7 @@ t.test("it returns undefined if matching but not enabled", async () => {
         [],
         [],
         false,
+        [],
         []
       )
     ),
@@ -103,6 +108,7 @@ t.test("it returns endpoint if matching and enabled", async () => {
         [],
         [],
         false,
+        [],
         []
       )
     ),
@@ -153,6 +159,7 @@ t.test("it returns endpoint with lowest max requests", async () => {
         [],
         [],
         false,
+        [],
         []
       )
     ),
@@ -203,6 +210,7 @@ t.test("it returns endpoint with smallest window size", async () => {
         [],
         [],
         false,
+        [],
         []
       )
     ),
@@ -253,6 +261,7 @@ t.test("it always returns exact matches first", async () => {
         [],
         [],
         false,
+        [],
         []
       )
     ),
diff --git a/library/sources/HTTPServer.test.ts b/library/sources/HTTPServer.test.ts
diff --git a/library/sources/Hono.test.ts b/library/sources/Hono.test.ts
diff --git a/library/sources/http-server/checkIfRequestIsBlocked.ts b/library/sources/http-server/checkIfRequestIsBlocked.ts
