Merge pull request #529 from AikidoSec/sliding-window

Use a sliding window for rate limiting

Author: hansott

.github/workflows/benchmark.yml

@@ -53,3 +53,5 @@ jobs:
         run: cd benchmarks/api-discovery && node benchmark.js
       - name: Run Express Benchmark
         run: cd benchmarks/express && node benchmark.js
+      - name: Check Rate Limiter memory usage
+        run: cd benchmarks/rate-limiting && node --expose-gc memory.js

benchmarks/rate-limiting/memory.js

@@ -0,0 +1,41 @@
+const { RateLimiter } = require("../../build/ratelimiting/RateLimiter");
+
+const ttl = 60000; // 1 minute in milliseconds
+const keyCount = 1_000_000;
+const checksPerKey = 10;
+
+(async () => {
+  const keys = Array.from({ length: keyCount }, (_, i) => `user${i}`);
+  const limiter = new RateLimiter(100_000_000, ttl);
+
+  // Warmup, block all keys
+  for (const key of keys) {
+    limiter.isAllowed(key, ttl, 2);
+    limiter.isAllowed(key, ttl, 2);
+  }
+
+  global.gc();
+
+  const heapUsedBefore = process.memoryUsage().heapUsed;
+
+  for (const key of keys) {
+    for (let i = 0; i < checksPerKey; i++) {
+      limiter.isAllowed(key, ttl, 2);
+    }
+  }
+
+  global.gc();
+
+  const heapUsedAfter = process.memoryUsage().heapUsed;
+  const heapUsedDiff = heapUsedAfter - heapUsedBefore;
+
+  if (heapUsedDiff > 0) {
+    console.error(
+      `Higher heap usage after rate limiting: +${heapUsedDiff} bytes`
+    );
+    process.exit(1);
+  }
+
+  console.info(`No memory leak detected`);
+  process.exit(0);
+})();

library/ratelimiting/RateLimiter.performance.test.ts

@@ -0,0 +1,76 @@
+import * as t from "tap";
+import { RateLimiter } from "./RateLimiter";
+
+const ttl = 60000; // 1 minute in milliseconds
+const keyCount = 50_000;
+
+const keys = Array.from({ length: keyCount }, (_, i) => `user${i}`);
+
+t.test("check performance for first check for a key", async (t) => {
+  const limiter = new RateLimiter(100_000_000, ttl);
+
+  const checkStart = performance.now();
+
+  for (const key of keys) {
+    limiter.isAllowed(key, ttl, 3);
+  }
+
+  const checkEnd = performance.now();
+  const timePerCheck = (checkEnd - checkStart) / keyCount;
+
+  if (timePerCheck > 0.005 /* ms */) {
+    t.fail(`Performance test failed: ${timePerCheck}ms per check`);
+  } else {
+    t.pass(`Performance test passed: ${timePerCheck}ms per check`);
+  }
+});
+
+t.test(
+  "check performance for second check for a key (still allowed)",
+  async (t) => {
+    const limiter = new RateLimiter(100_000_000, ttl);
+
+    for (const key of keys) {
+      limiter.isAllowed(key, ttl, 3);
+    }
+
+    const checkStart = performance.now();
+
+    for (const key of keys) {
+      limiter.isAllowed(key, ttl, 3);
+    }
+
+    const checkEnd = performance.now();
+    const timePerCheck = (checkEnd - checkStart) / keyCount;
+
+    if (timePerCheck > 0.005 /* ms */) {
+      t.fail(`Performance test failed: ${timePerCheck}ms per check`);
+    } else {
+      t.pass(`Performance test passed: ${timePerCheck}ms per check`);
+    }
+  }
+);
+
+t.test("check performance a blocked key", async (t) => {
+  const limiter = new RateLimiter(100_000_000, ttl);
+
+  for (const key of keys) {
+    limiter.isAllowed(key, ttl, 2);
+    limiter.isAllowed(key, ttl, 2);
+  }
+
+  const checkStart = performance.now();
+
+  for (const key of keys) {
+    limiter.isAllowed(key, ttl, 2);
+  }
+
+  const checkEnd = performance.now();
+  const timePerCheck = (checkEnd - checkStart) / keyCount;
+
+  if (timePerCheck > 0.005 /* ms */) {
+    t.fail(`Performance test failed: ${timePerCheck}ms per check`);
+  } else {
+    t.pass(`Performance test passed: ${timePerCheck}ms per check`);
+  }
+});

library/ratelimiting/RateLimiter.test.ts

@@ -86,3 +86,169 @@ t.test("should allow requests for different keys independently", async (t) => {
     `Request ${maxAmount + 1} for key2 should not be allowed`
   );
 });
