import IORedis from "ioredis"; import { prisma } from "./db"; import { isSelfHostedEdition } from "./edition"; export const DEFAULT_API_RATE_LIMIT = 60; export const API_RATE_WINDOW_SECONDS = 60; export const MAX_ADMIN_API_RATE_BONUS = 120; const SELFHOSTED_API_RATE_LIMIT = 100_000; type MemoryBucket = { count: number; resetAt: number; }; const memoryBuckets = new Map(); let redis: IORedis | null | undefined; function getRedis() { if (redis !== undefined) return redis; try { const url = process.env.REDIS_URL || "redis://localhost:6379"; redis = new IORedis(url, { maxRetriesPerRequest: 1, enableOfflineQueue: false, lazyConnect: true, }); return redis; } catch { redis = null; return null; } } function rateKey(userId: string) { return `s2yt:api-rate:${userId}`; } export async function getEffectiveApiRateLimit(userId: string) { if (isSelfHostedEdition()) return SELFHOSTED_API_RATE_LIMIT; const user = await prisma.user.findUnique({ where: { id: userId }, select: { apiRateLimitBonus: true, plan: true }, }); if (!user || user.plan !== "PREMIUM") return DEFAULT_API_RATE_LIMIT; return DEFAULT_API_RATE_LIMIT + Math.max(0, user.apiRateLimitBonus); } function memoryStatus(userId: string, limit: number) { const now = Date.now(); const bucket = memoryBuckets.get(userId); if (!bucket || now >= bucket.resetAt) { return { limit, used: 0, remaining: limit, windowSeconds: API_RATE_WINDOW_SECONDS, resetsInSeconds: API_RATE_WINDOW_SECONDS, bonus: Math.max(0, limit - DEFAULT_API_RATE_LIMIT), }; } const used = Math.min(bucket.count, limit); return { limit, used, remaining: Math.max(0, limit - used), windowSeconds: API_RATE_WINDOW_SECONDS, resetsInSeconds: Math.max(1, Math.ceil((bucket.resetAt - now) / 1000)), bonus: Math.max(0, limit - DEFAULT_API_RATE_LIMIT), }; } export async function getApiRateLimitStatus(userId: string) { const limit = await getEffectiveApiRateLimit(userId); const client = getRedis(); if (!client) return memoryStatus(userId, limit); try { if (client.status !== "ready") { await client.connect().catch(() => {}); } const key = rateKey(userId); const [countRaw, ttl] = await Promise.all([client.get(key), client.ttl(key)]); const used = Math.min(Number(countRaw || 0), limit); return { limit, used, remaining: Math.max(0, limit - used), windowSeconds: API_RATE_WINDOW_SECONDS, resetsInSeconds: ttl > 0 ? ttl : API_RATE_WINDOW_SECONDS, bonus: Math.max(0, limit - DEFAULT_API_RATE_LIMIT), }; } catch { return { ...memoryStatus(userId, limit), bonus: Math.max(0, limit - DEFAULT_API_RATE_LIMIT) }; } } function checkMemoryRateLimit(userId: string, limit: number) { const now = Date.now(); const bucket = memoryBuckets.get(userId); if (!bucket || now >= bucket.resetAt) { memoryBuckets.set(userId, { count: 1, resetAt: now + API_RATE_WINDOW_SECONDS * 1000, }); return { ok: true as const, limit, remaining: limit - 1 }; } if (bucket.count >= limit) { return { ok: false as const, retryAfterSeconds: Math.max(1, Math.ceil((bucket.resetAt - now) / 1000)), limit, remaining: 0, }; } bucket.count += 1; return { ok: true as const, limit, remaining: Math.max(0, limit - bucket.count) }; } export async function checkApiRateLimit(userId: string) { if (isSelfHostedEdition()) { return { ok: true as const, limit: SELFHOSTED_API_RATE_LIMIT, remaining: SELFHOSTED_API_RATE_LIMIT }; } const limit = await getEffectiveApiRateLimit(userId); const client = getRedis(); if (!client) return checkMemoryRateLimit(userId, limit); try { if (client.status !== "ready") { await client.connect().catch(() => {}); } const key = rateKey(userId); const count = await client.incr(key); if (count === 1) { await client.expire(key, API_RATE_WINDOW_SECONDS); } if (count > limit) { const ttl = await client.ttl(key); return { ok: false as const, retryAfterSeconds: Math.max(1, ttl > 0 ? ttl : API_RATE_WINDOW_SECONDS), limit, remaining: 0, }; } return { ok: true as const, limit, remaining: Math.max(0, limit - count) }; } catch { return checkMemoryRateLimit(userId, limit); } }