freezmo/lib/rate-limit.ts

/**
 * In-memory sliding-window rate limiter.
 *
 * Deliberately dependency-free and good enough for a self-hosted, single
 * instance app: each key (route + client IP) gets N requests per window.
 * Counters live in memory and reset on restart, and a busy instance cannot
 * share state with others — fine for Freezmo, documented in the README.
 */

import { DEFAULT_LOCALE, translate, type Locale } from "./i18n-dict";
type Hits = { timestamps: number[] };

/**
 * Kept on globalThis so dev-mode hot reloads (and any module re-evaluation)
 * don't silently reset everybody's counters.
 */
const globalStore = globalThis as unknown as { __freezmoRateLimit?: Map<string, Hits> };
const buckets: Map<string, Hits> = (globalStore.__freezmoRateLimit ??= new Map());
const MAX_BUCKETS = 5_000;

export type RateLimitResult = {
  allowed: boolean;
  limit: number;
  remaining: number;
  /** Seconds until the oldest hit falls out of the window (0 when allowed). */
  retryAfter: number;
  resetAt: number;
};

export function rateLimit(key: string, limit: number, windowMs: number): RateLimitResult {
  const now = Date.now();
  const windowStart = now - windowMs;

  let bucket = buckets.get(key);
  if (!bucket) {
    if (buckets.size >= MAX_BUCKETS) prune(now, windowMs);
    bucket = { timestamps: [] };
    buckets.set(key, bucket);
  }

  bucket.timestamps = bucket.timestamps.filter((timestamp) => timestamp > windowStart);

  if (bucket.timestamps.length >= limit) {
    const oldest = bucket.timestamps[0];
    const retryAfter = Math.max(1, Math.ceil((oldest + windowMs - now) / 1000));
    return {
      allowed: false,
      limit,
      remaining: 0,
      retryAfter,
      resetAt: oldest + windowMs,
    };
  }

  bucket.timestamps.push(now);
  return {
    allowed: true,
    limit,
    remaining: Math.max(0, limit - bucket.timestamps.length),
    retryAfter: 0,
    resetAt: now + windowMs,
  };
}

/** Drop stale buckets so long-running servers don't grow without bound. */
function prune(now: number, windowMs: number): void {
  const cutoff = now - windowMs;
  for (const [key, bucket] of buckets) {
    if (bucket.timestamps.length === 0 || bucket.timestamps[bucket.timestamps.length - 1] < cutoff) {
      buckets.delete(key);
    }
    if (buckets.size < MAX_BUCKETS / 2) break;
  }
}

/* ------------------------------------------------------------------ */
/* In-flight slots (concurrency, not per-window)                       */
/* ------------------------------------------------------------------ */

/**
 * Per-window limits alone let a client fire every allowed request at once:
 * ten parallel model calls still hammer Ollama (and hold ten sockets). These
 * counters cap what is running *right now*. Kept on globalThis for the same
 * reason as the buckets.
 */
const slotStore = globalThis as unknown as { __freezmoSlots?: Map<string, number> };
const slots: Map<string, number> = (slotStore.__freezmoSlots ??= new Map());

/** Take one of `max` in-flight slots for `key`. */
export function acquireSlot(key: string, max: number): boolean {
  const current = slots.get(key) ?? 0;
  if (current >= max) return false;
  slots.set(key, current + 1);
  return true;
}

/** Give an in-flight slot back. */
export function releaseSlot(key: string): void {
  const current = slots.get(key) ?? 0;
  if (current <= 1) slots.delete(key);
  else slots.set(key, current - 1);
}

/** How many slots are in flight for a key (diagnostics and tests). */
export function inFlight(key: string): number {
  return slots.get(key) ?? 0;
}

/** Best-effort client identity: proxy headers first, then a shared bucket. */
export function clientKey(request: Request): string {
  const forwarded = request.headers.get("x-forwarded-for");
  if (forwarded) {
    // Take the **last** hop, not the first: our reverse proxy appends the peer
    // it saw, so the last entry is the one connection we can trust. Reading the
    // first entry would let a client send its own `x-forwarded-for` and dodge
    // every per-IP limit.
    const hops = forwarded
      .split(",")
      .map((entry) => entry.trim())
      .filter(Boolean);
    const last = hops.at(-1);
    if (last) return last;
  }
  return (
    request.headers.get("x-real-ip") ??
    request.headers.get("cf-connecting-ip") ??
    "local"
  );
}

export function rateLimitResponse(result: RateLimitResult, locale: Locale = DEFAULT_LOCALE): Response {
  return Response.json(
    {
      error: translate(locale, "server.rateLimited", { limit: result.limit, retry: result.retryAfter }),
      kind: "rate_limited",
      retryAfter: result.retryAfter,
    },
    {
      status: 429,
      headers: {
        "Retry-After": String(result.retryAfter),
        "Cache-Control": "no-store",
      },
    },
  );
}