api/middleware/rate_limiter.py

import os, sys, time, hashlib
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from collections import defaultdict
from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from loguru import logger

LIMITS = {
    "/search":    (100, 60),
    "/profile":   (100, 60),
    "/risk":      (100, 60),
    "/graph":     (50,  60),
    "/investigate":(30, 60),
    "/export":    (10,  60),
    "/translate": (10,  60),
    "/health":    (1000,60),
    "/admin":     (5,   60),
    "default":    (200, 60),
}


class SlidingWindowRateLimiter(BaseHTTPMiddleware):

    def __init__(self, app):
        super().__init__(app)
        # BUG-14 NOTE: in-memory state -- each worker starts fresh on restart.
        # Rate limits are per-worker, not cross-worker. Acceptable for single-process HF Spaces.
        # For multi-worker: replace with Redis-backed sliding window.
        self._windows: dict[str, list[float]] = defaultdict(list)
        # BUG-14 FIX: add eviction to prevent unbounded memory growth on long-running server
        self._evict_at = 0.0

    def _get_limit(self, path: str) -> tuple[int, int]:
        for prefix, limit in LIMITS.items():
            if prefix != "default" and path.startswith(prefix):
                return limit
        return LIMITS["default"]

    def _get_ip(self, request: Request) -> str:
        forwarded = request.headers.get("X-Forwarded-For", "")
        raw = forwarded.split(",")[0].strip() if forwarded else (
            request.client.host if request.client else "unknown"
        )
        return hashlib.sha256(raw.encode()).hexdigest()[:16]

    async def dispatch(self, request: Request, call_next):
        ip  = self._get_ip(request)
        path = request.url.path
        max_req, window = self._get_limit(path)
        key = f"{ip}:{path.split('/')[1]}"
        now = time.time()

        # BACKEND-1 FIX: prune current key FIRST, THEN evict stale entries.
        # Previous order evicted before pruning -- lost tracking on burst-after-idle keys.
        self._windows[key] = [t for t in self._windows[key] if now - t < window]

        # Evict stale entries every 5 minutes to prevent memory leak
        if now - self._evict_at > 300:
            stale_keys = [k for k, ts in self._windows.items() if not ts or now - ts[0] > window]
            for k in stale_keys:
                del self._windows[k]
            self._evict_at = now

        if len(self._windows[key]) >= max_req:
            retry = int(window - (now - self._windows[key][0]))
            logger.warning(f"[RateLimit] {ip} exceeded {max_req}/min on {path}")
            raise HTTPException(
                status_code=429,
                detail=f"Rate limit exceeded. Retry after {retry} seconds.",
                headers={"Retry-After": str(retry)},
            )

        self._windows[key].append(now)
        return await call_next(request)