Skip to content

Commit 10ed01d

Browse files
authored
Merge pull request #529 from AikidoSec/sliding-window
Use a sliding window for rate limiting
2 parents 1a5a2a8 + 76d147c commit 10ed01d

File tree

5 files changed

+309
-27
lines changed

5 files changed

+309
-27
lines changed

.github/workflows/benchmark.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,3 +53,5 @@ jobs:
5353
run: cd benchmarks/api-discovery && node benchmark.js
5454
- name: Run Express Benchmark
5555
run: cd benchmarks/express && node benchmark.js
56+
- name: Check Rate Limiter memory usage
57+
run: cd benchmarks/rate-limiting && node --expose-gc memory.js

benchmarks/rate-limiting/memory.js

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
const { RateLimiter } = require("../../build/ratelimiting/RateLimiter");
2+
3+
const ttl = 60000; // 1 minute in milliseconds
4+
const keyCount = 1_000_000;
5+
const checksPerKey = 10;
6+
7+
(async () => {
8+
const keys = Array.from({ length: keyCount }, (_, i) => `user${i}`);
9+
const limiter = new RateLimiter(100_000_000, ttl);
10+
11+
// Warmup, block all keys
12+
for (const key of keys) {
13+
limiter.isAllowed(key, ttl, 2);
14+
limiter.isAllowed(key, ttl, 2);
15+
}
16+
17+
global.gc();
18+
19+
const heapUsedBefore = process.memoryUsage().heapUsed;
20+
21+
for (const key of keys) {
22+
for (let i = 0; i < checksPerKey; i++) {
23+
limiter.isAllowed(key, ttl, 2);
24+
}
25+
}
26+
27+
global.gc();
28+
29+
const heapUsedAfter = process.memoryUsage().heapUsed;
30+
const heapUsedDiff = heapUsedAfter - heapUsedBefore;
31+
32+
if (heapUsedDiff > 0) {
33+
console.error(
34+
`Higher heap usage after rate limiting: +${heapUsedDiff} bytes`
35+
);
36+
process.exit(1);
37+
}
38+
39+
console.info(`No memory leak detected`);
40+
process.exit(0);
41+
})();
Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
import * as t from "tap";
2+
import { RateLimiter } from "./RateLimiter";
3+
4+
const ttl = 60000; // 1 minute in milliseconds
5+
const keyCount = 50_000;
6+
7+
const keys = Array.from({ length: keyCount }, (_, i) => `user${i}`);
8+
9+
t.test("check performance for first check for a key", async (t) => {
10+
const limiter = new RateLimiter(100_000_000, ttl);
11+
12+
const checkStart = performance.now();
13+
14+
for (const key of keys) {
15+
limiter.isAllowed(key, ttl, 3);
16+
}
17+
18+
const checkEnd = performance.now();
19+
const timePerCheck = (checkEnd - checkStart) / keyCount;
20+
21+
if (timePerCheck > 0.005 /* ms */) {
22+
t.fail(`Performance test failed: ${timePerCheck}ms per check`);
23+
} else {
24+
t.pass(`Performance test passed: ${timePerCheck}ms per check`);
25+
}
26+
});
27+
28+
t.test(
29+
"check performance for second check for a key (still allowed)",
30+
async (t) => {
31+
const limiter = new RateLimiter(100_000_000, ttl);
32+
33+
for (const key of keys) {
34+
limiter.isAllowed(key, ttl, 3);
35+
}
36+
37+
const checkStart = performance.now();
38+
39+
for (const key of keys) {
40+
limiter.isAllowed(key, ttl, 3);
41+
}
42+
43+
const checkEnd = performance.now();
44+
const timePerCheck = (checkEnd - checkStart) / keyCount;
45+
46+
if (timePerCheck > 0.005 /* ms */) {
47+
t.fail(`Performance test failed: ${timePerCheck}ms per check`);
48+
} else {
49+
t.pass(`Performance test passed: ${timePerCheck}ms per check`);
50+
}
51+
}
52+
);
53+
54+
t.test("check performance a blocked key", async (t) => {
55+
const limiter = new RateLimiter(100_000_000, ttl);
56+
57+
for (const key of keys) {
58+
limiter.isAllowed(key, ttl, 2);
59+
limiter.isAllowed(key, ttl, 2);
60+
}
61+
62+
const checkStart = performance.now();
63+
64+
for (const key of keys) {
65+
limiter.isAllowed(key, ttl, 2);
66+
}
67+
68+
const checkEnd = performance.now();
69+
const timePerCheck = (checkEnd - checkStart) / keyCount;
70+
71+
if (timePerCheck > 0.005 /* ms */) {
72+
t.fail(`Performance test failed: ${timePerCheck}ms per check`);
73+
} else {
74+
t.pass(`Performance test passed: ${timePerCheck}ms per check`);
75+
}
76+
});

