Merge pull request #221 from AikidoSec/AIK-2984

Match wildcard endpoints

Author: willem-delbare

library/agent/ServiceConfig.test.ts

@@ -3,10 +3,13 @@ import { ServiceConfig } from "./ServiceConfig";
 
 t.test("it returns false if empty rules", async () => {
   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.same(
+    config.getEndpoint({ url: undefined, method: undefined, route: undefined }),
+    undefined
+  );
 });
 
 t.test("it works", async () => {
@@ -48,51 +51,28 @@ t.test("it works", async () => {
     []
   );
 
-  t.same(config.shouldProtectEndpoint("GET", "/foo"), true);
-  t.same(config.shouldProtectEndpoint("POST", "/foo"), false);
-  t.same(config.shouldProtectEndpoint("GET", "/unknown"), true);
-  t.same(config.shouldProtectEndpoint("POST", /fly+/), false);
   t.same(config.isUserBlocked("123"), true);
   t.same(config.isUserBlocked("567"), false);
-});
-
-t.test("it returns rate limiting", async () => {
-  const config = new ServiceConfig(
-    [
-      {
+  t.same(
+    config.getEndpoint({
+      url: undefined,
+      method: "GET",
+      route: "/foo",
+    }),
+    {
+      endpoint: {
         method: "GET",
         route: "/foo",
         forceProtectionOff: false,
-        rateLimiting: { enabled: true, maxRequests: 10, windowSizeInMS: 1000 },
-      },
-      {
-        method: "POST",
-        route: "/foo",
-        forceProtectionOff: true,
         rateLimiting: {
           enabled: false,
           maxRequests: 0,
           windowSizeInMS: 0,
         },
       },
-    ],
-    0,
-    [],
-    []
+      route: "/foo",
+    }
   );
-
-  t.same(config.getRateLimiting("GET", "/foo"), {
-    enabled: true,
-    maxRequests: 10,
-    windowSizeInMS: 1000,
-  });
-
-  t.same(config.getRateLimiting("GET", "/unknown"), undefined);
-  t.same(config.getRateLimiting("POST", "/foo"), {
-    enabled: false,
-    maxRequests: 0,
-    windowSizeInMS: 0,
-  });
 });
 
 t.test("it checks if IP is allowed", async () => {

library/agent/ServiceConfig.ts

@@ -1,25 +1,16 @@
+import { LimitedContext, matchEndpoint } from "../helpers/matchEndpoint";
 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 endpoints: Endpoint[],
     private readonly lastUpdatedAt: number,
     blockedUserIds: string[],
     allowedIPAddresses: string[]
   ) {
-    endpoints.forEach((rule) => {
-      this.endpoints.set(this.getKey(rule.method, rule.route), {
-        method: rule.method,
-        route: rule.route,
-        forceProtectionOff: rule.forceProtectionOff,
-        rateLimiting: rule.rateLimiting,
-      });
-    });
-
     blockedUserIds.forEach((userId) => {
       this.blockedUserIds.set(userId, userId);
     });
@@ -29,44 +20,14 @@ export class ServiceConfig {
     });
   }
 
-  private getKey(method: string, route: string) {
-    return `${method}:${route}`;
-  }
-
-  getRateLimiting(method: string, route: string | RegExp) {
-    const key = this.getKey(
-      method,
-      typeof route === "string" ? route : route.source
-    );
-
-    const rule = this.endpoints.get(key);
-
-    if (!rule || !rule.rateLimiting) {
-      return undefined;
-    }
-
-    return rule.rateLimiting;
+  getEndpoint(context: LimitedContext) {
+    return matchEndpoint(context, this.endpoints);
   }
 
   isAllowedIP(ip: string) {
     return this.allowedIPAddresses.has(ip);
   }
 
-  shouldProtectEndpoint(method: string, route: string | RegExp) {
-    const key = this.getKey(
-      method,
-      typeof route === "string" ? route : route.source
-    );
-
-    const rule = this.endpoints.get(key);
-
-    if (!rule) {
-      return true;
-    }
-
-    return !rule.forceProtectionOff;
-  }
-
   isUserBlocked(userId: string) {
     return this.blockedUserIds.has(userId);
   }

library/agent/applyHooks.ts

@@ -160,20 +160,17 @@ function wrapWithoutArgumentModification(
         const args = Array.from(arguments);
         const context = getContext();
 
-        if (
-          context &&
-          context.method &&
-          context.route &&
-          !agent
-            .getConfig()
-            .shouldProtectEndpoint(context.method, context.route)
-        ) {
-          return original.apply(
-            // @ts-expect-error We don't now the type of this
-            this,
-            // eslint-disable-next-line prefer-rest-params
-            arguments
-          );
+        if (context) {
+          const match = agent.getConfig().getEndpoint(context);
+
+          if (match && match.endpoint.forceProtectionOff) {
+            return original.apply(
+              // @ts-expect-error We don't now the type of this
+              this,
+              // eslint-disable-next-line prefer-rest-params
+              arguments
+            );
+          }
         }
 
         const start = performance.now();