Spaces:
Runtime error
Runtime error
File size: 7,479 Bytes
cd8bd0a | 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 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 | import Redis from "ioredis";
// Redis is optional. When REDIS_URL is unset, use a process-local fallback
// instead of probing localhost on every API request.
const REDIS_URL = process.env.REDIS_URL?.trim() || "";
if (process.env.NODE_ENV === "production" && !REDIS_URL) {
console.warn("[REDIS] REDIS_URL is not set in production. Using in-memory rate limiting.");
}
let redisClient: Redis | null = null;
export function isRedisConfigured(): boolean {
return REDIS_URL.length > 0;
}
/**
* State-change-gated log throttle for the Redis error handler.
*
* #4878: when REDIS_URL points at a non-running server, ioredis retries on a
* backoff and fires the "error" event on every attempt, flooding the logs with
* identical "[REDIS] Error:" lines. We only want to log when the error STATE
* actually changes (first occurrence, or a different error message), not on
* every retry of the same failure.
*/
export function createRedisLogThrottle() {
let lastLogged: string | null = null;
return {
shouldLog(message: string): boolean {
if (message === lastLogged) return false;
lastLogged = message;
return true;
},
reset(): void {
lastLogged = null;
},
};
}
const redisLogThrottle = createRedisLogThrottle();
// Exposed for unit tests β returns a fresh, isolated throttle instance.
export function _createRedisLogThrottleForTests() {
return createRedisLogThrottle();
}
export function getRedisClient() {
if (!isRedisConfigured()) {
throw new Error("Redis is not configured");
}
if (!redisClient) {
redisClient = new Redis(REDIS_URL, {
maxRetriesPerRequest: 3,
enableReadyCheck: false,
retryStrategy(times) {
return Math.min(times * 50, 2000); // Exponential backoff
},
});
redisClient.on("error", (err) => {
// Throttle: log once per error-state change instead of on every retry (#4878).
if (redisLogThrottle.shouldLog(err.message)) {
console.error("[REDIS] Error:", err.message);
}
});
// A successful connection resets the throttle so the next failure logs again.
redisClient.on("ready", () => redisLogThrottle.reset());
}
return redisClient;
}
export interface RateLimitRule {
limit: number;
window: number; // in seconds
}
export interface RateLimitResult {
allowed: boolean;
failedWindow?: number;
}
/**
* Atomic Lua script for multi-rule rate limiting using fixed window.
* Returns {1, 0} if allowed, or {0, failedWindow} if rejected.
*/
const RATE_LIMIT_SCRIPT = `
local key_prefix = KEYS[1]
local current_time = tonumber(ARGV[1])
local rules = {}
for i = 2, #ARGV, 2 do
table.insert(rules, {
limit = tonumber(ARGV[i]),
window = tonumber(ARGV[i+1])
})
end
-- First pass: check if any limit is exceeded
for i, rule in ipairs(rules) do
local current_window = math.floor(current_time / rule.window)
local window_key = key_prefix .. ":" .. rule.window .. ":" .. current_window
local count = tonumber(redis.call("GET", window_key) or "0")
if count >= rule.limit then
return { 0, rule.window } -- Reject, return which window failed
end
end
-- Second pass: increment all rules
for i, rule in ipairs(rules) do
local current_window = math.floor(current_time / rule.window)
local window_key = key_prefix .. ":" .. rule.window .. ":" .. current_window
local count = redis.call("INCR", window_key)
if count == 1 then
-- TTL is twice the window size to ensure it covers the current window safely
redis.call("EXPIRE", window_key, rule.window * 2)
end
end
return { 1, 0 } -- Accepted
`;
const TEST_MEMORY_STORE = new Map<string, number>();
const FALLBACK_MEMORY_STORE = new Map<string, number>();
let explicitTestMode = false;
// Minimum store size before we bother sweeping (avoids O(n) cost on tiny stores)
const EVICTION_THRESHOLD = 50;
/**
* Evict all in-memory rate-limit window keys whose window has already ended.
*
* Key format: `rl:api_key:{id}:{windowSize}:{windowNumber}`
* A key expires at epoch-second `(windowNumber + 1) * windowSize`.
*
* Exported so tests can exercise it directly and so callers can invoke it
* with any store (TEST_MEMORY_STORE or FALLBACK_MEMORY_STORE).
*
* Fixes: #4041 β FALLBACK_MEMORY_STORE accumulated indefinitely β OOM (#4771).
*/
export function evictStaleRateLimitWindows(store: Map<string, number>, nowSeconds: number): void {
for (const key of store.keys()) {
// Format: rl:api_key:{id}:{windowSize}:{windowNumber}
// Split only on the last two colons to handle ids that contain colons.
const lastColon = key.lastIndexOf(":");
if (lastColon === -1) continue;
const secondLastColon = key.lastIndexOf(":", lastColon - 1);
if (secondLastColon === -1) continue;
const windowNumber = Number(key.slice(lastColon + 1));
const windowSize = Number(key.slice(secondLastColon + 1, lastColon));
if (!Number.isFinite(windowNumber) || !Number.isFinite(windowSize) || windowSize <= 0) {
continue;
}
const windowEnd = (windowNumber + 1) * windowSize;
if (windowEnd <= nowSeconds) {
store.delete(key);
}
}
}
export function setRateLimiterTestMode(enabled: boolean) {
explicitTestMode = enabled;
if (enabled) TEST_MEMORY_STORE.clear();
}
function checkInMemoryRateLimit(
store: Map<string, number>,
keyId: string,
rules: RateLimitRule[]
): RateLimitResult {
const now = Math.floor(Date.now() / 1000);
// Opportunistic eviction: sweep stale windows when the store has grown past
// the threshold. Bounded O(n) sweep β no timer, no background work.
if (store.size > EVICTION_THRESHOLD) {
evictStaleRateLimitWindows(store, now);
}
for (const rule of rules) {
const currentWindow = Math.floor(now / rule.window);
const windowKey = `rl:api_key:${keyId}:${rule.window}:${currentWindow}`;
const count = store.get(windowKey) || 0;
if (count >= rule.limit) {
return { allowed: false, failedWindow: rule.window };
}
}
for (const rule of rules) {
const currentWindow = Math.floor(now / rule.window);
const windowKey = `rl:api_key:${keyId}:${rule.window}:${currentWindow}`;
store.set(windowKey, (store.get(windowKey) || 0) + 1);
}
return { allowed: true };
}
export async function checkRateLimit(
keyId: string,
rules: RateLimitRule[]
): Promise<RateLimitResult> {
if (!rules || rules.length === 0) return { allowed: true };
// ββ In-memory mock for unit tests ββ
const isTestMode =
explicitTestMode ||
process.env.NODE_ENV === "test" ||
process.env.DISABLE_SQLITE_AUTO_BACKUP === "true";
if (isTestMode) {
return checkInMemoryRateLimit(TEST_MEMORY_STORE, keyId, rules);
}
if (!isRedisConfigured()) {
return checkInMemoryRateLimit(FALLBACK_MEMORY_STORE, keyId, rules);
}
const redis = getRedisClient();
const args: (string | number)[] = [Math.floor(Date.now() / 1000)];
for (const rule of rules) {
args.push(rule.limit, rule.window);
}
try {
const result = (await redis.eval(RATE_LIMIT_SCRIPT, 1, `rl:api_key:${keyId}`, ...args)) as [
number,
number,
];
if (result[0] === 0) {
return { allowed: false, failedWindow: result[1] };
}
return { allowed: true };
} catch (error) {
// Fail-open strategy if Redis goes down to prevent complete API outage
console.error("[RATE_LIMITER] Redis eval failed, bypassing rate limit:", error);
return { allowed: true };
}
}
|