library/ratelimiting/RateLimiter.test.ts

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,3 +86,169 @@ t.test("should allow requests for different keys independently", async (t) => {
8686
`Request ${maxAmount + 1} for key2 should not be allowed`
8787
);
8888
});
89+
90+
t.test("should handle TTL expiration", async (t) => {
91+
const limiter = new RateLimiter(maxAmount, ttl);
92+
for (let i = 0; i < maxAmount; i++) {
93+
limiter.isAllowed(key, ttl, maxAmount);
94+
}
95+
96+
clock.tick(ttl + 1);
97+
98+
t.ok(
99+
limiter.isAllowed(key, ttl, maxAmount),
100+
`Request after TTL should be allowed`
101+
);
102+
});
103+
104+
t.test("should allow requests exactly at limit", async (t) => {
105+
const limiter = new RateLimiter(maxAmount, ttl);
106+
for (let i = 0; i < maxAmount; i++) {
107+
t.ok(
108+
limiter.isAllowed(key, ttl, maxAmount),
109+
`Request ${i + 1} should be allowed`
110+
);
111+
}
112+
t.notOk(
113+
limiter.isAllowed(key, ttl, maxAmount),
114+
`Request ${maxAmount + 1} should not be allowed`
115+
);
116+
});
117+
118+
t.test("should handle multiple rapid requests", async (t) => {
119+
const limiter = new RateLimiter(maxAmount, ttl);
120+
for (let i = 0; i < maxAmount; i++) {
121+
t.ok(
122+
limiter.isAllowed(key, ttl, maxAmount),
123+
`Request ${i + 1} should be allowed`
124+
);
125+
}
126+
127+
clock.tick(100);
128+
129+
t.notOk(
130+
limiter.isAllowed(key, ttl, maxAmount),
131+
`Request ${maxAmount + 1} should not be allowed`
132+
);
133+
});
134+
135+
t.test("should handle different window sizes", async (t) => {
136+
const limiter = new RateLimiter(maxAmount, ttl);
137+
const differentWindowSize = 1000; // 1 second window
138+
for (let i = 0; i < maxAmount; i++) {
139+
t.ok(
140+
limiter.isAllowed(key, differentWindowSize, maxAmount),
141+
`Request ${i + 1} should be allowed`
142+
);
143+
}
144+
t.notOk(
145+
limiter.isAllowed(key, differentWindowSize, maxAmount),
146+
`Request ${maxAmount + 1} should not be allowed`
147+
);
148+
});
149+
150+
t.test("should handle sliding window with intermittent requests", async (t) => {
151+
const limiter = new RateLimiter(maxAmount, ttl);
152+
for (let i = 0; i < maxAmount; i++) {
153+
t.ok(
154+
limiter.isAllowed(key, ttl, maxAmount),
155+
`Request ${i + 1} should be allowed`
156+
);
157+
clock.tick(100);
158+
}
159+
160+
clock.tick(ttl + 1);
161+
162+
t.ok(
163+
limiter.isAllowed(key, ttl, maxAmount),
164+
`Request after sliding window should be allowed`
165+
);
166+
});
167+
168+
t.test("should handle sliding window edge case", async (t) => {
169+
const limiter = new RateLimiter(maxAmount, ttl);
170+
for (let i = 0; i < maxAmount; i++) {
171+
t.ok(
172+
limiter.isAllowed(key, ttl, maxAmount),
173+
`Request ${i + 1} should be allowed`
174+
);
175+
}
176+
177+
clock.tick(ttl + 1);
178+
179+
t.ok(
180+
limiter.isAllowed(key, ttl, maxAmount),
181+
`Request after sliding window should be allowed`
182+
);
183+
184+
clock.tick(ttl + 1);
185+
186+
t.ok(
187+
limiter.isAllowed(key, ttl, maxAmount),
188+
`Request after sliding window should be allowed`
189+
);
190+
});
191+
192+
t.test("should handle sliding window with delayed requests", async (t) => {
193+
const limiter = new RateLimiter(maxAmount, ttl);
194+
for (let i = 0; i < maxAmount; i++) {
195+
t.ok(
196+
limiter.isAllowed(key, ttl, maxAmount),
197+
`Request ${i + 1} should be allowed`
198+
);
199+
clock.tick(100);
200+
}
201+
202+
clock.tick(ttl + 1);
203+
204+
t.ok(
205+
limiter.isAllowed(key, ttl, maxAmount),
206+
`Request after sliding window should be allowed`
207+
);
208+
});
209+
210+
t.test("should handle sliding window with burst requests", async (t) => {
211+
const limiter = new RateLimiter(maxAmount, ttl);
212+
for (let i = 0; i < maxAmount; i++) {
213+
t.ok(
214+
limiter.isAllowed(key, ttl, maxAmount),
215+
`Request ${i + 1} should be allowed`
216+
);
217+
}
218+
219+
clock.tick(ttl / 2 + 1);
220+
221+
t.notOk(
222+
limiter.isAllowed(key, ttl, maxAmount),
223+
`Request ${maxAmount + 1} should not be allowed`
224+
);
225+
t.notOk(
226+
limiter.isAllowed(key, ttl, maxAmount),
227+
`Request ${maxAmount + 2} should not be allowed`
228+
);
229+
t.notOk(
230+
limiter.isAllowed(key, ttl, maxAmount),
231+
`Request ${maxAmount + 3} should not be allowed`
232+
);
233+
234+
clock.tick(ttl / 2 + 1);
235+
236+
for (let i = 0; i < 2; i++) {
237+
t.ok(
238+
limiter.isAllowed(key, ttl, maxAmount),
239+
`Request ${i + 1} should be allowed`
240+
);
241+
}
242+
243+
t.notOk(
244+
limiter.isAllowed(key, ttl, maxAmount),
245+
`Request ${maxAmount + 1} should not be allowed`
246+
);
247+
248+
clock.tick(ttl + 1);
249+
250+
t.ok(
251+
limiter.isAllowed(key, ttl, maxAmount),
252+
`Request after sliding window should be allowed`
253+
);
254+
});

