62 lines
1.7 KiB
TypeScript
62 lines
1.7 KiB
TypeScript
/**
|
|
* Limits how often something may happen, with token buckets kept in memory,
|
|
* which a restart empties. Logins, password reset requests and the chat's
|
|
* typing messages use it.
|
|
*
|
|
* @module
|
|
*/
|
|
|
|
import type { Job } from "./scheduler.ts";
|
|
|
|
interface Bucket {
|
|
tokens: number;
|
|
at: number;
|
|
}
|
|
const buckets = new Map<string, Bucket>();
|
|
|
|
/**
|
|
* Passed to {@link takeToken} by each limited action, which declares its own
|
|
* as a constant.
|
|
*/
|
|
export interface RateLimit {
|
|
/** How many tokens a full bucket holds, and so the largest burst. */
|
|
capacity: number;
|
|
/** Tokens regained per second. */
|
|
refill: number;
|
|
}
|
|
|
|
/**
|
|
* Takes one token from `key`'s bucket, or returns false when it is empty. The
|
|
* caller picks the key, such as `typing:<userId>`, and keys by user or account
|
|
* so that more tabs or parallel requests don't buy more tokens.
|
|
*/
|
|
export function takeToken(key: string, limit: RateLimit): boolean {
|
|
const now = Date.now();
|
|
const bucket = buckets.get(key) ?? { tokens: limit.capacity, at: now };
|
|
bucket.tokens = Math.min(
|
|
limit.capacity,
|
|
bucket.tokens + ((now - bucket.at) / 1000) * limit.refill,
|
|
);
|
|
bucket.at = now;
|
|
buckets.set(key, bucket);
|
|
if (bucket.tokens < 1) return false;
|
|
bucket.tokens -= 1;
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Drops the buckets untouched for an hour, every minute. Every limit in use
|
|
* refills within the hour, so such a bucket is full and dropping it changes
|
|
* nothing. `main.ts` starts it.
|
|
*/
|
|
export const rateLimitSweepJob: Job = {
|
|
name: "rate-limit",
|
|
everyMs: 60_000,
|
|
run() {
|
|
const cutoff = Date.now() - 60 * 60_000;
|
|
for (const [key, bucket] of buckets) {
|
|
if (bucket.at < cutoff) buckets.delete(key);
|
|
}
|
|
},
|
|
};
|