EpicNext-Cms/src/lib/rate-limit.ts

import { headers } from "next/headers";
import { resolveClientIp, UNKNOWN_CLIENT_IP } from "@/lib/client-ip";
import { logger } from "@/lib/logger";
import { redis } from "@/lib/redis";

type Bucket = { count: number; resetAt: number };
const buckets = new Map<string, Bucket>();

export interface RateLimitResult {
	ok: boolean;
	retryAfter: number;
}

const CLEANUP_INTERVAL_MS = 300_000;
const MAX_BUCKETS = 10_000;

let lastCleanup = Date.now();
let redisFailWarned = false;

function cleanup(): void {
	const now = Date.now();
	if (now - lastCleanup < CLEANUP_INTERVAL_MS) return;
	lastCleanup = now;

	for (const [k, b] of buckets) {
		if (now >= b.resetAt) buckets.delete(k);
	}

	if (buckets.size > MAX_BUCKETS) {
		const sorted = [...buckets.entries()].sort(
			(a, b) => a[1].resetAt - b[1].resetAt,
		);
		const keysToRemove = sorted
			.slice(0, Math.floor(sorted.length * 0.2))
			.map((entry) => entry[0]);
		for (const key of keysToRemove) buckets.delete(key);
	}
}

export async function rateLimit(
	key: string,
	limit: number,
	windowMs: number,
): Promise<RateLimitResult> {
	const now = Date.now();

	if (redis) {
		try {
			const windowKey = `ratelimit:${key}`;
			// One atomic round trip: INCR + PEXPIRE in a single Lua call. Two
			// separate commands left keys without a TTL whenever the process died
			// in between, which permanently locked that client out.
			const [current, ttl] = (await redis.eval(
				`local c = redis.call('INCR', KEYS[1])
				 if c == 1 then redis.call('PEXPIRE', KEYS[1], ARGV[1]) end
				 local t = redis.call('PTTL', KEYS[1])
				 if t < 0 then t = tonumber(ARGV[1]) end
				 return {c, t}`,
				1,
				windowKey,
				String(windowMs),
			)) as [number, number];
			if (current > limit) {
				return { ok: false, retryAfter: Math.max(1, Math.ceil(ttl / 1000)) };
			}
			return { ok: true, retryAfter: 0 };
		} catch {
			// Redis unavailable — fall through to in-memory
			if (process.env.NODE_ENV === "production" && !redisFailWarned) {
				redisFailWarned = true;
				logger.error(
					"[rate-limit] Redis error — falling back to in-process buckets. Limits are not shared across instances until Redis recovers.",
				);
			}
		}
	} else if (redisFailWarned) {
		redisFailWarned = false;
		logger.info("[rate-limit] Redis available again — shared limits restored.");
	}

	cleanup();

	const windowKey = `mem:${key}`;
	const bucket = buckets.get(windowKey);

	if (!bucket || now >= bucket.resetAt) {
		buckets.set(windowKey, { count: 1, resetAt: now + windowMs });
		return { ok: true, retryAfter: 0 };
	}

	const newCount = bucket.count + 1;
	if (newCount > limit) {
		return {
			ok: false,
			retryAfter: Math.max(1, Math.ceil((bucket.resetAt - now) / 1000)),
		};
	}

	bucket.count = newCount;
	return { ok: true, retryAfter: 0 };
}

export async function clientIp(): Promise<string> {
	try {
		return resolveClientIp(await headers());
	} catch {
		return UNKNOWN_CLIENT_IP;
	}
}