File size: 3,377 Bytes
f2d1fea 2b62b0c f2d1fea 2b62b0c f2d1fea 83e82db f2d1fea 83e82db bede2fc f2d1fea 83e82db f2d1fea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | import IORedis from "ioredis";
import { env } from "@pram/env/server";
import { TRPCError } from "@trpc/server";
import { logger } from "@pram/api/logger";
let redis: IORedis | null = null;
function getRedis(): IORedis {
if (!redis) {
redis = new IORedis(env.REDIS_URL, { maxRetriesPerRequest: 3, enableOfflineQueue: false });
}
return redis;
}
const SCRIPT = `
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local window_ms = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
-- Remove expired entries
redis.call("ZREMRANGEBYSCORE", key, 0, now - window_ms)
-- Count current entries
local count = redis.call("ZCARD", key)
if count >= limit then
return 0
end
-- Add current request
redis.call("ZADD", key, now, now .. "-" .. math.random())
redis.call("PEXPIRE", key, window_ms)
return 1
`;
export interface RateLimitConfig {
key: string;
limit: number;
windowMs: number;
/** When true, throw a 503 if Redis is unavailable instead of silently skipping. */
strict?: boolean;
}
export async function checkRateLimit(config: RateLimitConfig): Promise<void> {
const r = getRedis();
const key = `ratelimit:${config.key}`;
const now = Date.now();
try {
const allowed = await r.eval(SCRIPT, 1, key, config.limit.toString(), config.windowMs.toString(), now.toString());
if (allowed === 0) {
throw new TRPCError({
code: "TOO_MANY_REQUESTS",
message: "Too many requests. Please try again later.",
});
}
} catch (err) {
if (err instanceof TRPCError) throw err;
if (config.strict) {
logger.error("[RATELIMIT] Redis unavailable, rejecting request (strict mode)", { error: (err as Error).message });
throw new TRPCError({
code: "INTERNAL_SERVER_ERROR",
message: "Service temporarily unavailable. Please try again.",
});
}
logger.warn("[RATELIMIT] Redis unavailable, rate limit skipped", { error: (err as Error).message });
}
}
/**
* Returns true if the request is allowed, false if rate-limited.
* Used by non-tRPC contexts (e.g. Hono middleware) where throwing TRPCError is not appropriate.
* Fails open (returns true) when Redis is unavailable unless `strict` is set.
*/
export async function checkRateLimitAllowed(config: RateLimitConfig): Promise<boolean> {
const r = getRedis();
const key = `ratelimit:${config.key}`;
const now = Date.now();
try {
const allowed = await r.eval(SCRIPT, 1, key, config.limit.toString(), config.windowMs.toString(), now.toString());
return allowed !== 0;
} catch (err) {
if (config.strict) {
logger.error("[RATELIMIT] Redis unavailable (strict mode)", { error: (err as Error).message });
return false;
}
logger.warn("[RATELIMIT] Redis unavailable, rate limit skipped", { error: (err as Error).message });
return true;
}
}
export async function checkDailyBudget(key: string, limit: number): Promise<boolean> {
const r = getRedis();
const now = new Date();
const endOfDay = new Date(now.getFullYear(), now.getMonth(), now.getDate() + 1);
const ttlMs = endOfDay.getTime() - now.getTime();
try {
const current = await r.incrby(`budget:daily:${key}`, 0);
if (current > limit) return false;
await r.incr(`budget:daily:${key}`);
await r.pexpire(`budget:daily:${key}`, ttlMs);
return true;
} catch {
// Fail open
return true;
}
}
|