library/ratelimiting/RateLimiter.ts

Lines changed: 24 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,45 +1,42 @@
11
import { LRUMap } from "./LRUMap";
22

3+
/**
4+
* Sliding window rate limiter implementation
5+
*/
36
export class RateLimiter {
4-
private rateLimitedItems: LRUMap<
5-
string,
6-
{ count: number; startTime: number }
7-
>;
7+
private rateLimitedItems: LRUMap<string, number[]>;
88

99
constructor(
1010
readonly maxItems: number,
1111
readonly timeToLiveInMS: number
1212
) {
13-
this.rateLimitedItems = new LRUMap<
14-
string,
15-
{ count: number; startTime: number }
16-
>(maxItems, timeToLiveInMS);
13+
this.rateLimitedItems = new LRUMap(maxItems, timeToLiveInMS);
1714
}
1815

1916
isAllowed(key: string, windowSizeInMS: number, maxRequests: number): boolean {
2017
const currentTime = performance.now();
21-
const requestInfo = this.rateLimitedItems.get(key);
22-
23-
if (!requestInfo) {
24-
this.rateLimitedItems.set(key, { count: 1, startTime: currentTime });
25-
return true;
18+
const requestTimestamps = this.rateLimitedItems.get(key) || [];
19+
20+
// Filter out timestamps that are older than windowSizeInMS and already expired
21+
const filteredTimestamps = requestTimestamps.filter(
22+
(timestamp) => currentTime - timestamp <= windowSizeInMS
23+
);
24+
25+
// Ensure the number of entries exceeds maxRequests by only 1
26+
if (filteredTimestamps.length > maxRequests + 1) {
27+
filteredTimestamps.splice(
28+
0,
29+
filteredTimestamps.length - (maxRequests + 1)
30+
);
2631
}
2732

28-
const elapsedTime = currentTime - requestInfo.startTime;
29-
30-
if (elapsedTime > windowSizeInMS) {
31-
// Reset the counter and timestamp if windowSizeInMS has expired
32-
this.rateLimitedItems.set(key, { count: 1, startTime: currentTime });
33-
return true;
34-
}
33+
// Add current request timestamp to the list
34+
filteredTimestamps.push(currentTime);
3535

36-
if (requestInfo.count < maxRequests) {
37-
// Increment the counter if it is within the windowSizeInMS and maxRequests
38-
requestInfo.count += 1;
39-
return true;
40-
}
36+
// Update the list of timestamps for the key
37+
this.rateLimitedItems.set(key, filteredTimestamps);
4138

42-
// Deny the request if the maxRequests is reached within windowSizeInMS
43-
return false;
39+
// Check if the number of requests is less or equal to the maxRequests
40+
return filteredTimestamps.length <= maxRequests;
4441
}
4542
}

0 commit comments

Comments
 (0)