src/middleware/rateLimiter.ts

import { IncomingMessage } from 'http';

/**
 * Simple in-process token bucket rate limiter for WS upgrade requests.
 * Prevents a single IP from opening thousands of connections on HF free tier.
 *
 * Limits: 20 connections per IP per 60 seconds.
 */

interface Bucket {
  tokens: number;
  lastRefill: number;
}

const CAPACITY = 20;
const REFILL_RATE = 20;      // tokens per window
const WINDOW_MS = 60_000;    // 60 seconds

const buckets = new Map<string, Bucket>();

// Prune old entries every 5 minutes to prevent memory leak
setInterval(() => {
  const cutoff = Date.now() - WINDOW_MS * 2;
  for (const [ip, bucket] of buckets) {
    if (bucket.lastRefill < cutoff) buckets.delete(ip);
  }
}, 300_000).unref();

export function isRateLimited(req: IncomingMessage): boolean {
  const ip =
    (req.headers['x-forwarded-for'] as string | undefined)?.split(',')[0]?.trim() ??
    req.socket.remoteAddress ??
    'unknown';

  const now = Date.now();
  let bucket = buckets.get(ip);

  if (!bucket) {
    bucket = { tokens: CAPACITY - 1, lastRefill: now };
    buckets.set(ip, bucket);
    return false;
  }

  // Refill based on elapsed time
  const elapsed = now - bucket.lastRefill;
  if (elapsed >= WINDOW_MS) {
    bucket.tokens = CAPACITY - 1;
    bucket.lastRefill = now;
    return false;
  }

  if (bucket.tokens <= 0) return true;

  bucket.tokens--;
  return false;
}