71 lines
2.3 KiB
TypeScript
71 lines
2.3 KiB
TypeScript
// WebSocket middleware: rate limiting, origin allow-list, message size cap.
|
|
// See PROTOCOL.md for full spec.
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Rate limiter — token-bucket algorithm, per connection
|
|
// ---------------------------------------------------------------------------
|
|
|
|
export class RateLimiter {
|
|
readonly #capacity: number;
|
|
readonly #refillRate: number; // tokens per second
|
|
#tokens: number;
|
|
#lastRefill: number;
|
|
|
|
constructor(capacity = 20, refillRate = 100) {
|
|
this.#capacity = capacity;
|
|
this.#refillRate = refillRate;
|
|
this.#tokens = capacity;
|
|
this.#lastRefill = performance.now();
|
|
}
|
|
|
|
/**
|
|
* Returns true if the message is allowed; false if rate-limited.
|
|
* Refills tokens based on elapsed time since last call.
|
|
*/
|
|
consume(): boolean {
|
|
const now = performance.now();
|
|
const elapsed = (now - this.#lastRefill) / 1000; // seconds
|
|
this.#tokens = Math.min(
|
|
this.#capacity,
|
|
this.#tokens + elapsed * this.#refillRate,
|
|
);
|
|
this.#lastRefill = now;
|
|
|
|
if (this.#tokens < 1) return false;
|
|
this.#tokens -= 1;
|
|
return true;
|
|
}
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Origin allow-list
|
|
// ---------------------------------------------------------------------------
|
|
|
|
export function getAllowedOrigins(): Set<string> {
|
|
const env = process.env.ALLOWED_ORIGINS ?? "http://localhost:5173";
|
|
return new Set(env.split(",").map((o) => o.trim()));
|
|
}
|
|
|
|
/**
|
|
* Returns true when requestOrigin is present and in the allow-list.
|
|
* A null/empty origin is always rejected.
|
|
*/
|
|
export function checkOrigin(requestOrigin: string | null): boolean {
|
|
if (!requestOrigin) return false;
|
|
return getAllowedOrigins().has(requestOrigin);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Message size cap — 64 KB
|
|
// ---------------------------------------------------------------------------
|
|
|
|
export const MAX_MESSAGE_BYTES = 64 * 1024; // 64 KB
|
|
|
|
/**
|
|
* Returns true when the message is within the 64 KB limit; false otherwise.
|
|
*/
|
|
export function checkMessageSize(msg: string | Buffer): boolean {
|
|
const size =
|
|
typeof msg === "string" ? Buffer.byteLength(msg, "utf8") : msg.length;
|
|
return size <= MAX_MESSAGE_BYTES;
|
|
}
|