Merge pull request #220 from AikidoSec/allow-list-ip

Don't block IPs from allow list

library/agent/Agent.ts

@@ -37,7 +37,7 @@ export class Agent {
   private timeoutInMS = 5000;
   private hostnames = new Hostnames(200);
   private users = new Users(1000);
-  private serviceConfig = new ServiceConfig([], Date.now(), []);
+  private serviceConfig = new ServiceConfig([], Date.now(), [], []);
   private routes: Routes = new Routes(200);
   private rateLimiter: RateLimiter = new RateLimiter(5000, 120 * 60 * 1000);
   private statistics = new InspectionStatistics({
@@ -218,7 +218,8 @@ export class Agent {
           typeof response.configUpdatedAt === "number"
             ? response.configUpdatedAt
             : Date.now(),
-          response.blockedUserIds ? response.blockedUserIds : []
+          response.blockedUserIds ? response.blockedUserIds : [],
+          response.allowedIPAddresses ? response.allowedIPAddresses : []
         );
       }
 

library/agent/Config.ts

@@ -14,4 +14,5 @@ export type Config = {
   heartbeatIntervalInMS: number;
   configUpdatedAt: number;
   blockedUserIds: string[];
+  allowedIPAddresses: string[];
 };

library/agent/ServiceConfig.test.ts

@@ -2,10 +2,11 @@ import * as t from "tap";
 import { ServiceConfig } from "./ServiceConfig";
 
 t.test("it returns false if empty rules", async () => {
-  const config = new ServiceConfig([], 0, []);
+  const config = new ServiceConfig([], 0, [], []);
   t.same(config.shouldProtectEndpoint("GET", "/foo"), true);
   t.same(config.getLastUpdatedAt(), 0);
   t.same(config.isUserBlocked("id"), false);
+  t.same(config.isAllowedIP("1.2.3.4"), false);
 });
 
 t.test("it works", async () => {
@@ -43,7 +44,8 @@ t.test("it works", async () => {
       },
     ],
     0,
-    ["123"]
+    ["123"],
+    []
   );
 
   t.same(config.shouldProtectEndpoint("GET", "/foo"), true);
@@ -75,6 +77,7 @@ t.test("it returns rate limiting", async () => {
       },
     ],
     0,
+    [],
     []
   );
 
@@ -91,3 +94,9 @@ t.test("it returns rate limiting", async () => {
     windowSizeInMS: 0,
   });
 });
+
+t.test("it checks if IP is allowed", async () => {
+  const config = new ServiceConfig([], 0, [], ["1.2.3.4"]);
+  t.same(config.isAllowedIP("1.2.3.4"), true);
+  t.same(config.isAllowedIP("1.2.3.5"), false);
+});

library/agent/ServiceConfig.ts

