Files
WRNexusJS/packages/core/src/ratelimit.ts
T
2026-07-12 15:55:18 +05:30

127 lines
4.6 KiB
TypeScript

/**
* Fixed-window rate limiting middleware. Keeps an in-memory counter per key
* (client IP by default, read from `x-forwarded-for` / `x-real-ip`) and rejects
* requests over the limit with a 429 and a `Retry-After` header. Sets the
* `RateLimit-Limit` / `RateLimit-Remaining` / `RateLimit-Reset` headers.
*
* The store is process-local; behind multiple instances use a shared store
* (out of scope here). Suitable as-is for single-process apps and dev.
*/
import type { Context, Middleware } from "./context.ts";
export interface RateLimitOptions {
/** Window length in milliseconds. Default 60_000 (1 minute). */
windowMs?: number;
/** Max requests allowed per key per window. Default 60. */
max?: number;
/** Derive the bucket key from the request. Default: client IP. */
key?: (ctx: Context) => string;
/**
* Trust `x-forwarded-for` / `x-real-ip` for the client IP. Default false —
* those headers are attacker-spoofable, so by default we key on the direct
* socket peer (`ctx.ip`). Enable ONLY when behind a proxy that overwrites
* these headers (nginx, a load balancer, Cloudflare).
*/
trustProxy?: boolean;
/** Body returned on 429. Default "Too Many Requests". */
message?: string;
/** Emit RateLimit-* headers. Default true. */
headers?: boolean;
/** Persistence for the counters. Default: process-local memory. */
store?: RateLimitStore;
/** Maximum in-memory keys before oldest buckets are evicted. Ignored for custom stores. */
maxKeys?: number;
}
export interface Bucket {
count: number;
resetAt: number;
}
/**
* Pluggable rate-limit counter store. The default is process-local memory; swap
* in a shared store (Redis/SQL) so limits hold across instances. `hit` records
* one request for `key` in the current window and returns the running bucket.
* It may be async (e.g. a Redis INCR + PEXPIRE) — the middleware awaits it.
*/
export interface RateLimitStore {
hit(key: string, windowMs: number, now: number): Bucket | Promise<Bucket>;
}
function createMemoryRateLimitStore(maxKeys: number): RateLimitStore {
const buckets = new Map<string, Bucket>();
return {
hit(key, windowMs, now) {
let bucket = buckets.get(key);
if (!bucket || bucket.resetAt <= now) {
if (!bucket && buckets.size >= maxKeys) {
for (const [k, b] of buckets) if (b.resetAt <= now) buckets.delete(k);
while (buckets.size >= maxKeys) buckets.delete(buckets.keys().next().value!);
}
bucket = { count: 0, resetAt: now + windowMs };
buckets.set(key, bucket);
}
bucket.count++;
return bucket;
},
};
}
export function rateLimit(options: RateLimitOptions = {}): Middleware {
const windowMs = options.windowMs ?? 60_000;
const max = options.max ?? 60;
const emitHeaders = options.headers ?? true;
const maxKeys = options.maxKeys ?? 10_000;
if (!Number.isInteger(maxKeys) || maxKeys < 1)
throw new RangeError("rateLimit maxKeys must be a positive integer");
const keyOf = options.key ?? (options.trustProxy ? proxyKey : peerKey);
const store = options.store ?? createMemoryRateLimitStore(maxKeys);
return async (ctx, next) => {
const now = Date.now();
const bucket = await store.hit(keyOf(ctx), windowMs, now);
const resetSec = Math.max(0, Math.ceil((bucket.resetAt - now) / 1000));
const remaining = Math.max(0, max - bucket.count);
if (bucket.count > max) {
const res = new Response(options.message ?? "Too Many Requests", {
status: 429,
headers: { "content-type": "text/plain", "retry-after": String(resetSec) },
});
if (emitHeaders) applyHeaders(res, max, 0, resetSec);
return res;
}
const res = await next();
if (emitHeaders) applyHeaders(res, max, remaining, resetSec);
return res;
};
}
function applyHeaders(res: Response, limit: number, remaining: number, resetSec: number): void {
try {
res.headers.set("RateLimit-Limit", String(limit));
res.headers.set("RateLimit-Remaining", String(remaining));
res.headers.set("RateLimit-Reset", String(resetSec));
} catch {
/* immutable response — skip */
}
}
/** Non-spoofable key: the direct socket peer IP (set by the server). */
export function peerKey(ctx: Context): string {
return ctx.ip ?? "global";
}
/** Proxy-aware key: trusts `x-forwarded-for` / `x-real-ip`, else the peer IP. */
export function proxyKey(ctx: Context): string {
const xff = ctx.req.headers.get("x-forwarded-for");
if (xff) return xff.split(",")[0]!.trim();
return ctx.req.headers.get("x-real-ip") ?? ctx.ip ?? "global";
}
/** @deprecated Use `peerKey` (default) or `proxyKey`. Kept for compatibility. */
export const defaultKey = proxyKey;