blob: c47889d371330ccef57848629f5c513eea82c7c4 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
|
import type {
RateLimitClient,
RateLimitConfig,
RateLimitResult,
} from "@karakeep/shared/ratelimiting";
import { PluginProvider } from "@karakeep/shared/plugins";
interface RateLimitEntry {
count: number;
resetTime: number;
}
export class RateLimiter implements RateLimitClient {
private store = new Map<string, RateLimitEntry>();
private cleanupProbability: number;
constructor(cleanupProbability = 0.01) {
// Probability of cleanup on each check (default 1%)
this.cleanupProbability = cleanupProbability;
}
private cleanupExpiredEntries() {
const now = Date.now();
for (const [key, entry] of this.store.entries()) {
if (now > entry.resetTime) {
this.store.delete(key);
}
}
}
checkRateLimit(config: RateLimitConfig, key: string): RateLimitResult {
if (!key) {
return { allowed: true };
}
// Probabilistic cleanup
if (Math.random() < this.cleanupProbability) {
this.cleanupExpiredEntries();
}
const rateLimitKey = `${config.name}:${key}`;
const now = Date.now();
let entry = this.store.get(rateLimitKey);
if (!entry || now > entry.resetTime) {
entry = {
count: 1,
resetTime: now + config.windowMs,
};
this.store.set(rateLimitKey, entry);
return { allowed: true };
}
if (entry.count >= config.maxRequests) {
const resetInSeconds = Math.ceil((entry.resetTime - now) / 1000);
return {
allowed: false,
resetInSeconds,
};
}
entry.count++;
return { allowed: true };
}
reset(config: RateLimitConfig, key: string) {
const rateLimitKey = `${config.name}:${key}`;
this.store.delete(rateLimitKey);
}
clear() {
this.store.clear();
}
}
export class RateLimitProvider implements PluginProvider<RateLimitClient> {
private client: RateLimiter | null = null;
async getClient(): Promise<RateLimitClient | null> {
if (!this.client) {
this.client = new RateLimiter();
}
return this.client;
}
}
|
