95 lines
2.2 KiB
TypeScript
95 lines
2.2 KiB
TypeScript
export type RateLimitResult = {
|
|
allowed: boolean;
|
|
limit: number;
|
|
used: number;
|
|
remaining: number;
|
|
resetAt: number;
|
|
};
|
|
|
|
type WindowEntry = {
|
|
used: number;
|
|
expiresAt: number;
|
|
};
|
|
|
|
type RateLimitState = {
|
|
limit: number;
|
|
windowMs: number;
|
|
entries: Map<string, WindowEntry>;
|
|
};
|
|
|
|
const rateLimitBuckets = new Map<string, RateLimitState>();
|
|
|
|
function getState(scope: string, limit: number, windowMs: number) {
|
|
const key = `${scope}|${limit}|${windowMs}`;
|
|
const existing = rateLimitBuckets.get(key);
|
|
if (existing) {
|
|
return existing;
|
|
}
|
|
|
|
const state = { limit, windowMs, entries: new Map<string, WindowEntry>() };
|
|
rateLimitBuckets.set(key, state);
|
|
return state;
|
|
}
|
|
|
|
export function consumeRateLimit(
|
|
identifier: string,
|
|
options: {
|
|
scope: string;
|
|
limit: number;
|
|
windowMs: number;
|
|
}
|
|
): RateLimitResult {
|
|
const { scope, limit, windowMs } = options;
|
|
const now = Date.now();
|
|
const key = `${identifier}`;
|
|
const state = getState(scope, limit, windowMs);
|
|
if (state.entries.size > 200) {
|
|
for (const [storedIdentifier, entry] of state.entries) {
|
|
if (entry.expiresAt <= now) {
|
|
state.entries.delete(storedIdentifier);
|
|
}
|
|
}
|
|
}
|
|
|
|
const current = state.entries.get(key);
|
|
if (!current || current.expiresAt <= now) {
|
|
const fresh = { used: 1, expiresAt: now + windowMs };
|
|
state.entries.set(key, fresh);
|
|
return {
|
|
allowed: true,
|
|
limit,
|
|
used: fresh.used,
|
|
remaining: limit - fresh.used,
|
|
resetAt: fresh.expiresAt
|
|
};
|
|
}
|
|
|
|
if (current.used >= limit) {
|
|
return {
|
|
allowed: false,
|
|
limit,
|
|
used: current.used,
|
|
remaining: Math.max(0, limit - current.used),
|
|
resetAt: current.expiresAt
|
|
};
|
|
}
|
|
|
|
current.used += 1;
|
|
return {
|
|
allowed: true,
|
|
limit,
|
|
used: current.used,
|
|
remaining: limit - current.used,
|
|
resetAt: current.expiresAt
|
|
};
|
|
}
|
|
|
|
export function getRateLimitHeaders(result: RateLimitResult) {
|
|
return {
|
|
"RateLimit-Limit": String(result.limit),
|
|
"RateLimit-Remaining": String(result.remaining),
|
|
"RateLimit-Reset": String(Math.max(0, Math.ceil(result.resetAt / 1000))),
|
|
"Retry-After": String(Math.max(1, Math.ceil((result.resetAt - Date.now()) / 1000)))
|
|
};
|
|
}
|