+
+t.test("should handle TTL expiration", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    limiter.isAllowed(key, ttl, maxAmount);
+  }
+
+  clock.tick(ttl + 1);
+
+  t.ok(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request after TTL should be allowed`
+  );
+});
+
+t.test("should allow requests exactly at limit", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+  }
+  t.notOk(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request ${maxAmount + 1} should not be allowed`
+  );
+});
+
+t.test("should handle multiple rapid requests", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+  }
+
+  clock.tick(100);
+
+  t.notOk(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request ${maxAmount + 1} should not be allowed`
+  );
+});
+
+t.test("should handle different window sizes", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  const differentWindowSize = 1000; // 1 second window
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, differentWindowSize, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+  }
+  t.notOk(
+    limiter.isAllowed(key, differentWindowSize, maxAmount),
+    `Request ${maxAmount + 1} should not be allowed`
+  );
+});
+
+t.test("should handle sliding window with intermittent requests", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+    clock.tick(100);
+  }
+
+  clock.tick(ttl + 1);
+
+  t.ok(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request after sliding window should be allowed`
+  );
+});
+
+t.test("should handle sliding window edge case", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+  }
+
+  clock.tick(ttl + 1);
+
+  t.ok(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request after sliding window should be allowed`
+  );
+
+  clock.tick(ttl + 1);
+
+  t.ok(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request after sliding window should be allowed`
+  );
+});
+
+t.test("should handle sliding window with delayed requests", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+    clock.tick(100);
+  }
+
+  clock.tick(ttl + 1);
+
+  t.ok(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request after sliding window should be allowed`
+  );
+});
+
+t.test("should handle sliding window with burst requests", async (t) => {
+  const limiter = new RateLimiter(maxAmount, ttl);
+  for (let i = 0; i < maxAmount; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+  }
+
+  clock.tick(ttl / 2 + 1);
+
+  t.notOk(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request ${maxAmount + 1} should not be allowed`
+  );
+  t.notOk(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request ${maxAmount + 2} should not be allowed`
+  );
+  t.notOk(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request ${maxAmount + 3} should not be allowed`
+  );
+
+  clock.tick(ttl / 2 + 1);
+
+  for (let i = 0; i < 2; i++) {
+    t.ok(
+      limiter.isAllowed(key, ttl, maxAmount),
+      `Request ${i + 1} should be allowed`
+    );
+  }
+
+  t.notOk(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request ${maxAmount + 1} should not be allowed`
+  );
+
+  clock.tick(ttl + 1);
+
+  t.ok(
+    limiter.isAllowed(key, ttl, maxAmount),
+    `Request after sliding window should be allowed`
+  );
+});

library/ratelimiting/RateLimiter.ts

@@ -1,45 +1,42 @@
 import { LRUMap } from "./LRUMap";
 
+/**
+ * Sliding window rate limiter implementation
+ */
 export class RateLimiter {
-  private rateLimitedItems: LRUMap<
-    string,
-    { count: number; startTime: number }
-  >;
+  private rateLimitedItems: LRUMap<string, number[]>;
 
   constructor(
     readonly maxItems: number,
     readonly timeToLiveInMS: number
   ) {
-    this.rateLimitedItems = new LRUMap<
-      string,
-      { count: number; startTime: number }
-    >(maxItems, timeToLiveInMS);
+    this.rateLimitedItems = new LRUMap(maxItems, timeToLiveInMS);
   }
 
   isAllowed(key: string, windowSizeInMS: number, maxRequests: number): boolean {
     const currentTime = performance.now();
-    const requestInfo = this.rateLimitedItems.get(key);
-
-    if (!requestInfo) {
-      this.rateLimitedItems.set(key, { count: 1, startTime: currentTime });
-      return true;
+    const requestTimestamps = this.rateLimitedItems.get(key) || [];
+
+    // Filter out timestamps that are older than windowSizeInMS and already expired
+    const filteredTimestamps = requestTimestamps.filter(
+      (timestamp) => currentTime - timestamp <= windowSizeInMS
+    );
+
+    // Ensure the number of entries exceeds maxRequests by only 1
+    if (filteredTimestamps.length > maxRequests + 1) {
+      filteredTimestamps.splice(
+        0,
+        filteredTimestamps.length - (maxRequests + 1)
+      );
     }
 
-    const elapsedTime = currentTime - requestInfo.startTime;
-
-    if (elapsedTime > windowSizeInMS) {
-      // Reset the counter and timestamp if windowSizeInMS has expired
-      this.rateLimitedItems.set(key, { count: 1, startTime: currentTime });
-      return true;
-    }
+    // Add current request timestamp to the list
+    filteredTimestamps.push(currentTime);
 
-    if (requestInfo.count < maxRequests) {
-      // Increment the counter if it is within the windowSizeInMS and maxRequests
-      requestInfo.count += 1;
-      return true;
-    }
+    // Update the list of timestamps for the key
+    this.rateLimitedItems.set(key, filteredTimestamps);
 
-    // Deny the request if the maxRequests is reached within windowSizeInMS
-    return false;
+    // Check if the number of requests is less or equal to the maxRequests
+    return filteredTimestamps.length <= maxRequests;
   }
 }