123 lines
3.2 KiB
TypeScript
123 lines
3.2 KiB
TypeScript
import type { NextFunction, Request, Response } from "express";
|
|
|
|
import { buildErrorPayload, createHttpError } from "../utils/http.js";
|
|
import { logWarn } from "../utils/logger.js";
|
|
|
|
type RateLimitKeyResolver = (req: Request) => string;
|
|
|
|
type RateLimitExceededContext = {
|
|
scope: string;
|
|
retryAfterSeconds: number;
|
|
};
|
|
|
|
type RateLimitOptions = {
|
|
scope: string;
|
|
windowMs: number;
|
|
max: number;
|
|
key?: RateLimitKeyResolver;
|
|
onLimit?: (req: Request, res: Response, context: RateLimitExceededContext) => void;
|
|
};
|
|
|
|
type RateLimitBucket = {
|
|
count: number;
|
|
resetAt: number;
|
|
};
|
|
|
|
const buckets = new Map<string, RateLimitBucket>();
|
|
let lastCleanupAt = 0;
|
|
|
|
export function createRateLimitMiddleware({
|
|
scope,
|
|
windowMs,
|
|
max,
|
|
key = defaultRateLimitKey,
|
|
onLimit,
|
|
}: RateLimitOptions) {
|
|
const normalizedScope = String(scope || "default").trim() || "default";
|
|
const normalizedWindowMs = Math.max(1000, Number(windowMs || 0));
|
|
const normalizedMax = Math.max(1, Number(max || 0));
|
|
|
|
return (req: Request, res: Response, next: NextFunction): void => {
|
|
const now = Date.now();
|
|
cleanupExpiredBuckets(now);
|
|
|
|
const bucketKey = `${normalizedScope}:${key(req)}`;
|
|
const current = buckets.get(bucketKey);
|
|
const bucket = current && current.resetAt > now
|
|
? current
|
|
: { count: 0, resetAt: now + normalizedWindowMs };
|
|
|
|
bucket.count += 1;
|
|
buckets.set(bucketKey, bucket);
|
|
|
|
if (bucket.count <= normalizedMax) {
|
|
next();
|
|
return;
|
|
}
|
|
|
|
const retryAfterSeconds = Math.max(1, Math.ceil((bucket.resetAt - now) / 1000));
|
|
res.setHeader("Retry-After", String(retryAfterSeconds));
|
|
|
|
logWarn("[rate-limit]", "请求触发限流", {
|
|
scope: normalizedScope,
|
|
ip: req.ip,
|
|
originalUrl: req.originalUrl,
|
|
retryAfterSeconds,
|
|
});
|
|
|
|
if (onLimit) {
|
|
onLimit(req, res, {
|
|
scope: normalizedScope,
|
|
retryAfterSeconds,
|
|
});
|
|
return;
|
|
}
|
|
|
|
const error = createHttpError("请求过于频繁，请稍后再试", {
|
|
statusCode: 429,
|
|
errorCode: "rate_limited",
|
|
});
|
|
res.status(429).json(buildErrorPayload(error, "请求过于频繁，请稍后再试"));
|
|
};
|
|
}
|
|
|
|
export function resetRateLimitBucketsForTest(): void {
|
|
buckets.clear();
|
|
lastCleanupAt = 0;
|
|
}
|
|
|
|
export function getBodyFieldRateLimitKey(fieldName: string): RateLimitKeyResolver {
|
|
return (req) => [defaultRateLimitKey(req), normalizeKeyPart(req.body?.[fieldName])].join(":");
|
|
}
|
|
|
|
export function getParamRateLimitKey(paramName: string): RateLimitKeyResolver {
|
|
return (req) => [defaultRateLimitKey(req), normalizeKeyPart(req.params?.[paramName])].join(":");
|
|
}
|
|
|
|
function defaultRateLimitKey(req: Request): string {
|
|
return normalizeKeyPart(
|
|
req.ip ||
|
|
String(req.headers["x-forwarded-for"] || "").split(",")[0] ||
|
|
req.socket.remoteAddress ||
|
|
"unknown",
|
|
);
|
|
}
|
|
|
|
function normalizeKeyPart(value: unknown): string {
|
|
const normalized = String(value || "").trim().toLowerCase();
|
|
return normalized || "unknown";
|
|
}
|
|
|
|
function cleanupExpiredBuckets(now: number): void {
|
|
if (now - lastCleanupAt < 60_000) {
|
|
return;
|
|
}
|
|
|
|
lastCleanupAt = now;
|
|
for (const [key, bucket] of buckets.entries()) {
|
|
if (bucket.resetAt <= now) {
|
|
buckets.delete(key);
|
|
}
|
|
}
|
|
}
|