Calbook/lib/rate-limit.ts

import { isIP } from "net";

type RateLimitEntry = {
  count: number;
  resetAt: number;
};

type HeaderBag = Headers | Record<string, string | string[] | undefined>;

type EnforceRateLimitInput = {
  req: Request;
  scope: string;
  limit: number;
  windowMs: number;
  keySuffix?: string;
};

type RateLimitResult = {
  ok: boolean;
  remaining: number;
  retryAfterSeconds: number;
};

declare global {
  // eslint-disable-next-line no-var
  var calbookRateLimitStore: Map<string, RateLimitEntry> | undefined;
}

const store = global.calbookRateLimitStore ?? new Map<string, RateLimitEntry>();
if (!global.calbookRateLimitStore) {
  global.calbookRateLimitStore = store;
}

const CLEANUP_EVERY = 200;
let writes = 0;

function nowMs() {
  return Date.now();
}

function cleanupExpired() {
  const now = nowMs();
  for (const [key, value] of store.entries()) {
    if (value.resetAt <= now) {
      store.delete(key);
    }
  }
}

function readHeader(headers: HeaderBag, key: string): string | undefined {
  if (headers instanceof Headers) {
    return headers.get(key) ?? undefined;
  }
  const value = headers[key] ?? headers[key.toLowerCase()] ?? headers[key.toUpperCase()];
  if (Array.isArray(value)) return value[0];
  return value ?? undefined;
}

function trustProxyHeaders() {
  const raw = (process.env.TRUST_PROXY_HEADERS ?? "").trim().toLowerCase();
  return raw === "1" || raw === "true" || raw === "yes";
}

function normalizeIpCandidate(value: string): string | null {
  let candidate = value.trim().replace(/^for=/i, "").replace(/^"|"$/g, "");
  if (!candidate) return null;

  if (candidate.startsWith("[")) {
    const end = candidate.indexOf("]");
    if (end > 1) {
      candidate = candidate.slice(1, end);
    }
  }

  if (candidate.includes("%")) {
    candidate = candidate.split("%")[0] ?? candidate;
  }

  if (candidate.includes(".") && candidate.includes(":") && isIP(candidate) === 0) {
    const withoutPort = candidate.split(":")[0];
    if (withoutPort) {
      candidate = withoutPort;
    }
  }

  return isIP(candidate) ? candidate : null;
}

export function getClientIpFromHeaders(headers: HeaderBag): string {
  if (!trustProxyHeaders()) {
    return "proxy-untrusted";
  }

  const forwarded = readHeader(headers, "x-forwarded-for");
  if (forwarded) {
    const parsed = forwarded
      .split(",")
      .map((part) => normalizeIpCandidate(part))
      .find((value): value is string => Boolean(value));
    if (parsed) return parsed;
  }

  const realIp = readHeader(headers, "x-real-ip");
  if (realIp) {
    const parsed = normalizeIpCandidate(realIp);
    if (parsed) return parsed;
  }

  const cfIp = readHeader(headers, "cf-connecting-ip");
  if (cfIp) {
    const parsed = normalizeIpCandidate(cfIp);
    if (parsed) return parsed;
  }

  return "unknown";
}

export function getClientIp(req: Request): string {
  return getClientIpFromHeaders(req.headers);
}

export function consumeRateLimit(
  key: string,
  limit: number,
  windowMs: number
): RateLimitResult {
  writes += 1;
  if (writes % CLEANUP_EVERY === 0) {
    cleanupExpired();
  }

  const now = nowMs();
  const existing = store.get(key);

  if (!existing || existing.resetAt <= now) {
    store.set(key, {
      count: 1,
      resetAt: now + windowMs
    });
    return {
      ok: true,
      remaining: Math.max(0, limit - 1),
      retryAfterSeconds: Math.ceil(windowMs / 1000)
    };
  }

  existing.count += 1;
  store.set(key, existing);

  if (existing.count > limit) {
    return {
      ok: false,
      remaining: 0,
      retryAfterSeconds: Math.max(1, Math.ceil((existing.resetAt - now) / 1000))
    };
  }

  return {
    ok: true,
    remaining: Math.max(0, limit - existing.count),
    retryAfterSeconds: Math.max(1, Math.ceil((existing.resetAt - now) / 1000))
  };
}

export function enforceRateLimit(input: EnforceRateLimitInput): RateLimitResult {
  const ip = getClientIp(input.req);
  const suffix = input.keySuffix ? `:${input.keySuffix}` : "";
  const key = `${input.scope}:${ip}${suffix}`;
  return consumeRateLimit(key, input.limit, input.windowMs);
}