Add support for IP address blocking

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,
+  });
+});

library/agent/Agent.ts

@@ -7,6 +7,7 @@
 import { filterEmptyRequestHeaders } from "../helpers/filterEmptyRequestHeaders";
 import { limitLengthMetadata } from "../helpers/limitLengthMetadata";
 import { RateLimiter } from "../ratelimiting/RateLimiter";
+import { fetchBlockedIPAddresses } from "./api/fetchBlockedIPAddresses";
 import { ReportingAPI, ReportingAPIResponse } from "./api/ReportingAPI";
 import { AgentInfo } from "./api/Event";
 import { Token } from "./api/Token";
@@ -39,7 +40,7 @@
   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({
@@ -111,6 +112,8 @@
       );
 
       this.updateServiceConfig(result);
+
+      await this.updateBlockedIPAddresses();
     }
   }
 
@@ -225,7 +228,7 @@
       }
 
       if (response.endpoints) {
-        this.serviceConfig = new ServiceConfig(
+        this.serviceConfig.updateConfig(
           response.endpoints && Array.isArray(response.endpoints)
             ? response.endpoints
             : [],
@@ -236,7 +239,7 @@
             ? response.blockedUserIds
             : [],
           response.allowedIPAddresses &&
-          Array.isArray(response.allowedIPAddresses)
+            Array.isArray(response.allowedIPAddresses)
             ? response.allowedIPAddresses
             : [],
           typeof response.receivedAnyStats === "boolean"
@@ -287,6 +290,7 @@
         },
         timeoutInMS
       );
+
       this.updateServiceConfig(response);
     }
   }
@@ -333,6 +337,26 @@
     this.interval.unref();
   }
 
+  private async updateBlockedIPAddresses() {
+    if (!this.token) {
+      return;
+    }
+
+    if (this.serverless) {
+      // Not supported in serverless mode
+      return;
+    }
+
+    try {
+      const blockedIps = await fetchBlockedIPAddresses(this.token);
+      this.serviceConfig.updateBlockedIPAddresses(blockedIps);
+    } catch (error: any) {
+      this.logger.log(
+        `Failed to update blocked IP addresses: ${error.message}`
+      );
+    }
+  }
+
   private startPollingForConfigChanges() {
     pollForChanges({
       token: this.token,
@@ -341,6 +365,7 @@
       lastUpdatedAt: this.serviceConfig.getLastUpdatedAt(),
       onConfigUpdate: (config) => {
         this.updateServiceConfig({ success: true, ...config });
+        this.updateBlockedIPAddresses().catch(() => {});
       },
     });
   }

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);
@@ -53,7 +53,8 @@ t.test("it works", async () => {
     0,
     ["123"],
     [],
-    false
+    false,
+    []
   );
 
   t.same(config.isUserBlocked("123"), true);
@@ -80,7 +81,52 @@ 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",
+      ],
+    },
+  ]);
+  t.same(config.isIPAddressBlocked("1.2.3.4"), {
+    blocked: true,
+    reason: "description",
+  });
+  t.same(config.isIPAddressBlocked("2.3.4.5"), { blocked: false });
+  t.same(config.isIPAddressBlocked("192.168.2.2"), {
+    blocked: true,
+    reason: "description",
+  });
+  t.same(config.isIPAddressBlocked("fd00:1234:5678:9abc::1"), {
+    blocked: true,
+    reason: "description",
+  });
+  t.same(config.isIPAddressBlocked("fd00:1234:5678:9abc::2"), {
+    blocked: false,
+  });
+  t.same(config.isIPAddressBlocked("fd00:3234:5678:9abc::1"), {
+    blocked: true,
+    reason: "description",
+  });
+  t.same(config.isIPAddressBlocked("fd00:3234:5678:9abc::2"), {
+    blocked: true,
+    reason: "description",
+  });
+  t.same(config.isIPAddressBlocked("5.6.7.8"), {
+    blocked: true,
+    reason: "description",
+  });
+  t.same(config.isIPAddressBlocked("1.2"), { blocked: false });
+});