songs2yt/lib/api-rate-limit.ts

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<string, MemoryBucket>();
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);
  }
}