@@ -3,11 +3,13 @@ import { Endpoint } from "./Config";
 export class ServiceConfig {
   private endpoints: Map<string, Endpoint> = new Map();
   private blockedUserIds: Map<string, string> = new Map();
+  private allowedIPAddresses: Map<string, string> = new Map();
 
   constructor(
     endpoints: Endpoint[],
     private readonly lastUpdatedAt: number,
-    blockedUserIds: string[]
+    blockedUserIds: string[],
+    allowedIPAddresses: string[]
   ) {
     endpoints.forEach((rule) => {
       this.endpoints.set(this.getKey(rule.method, rule.route), {
@@ -21,6 +23,10 @@ export class ServiceConfig {
     blockedUserIds.forEach((userId) => {
       this.blockedUserIds.set(userId, userId);
     });
+
+    allowedIPAddresses.forEach((ip) => {
+      this.allowedIPAddresses.set(ip, ip);
+    });
   }
 
   private getKey(method: string, route: string) {
@@ -42,6 +48,10 @@ export class ServiceConfig {
     return rule.rateLimiting;
   }
 
+  isAllowedIP(ip: string) {
+    return this.allowedIPAddresses.has(ip);
+  }
+
   shouldProtectEndpoint(method: string, route: string | RegExp) {
     const key = this.getKey(
       method,

library/agent/api/ReportingAPIForTesting.ts

@@ -12,6 +12,7 @@ export class ReportingAPIForTesting implements ReportingAPI {
       configUpdatedAt: 0,
       heartbeatIntervalInMS: 10 * 60 * 1000,
       blockedUserIds: [],
+      allowedIPAddresses: [],
     }
   ) {}
 

library/agent/api/ReportingAPIRateLimitedClientSide.test.ts

@@ -28,6 +28,7 @@ function generateAttackEvent(): Event {
       metadata: {},
       operation: "operation",
       payload: "payload",
+      user: undefined,
     },
     agent: {
       version: "1.0.0",
@@ -137,6 +138,7 @@ function generateHeartbeatEvent(): Event {
       sinks: {},
       requests: {
         total: 0,
+        aborted: 0,
         attacksDetected: {
           blocked: 0,
           total: 0,
@@ -206,6 +208,7 @@ t.test("it does not blow memory", async () => {
       configUpdatedAt: 0,
       heartbeatIntervalInMS: 10 * 60 * 1000,
       blockedUserIds: [],
+      allowedIPAddresses: [],
     });
   }
 

library/agent/api/ReportingAPIRateLimitedServerSide.test.ts

@@ -43,6 +43,7 @@ t.test("it stops sending requests if rate limited", async (t) => {
     heartbeatIntervalInMS: 10 * 60 * 1000,
     configUpdatedAt: 0,
     blockedUserIds: [],
+    allowedIPAddresses: [],
   });
   t.match(api.getEvents(), [{ type: "started" }]);
 
@@ -93,13 +94,15 @@ t.test("it stops sending requests if rate limited", async (t) => {
     heartbeatIntervalInMS: 10 * 60 * 1000,
     configUpdatedAt: 0,
     blockedUserIds: [],
+    allowedIPAddresses: [],
   });
   t.same(await rateLimitedAPI.report(token, generateStartedEvent(), 5000), {
     success: true,
     endpoints: [],
     heartbeatIntervalInMS: 10 * 60 * 1000,
     configUpdatedAt: 0,
     blockedUserIds: [],
+    allowedIPAddresses: [],
   });
   t.match(api.getEvents(), [
     { type: "started" },

library/agent/api/ReportingAPIThatValidatesToken.test.ts

@@ -41,6 +41,7 @@ t.test("it ignores valid tokens", async () => {
     configUpdatedAt: 0,
     heartbeatIntervalInMS: 10 * 60 * 1000,
     blockedUserIds: [],
+    allowedIPAddresses: [],
   });
   t.same(api.getEvents(), [event]);
 
@@ -50,6 +51,7 @@ t.test("it ignores valid tokens", async () => {
     configUpdatedAt: 0,
     heartbeatIntervalInMS: 10 * 60 * 1000,
     blockedUserIds: [],
+    allowedIPAddresses: [],
   });
   t.same(api.getEvents(), [event, event]);
 });

library/agent/applyHooks.test.ts

@@ -246,7 +246,18 @@ t.test("it ignores route if force protection off is on", async (t) => {
 
   api.setResult({
     success: true,
-    endpoints: [{ method: "GET", route: "/route", forceProtectionOff: true }],
+    endpoints: [
+      {
+        method: "GET",
+        route: "/route",
+        forceProtectionOff: true,
+        rateLimiting: undefined,
+      },
+    ],
+    heartbeatIntervalInMS: 10 * 60 * 1000,
+    blockedUserIds: [],
+    allowedIPAddresses: [],
+    configUpdatedAt: 0,
   });
 
   // Read rules from API
@@ -282,3 +293,45 @@ t.test("it ignores route if force protection off is on", async (t) => {
     { args: ["www.aikido.dev"] },
   ]);
 });
+
+t.test("it does not report attack if IP is allowed", async (t) => {
+  const hooks = new Hooks();
+  hooks
+    .addBuiltinModule("os")
+    .addSubject((exports) => exports)
+    .inspect("hostname", (args, subject, agent) => {
+      return {
+        operation: "os.hostname",
+        source: "body",
+        pathToPayload: "path",
+        payload: "payload",
+        metadata: {},
+        kind: "path_traversal",
+      };
+    });
+
+  const { agent, api } = createAgent();
+  applyHooks(hooks, agent);
+
+  api.setResult({
+    success: true,
+    endpoints: [],
+    configUpdatedAt: 0,
+    heartbeatIntervalInMS: 10 * 60 * 1000,
+    blockedUserIds: [],
+    allowedIPAddresses: ["::1"],
+  });
+
+  // Read rules from API
+  await agent.flushStats(1000);
+  api.clear();
+
+  const { hostname } = require("os");
+
+  await runWithContext(context, async () => {
+    const name = hostname();
+    t.ok(typeof name === "string");
+  });
+
+  t.same(api.getEvents(), []);
+});

library/agent/applyHooks.ts

@@ -200,7 +200,12 @@ function wrapWithoutArgumentModification(
           withoutContext: !context,
         });
 
-        if (result && context) {
+        const isAllowedIP =
+          context &&
+          context.remoteAddress &&
+          agent.getConfig().isAllowedIP(context.remoteAddress);
+
+        if (result && context && !isAllowedIP) {
           // Flag request as having an attack detected
           context.attackDetected = true;
 

library/sources/Express.test.ts

@@ -46,10 +46,21 @@ const agent = new Agent(
           enabled: true,
         },
       },
+      {
+        method: "GET",
+        route: "/white-listed-ip-address",
+        forceProtectionOff: false,
+        rateLimiting: {
+          windowSizeInMS: 2000,
+          maxRequests: 3,
+          enabled: true,
+        },
+      },
     ],
     blockedUserIds: ["567"],
     configUpdatedAt: 0,
     heartbeatIntervalInMS: 10 * 60 * 1000,
+    allowedIPAddresses: ["4.3.2.1"],
   }),
   new Token("123"),
   "lambda"
@@ -189,6 +200,10 @@ function getApp(userMiddleware = true) {
     res.send({ hello: "world" });
   });
 
+  app.get("/white-listed-ip-address", (req, res) => {
+    res.send({ hello: "world" });
+  });
+
   app.use("/middleware-rate-limited", (req, res, next) => {
     res.send({ hello: "world" });
   });
@@ -429,3 +444,12 @@ t.test("it rate limits by middleware", async () => {
   const res2 = await request(getApp()).get("/middleware-rate-limited");
   t.same(res2.statusCode, 200);
 });
+
+t.test("it allows white-listed IP address", async () => {
+  for (const _ of Array.from({ length: 5 })) {
+    const res = await request(getApp(false))
+      .get("/white-listed-ip-address")
+      .set("x-forwarded-for", "4.3.2.1");
+    t.same(res.statusCode, 200);
+  }
+});

library/sources/express/shouldRateLimitRequest.ts

@@ -28,7 +28,8 @@ export function shouldRateLimitRequest(context: Context, agent: Agent): Result {
   if (
     context.remoteAddress &&
     !context.consumedRateLimitForIP &&
-    !isLocalhostIP(context.remoteAddress)
+    !isLocalhostIP(context.remoteAddress) &&
+    !agent.getConfig().isAllowedIP(context.remoteAddress)
   ) {
     const allowed = agent
       .getRateLimiter()