Update to new API response

library/agent/Agent.test.ts

@@ -17,13 +17,20 @@ 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: ["1.3.2.4", "fe80::1234:5678:abcd:ef12/64"],
+        blockedIPAddresses: [
+          {
+            source: "name",
+            description: "Description",
+            ips: ["1.3.2.0/24", "fe80::1234:5678:abcd:ef12/64"],
+          },
+        ],
       }),
     };
   };
@@ -510,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({
@@ -524,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",
@@ -559,7 +561,7 @@ t.test("it logs when failed to report event", async () => {
     },
   });
 
-  await waitForCalls();
+  await setTimeout(0);
 
   t.same(logger.getMessages(), [
     "Starting agent...",
@@ -818,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);
   }
@@ -839,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);
   }
@@ -866,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);
 });
@@ -891,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);
 });
@@ -925,3 +927,37 @@ t.test("it sends middleware installed with heartbeat", async () => {
 
   clock.uninstall();
 });
+
+t.test("it fetched block 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

@@ -342,6 +342,11 @@ export class Agent {
       return;
     }
 
+    if (this.serverless) {
+      // Not supported in serverless mode
+      return;
+    }
+
     try {
       const blockedIps = await fetchBlockedIPAddresses(this.token);
       this.serviceConfig.updateBlockedIPAddresses(blockedIps);

library/agent/ServiceConfig.test.ts

@@ -88,19 +88,45 @@ t.test("it checks if IP is allowed", async () => {
 
 t.test("ip blocking works", async () => {
   const config = new ServiceConfig([], 0, [], [], false, [
-    "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",
+    {
+      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"), true);
-  t.same(config.isIPAddressBlocked("2.3.4.5"), false);
-  t.same(config.isIPAddressBlocked("192.168.2.2"), true);
-  t.same(config.isIPAddressBlocked("fd00:1234:5678:9abc::1"), true);
-  t.same(config.isIPAddressBlocked("fd00:1234:5678:9abc::2"), false);
-  t.same(config.isIPAddressBlocked("fd00:3234:5678:9abc::1"), true);
-  t.same(config.isIPAddressBlocked("fd00:3234:5678:9abc::2"), true);
-  t.same(config.isIPAddressBlocked("5.6.7.8"), true);
-  t.same(config.isIPAddressBlocked("1.2"), false);
+  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 });
 });

library/agent/ServiceConfig.ts

@@ -2,21 +2,23 @@ import { BlockList, isIPv4, isIPv6 } from "net";
 import { LimitedContext, matchEndpoints } from "../helpers/matchEndpoints";
 import { Endpoint } from "./Config";
 import { addIPAddressToBlocklist } from "../helpers/addIPAddressToBlocklist";
+import { Blocklist as BlocklistType } from "./api/fetchBlockedIPAddresses";
 
 export class ServiceConfig {
   private blockedUserIds: Map<string, string> = new Map();
   private allowedIPAddresses: Map<string, string> = new Map();
   private nonGraphQLEndpoints: Endpoint[] = [];
   private graphqlFields: Endpoint[] = [];
-  private blockedIPAddresses = new BlockList();
+  private blockedIPAddresses: { blocklist: BlockList; description: string }[] =
+    [];
 
   constructor(
     endpoints: Endpoint[],
     private lastUpdatedAt: number,
     blockedUserIds: string[],
     allowedIPAddresses: string[],
     private receivedAnyStats: boolean,
-    blockedIPAddresses: string[]
+    blockedIPAddresses: BlocklistType[]
   ) {
     this.setBlockedUserIds(blockedUserIds);
     this.setAllowedIPAddresses(allowedIPAddresses);
@@ -80,24 +82,48 @@ export class ServiceConfig {
     return this.blockedUserIds.has(userId);
   }
 
-  isIPAddressBlocked(ip: string) {
+  isIPAddressBlocked(
+    ip: string
+  ): { blocked: true; reason: string } | { blocked: false } {
+    let blocklist: { blocklist: BlockList; description: string } | undefined =
+      undefined;
+
     if (isIPv4(ip)) {
-      return this.blockedIPAddresses.check(ip, "ipv4");
+      blocklist = this.blockedIPAddresses.find((blocklist) =>
+        blocklist.blocklist.check(ip, "ipv4")
+      );
     }
+
     if (isIPv6(ip)) {
-      return this.blockedIPAddresses.check(ip, "ipv6");
+      blocklist = this.blockedIPAddresses.find((blocklist) =>
+        blocklist.blocklist.check(ip, "ipv6")
+      );
     }
-    return false;
+
+    if (blocklist) {
+      return { blocked: true, reason: blocklist.description };
+    }
+
+    return { blocked: false };
   }
 
-  private setBlockedIPAddresses(blockedIPAddresses: string[]) {
-    this.blockedIPAddresses = new BlockList();
-    blockedIPAddresses.forEach((ip) => {
-      addIPAddressToBlocklist(ip, this.blockedIPAddresses);
-    });
+  private setBlockedIPAddresses(blockedIPAddresses: BlocklistType[]) {
+    this.blockedIPAddresses = [];
+
+    for (const source of blockedIPAddresses) {
+      const blocklist = new BlockList();
+      for (const ip of source.ips) {
+        addIPAddressToBlocklist(ip, blocklist);
+      }
+
+      this.blockedIPAddresses.push({
+        blocklist,
+        description: source.description,
+      });
+    }
   }
 
-  updateBlockedIPAddresses(blockedIPAddresses: string[]) {
+  updateBlockedIPAddresses(blockedIPAddresses: BlocklistType[]) {
     this.setBlockedIPAddresses(blockedIPAddresses);
   }
 

library/agent/api/fetchBlockedIPAddresses.ts

@@ -2,7 +2,15 @@ import { fetch } from "../../helpers/fetch";
 import { getAPIURL } from "../getAPIURL";
 import { Token } from "./Token";
 
-export async function fetchBlockedIPAddresses(token: Token): Promise<string[]> {
+export type Blocklist = {
+  source: string;
+  description: string;
+  ips: string[];
+};
+
+export async function fetchBlockedIPAddresses(
+  token: Token
+): Promise<Blocklist[]> {
   const baseUrl = getAPIURL();
   const { body, statusCode } = await fetch({
     url: new URL(`${baseUrl.toString()}api/runtime/firewall/lists`),
@@ -19,7 +27,9 @@ export async function fetchBlockedIPAddresses(token: Token): Promise<string[]> {
     throw new Error(`Failed to fetch blocked IP addresses: ${statusCode}`);
   }
 
-  const result: { blockedIPAddresses: string[] } = JSON.parse(body);
+  const result: {
+    blockedIPAddresses: Blocklist[];
+  } = JSON.parse(body);
 
   return result && result.blockedIPAddresses ? result.blockedIPAddresses : [];
 }

library/middleware/shouldBlockRequest.ts

@@ -2,12 +2,14 @@ import { getInstance } from "../agent/AgentSingleton";
 import { getContext, updateContext } from "../agent/Context";
 import { shouldRateLimitRequest } from "../ratelimiting/shouldRateLimitRequest";
 
-export function shouldBlockRequest(): {
+type Result = {
   block: boolean;
   type?: "ratelimited" | "blocked";
   trigger?: "ip" | "user";
   ip?: string;
-} {
+};
+
+export function shouldBlockRequest(): Result {
   const context = getContext();
   if (!context) {
     return { block: false };
@@ -25,18 +27,6 @@ export function shouldBlockRequest(): {
     return { block: true, type: "blocked", trigger: "user" };
   }
 
-  if (
-    context.remoteAddress &&
-    agent.getConfig().isIPAddressBlocked(context.remoteAddress)
-  ) {
-    return {
-      block: true,
-      type: "blocked",
-      trigger: "ip",
-      ip: context.remoteAddress,
-    };
-  }
-
   const rateLimitResult = shouldRateLimitRequest(context, agent);
   if (rateLimitResult.block) {
     return {