Merge pull request #211 from AikidoSec/patch-rate-limiting

Also match on context.url.pathname for rate limiting

Author: willem-delbare

library/agent/Context.ts

@@ -13,6 +13,8 @@ export type Context = {
   body: unknown; // Can be an object, string or undefined (the body is parsed by something like body-parser)
   cookies: Record<string, string>;
   attackDetected?: boolean;
+  consumedRateLimitForIP?: boolean;
+  consumedRateLimitForUser?: boolean;
   user?: { id: string; name?: string };
   source: string;
   route: string | undefined;

library/helpers/tryParseURL.test.ts

@@ -0,0 +1,12 @@
+import * as t from "tap";
+import { tryParseURL } from "./tryParseURL";
+
+t.test("it returns undefined if invalid URL", async () => {
+  const url = tryParseURL("invalid");
+  t.same(url, undefined);
+});
+
+t.test("it returns URL if valid URL", async () => {
+  const url = tryParseURL("https://example.com");
+  t.same(url, new URL("https://example.com/"));
+});

library/helpers/tryParseURL.ts

@@ -0,0 +1,7 @@
+export function tryParseURL(url: string) {
+  try {
+    return new URL(url);
+  } catch {
+    return undefined;
+  }
+}

library/sources/Express.test.ts

@@ -36,6 +36,16 @@ const agent = new Agent(
           enabled: true,
         },
       },
+      {
+        method: "GET",
+        route: "/middleware-rate-limited",
+        forceProtectionOff: false,
+        rateLimiting: {
+          windowSizeInMS: 2000,
+          maxRequests: 3,
+          enabled: true,
+        },
+      },
     ],
     blockedUserIds: ["567"],
     configUpdatedAt: 0,
@@ -179,6 +189,10 @@ function getApp(userMiddleware = true) {
     res.send({ hello: "world" });
   });
 
+  app.use("/middleware-rate-limited", (req, res, next) => {
+    res.send({ hello: "world" });
+  });
+
   return app;
 }
 
@@ -390,3 +404,18 @@ t.test("it rate limits by user", async () => {
   const res2 = await request(getApp()).get("/user-rate-limited");
   t.same(res2.statusCode, 200);
 });
+
+t.test("it rate limits by middleware", async () => {
+  for (const _ of Array.from({ length: 3 })) {
+    const res = await request(getApp(false)).get("/middleware-rate-limited");
+    t.same(res.statusCode, 200);
+  }
+
+  const res = await request(getApp(false)).get("/middleware-rate-limited");
+  t.same(res.statusCode, 429);
+
+  await sleep(2000);
+
+  const res2 = await request(getApp(false)).get("/middleware-rate-limited");
+  t.same(res2.statusCode, 200);
+});

library/sources/express/shouldRateLimitRequest.ts

@@ -1,46 +1,82 @@
 import { Agent } from "../../agent/Agent";
 import { Context } from "../../agent/Context";
+import { tryParseURL } from "../../helpers/tryParseURL";
 
 export function shouldRateLimitRequest(context: Context, agent: Agent) {
-  if (!context.route || !context.method) {
-    return false;
-  }
-
-  const rateLimiting = agent
-    .getConfig()
-    .getRateLimiting(context.method, context.route);
+  const rateLimiting = getRateLimitingForContext(context, agent);
 
   if (!rateLimiting) {
     return false;
   }
 
-  if (context.remoteAddress) {
+  const { config, route } = rateLimiting;
+
+  if (context.remoteAddress && !context.consumedRateLimitForIP) {
     const allowed = agent
       .getRateLimiter()
       .isAllowed(
-        `${context.method}:${context.route}:ip:${context.remoteAddress}`,
-        rateLimiting.windowSizeInMS,
-        rateLimiting.maxRequests
+        `${context.method}:${route}:ip:${context.remoteAddress}`,
+        config.windowSizeInMS,
+        config.maxRequests
       );
 
+    // This function is executed for every middleware and route handler
+    // We want to count the request only once
+    context.consumedRateLimitForIP = true;
+
     if (!allowed) {
       return true;
     }
   }
 
-  if (context.user) {
+  if (context.user && !context.consumedRateLimitForUser) {
     const allowed = agent
       .getRateLimiter()
       .isAllowed(
-        `${context.method}:${context.route}:user:${context.user.id}`,
-        rateLimiting.windowSizeInMS,
-        rateLimiting.maxRequests
+        `${context.method}:${route}:user:${context.user.id}`,
+        config.windowSizeInMS,
+        config.maxRequests
       );
 
+    // This function is executed for every middleware and route handler
+    // We want to count the request only once
+    context.consumedRateLimitForUser = true;
+
     if (!allowed) {
       return true;
     }
   }
 
   return false;
 }
+
+function getRateLimitingForContext(context: Context, agent: Agent) {
+  if (!context.method) {
+    return undefined;
+  }
+
+  if (context.route) {
+    const rateLimiting = agent
+      .getConfig()
+      .getRateLimiting(context.method, context.route);
+
+    if (rateLimiting) {
+      return { config: rateLimiting, route: context.route };
+    }
+  }
+
+  if (context.url) {
+    const url = tryParseURL(context.url);
+    if (url && url.pathname) {
+      const rateLimiting = agent
+        .getConfig()
+        .getRateLimiting(context.method, url.pathname);
+
+      if (rateLimiting) {
+        return { config: rateLimiting, route: url.pathname };
+      }
+    }
+  }
+
+  return undefined;
+}