siftlode/backend/app/ratelimit.py

"""A tiny in-process sliding-window rate limiter.

Generic groundwork: keyed by an arbitrary string (e.g. a client — `client_key`), so it can throttle any
endpoint. Single-worker uvicorn → in-memory state is sufficient; it resets on restart, which
is fine for abuse throttling (not for anything that must survive a deploy). Not shared across
processes — if we ever run multiple workers, swap the backing store for Redis behind the same
``allow()`` interface.
"""

import ipaddress
import threading
import time

from fastapi import Request

from app.config import settings


class RateLimiter:
    """At most `max_events` attempts per `window_seconds` per key — or, when each attempt passes a
    `cost`, at most that much total cost (a byte budget counts bytes, not calls)."""

    def __init__(self, max_events: float, window_seconds: float):
        self.max_events = max_events
        self.window = window_seconds
        # key → (monotonic time, cost) of every counted attempt still inside the window
        self._hits: dict[str, list[tuple[float, float]]] = {}
        self._lock = threading.Lock()

    def reset(self) -> None:
        """Forget every recorded attempt. Each limiter is a module-level singleton, so a test suite
        sharing one process would otherwise pool all of its attempts into a single window and
        eventually trip an unrelated test with a 429."""
        with self._lock:
            self._hits.clear()

    def exhausted(self, key: str) -> bool:
        """Whether ``key``'s window is already full (``allow`` would refuse it), WITHOUT counting an
        attempt — for a limit that counts only some outcomes: ask it first, ``allow`` the ones that
        count."""
        cutoff = time.monotonic() - self.window
        with self._lock:
            return sum(c for t, c in self._hits.get(key, []) if t > cutoff) >= self.max_events

    def allow(self, key: str, cost: float = 1) -> bool:
        """Record an attempt costing ``cost`` for ``key`` and return whether it is within the limit.

        True  -> within the cap (the attempt is counted).
        False -> the attempt would pass the cap for the current window (attempt NOT counted, so a
                 blocked caller can't keep pushing the window forward)."""
        now = time.monotonic()
        cutoff = now - self.window
        with self._lock:
            hits = [(t, c) for t, c in self._hits.get(key, []) if t > cutoff]
            if sum(c for _, c in hits) + cost > self.max_events:
                self._hits[key] = hits
                return False
            hits.append((now, cost))
            self._hits[key] = hits
            # Opportunistic cleanup so abandoned keys don't accumulate unboundedly.
            if len(self._hits) > 4096:
                for k in list(self._hits):
                    fresh = [(t, c) for t, c in self._hits[k] if t > cutoff]
                    if fresh:
                        self._hits[k] = fresh
                    else:
                        del self._hits[k]
            return True


def _norm_ip(s: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None:
    """Parse an IP for comparison, unwrapping IPv4-mapped IPv6 (`::ffff:10.0.0.1` → `10.0.0.1`) so a
    plain-IPv4 allowlist still matches a peer that surfaces in mapped form. None for a non-IP string."""
    try:
        ip = ipaddress.ip_address(s.strip())
    except ValueError:
        return None
    return ip.ipv4_mapped if getattr(ip, "ipv4_mapped", None) else ip


def client_key(request: Request) -> str:
    """Best-effort client identity for rate limiting: its IP, or for IPv6 the /64 the address sits in — a
    subscriber is handed a whole /64, so keyed on the full address one client would have endless identities
    to spread its attempts over. X-Forwarded-For is client-controlled, so we trust it
    ONLY when the request actually arrived from a configured reverse proxy (settings.trusted_proxy_ips
    — the address the proxy connects FROM, e.g. the VPS Caddy's WireGuard peer IP). Our proxy APPENDS
    the client it saw to XFF, so the RIGHTMOST entry is the real client and can't be forged by a
    client pre-seeding the header. Any other peer (someone hitting the app port directly, bypassing the
    proxy) is untrusted — we use its real socket IP, so it can't spoof its rate-limit identity.

    SAFETY: this relies on `request.client` being the TRUE socket peer. uvicorn is started WITHOUT
    `--proxy-headers` (see entrypoint.sh) precisely so it never rewrites `request.client` from XFF —
    do NOT add that flag, or the trust check below would compare against a forgeable value."""
    peer = request.client.host if request.client else "unknown"
    client = peer
    trusted = {_norm_ip(p) for p in settings.trusted_proxy_ip_set} - {None}
    if _norm_ip(peer) in trusted:
        xff = request.headers.get("x-forwarded-for")
        parts = [p.strip() for p in (xff or "").split(",") if p.strip()]
        if parts:
            client = parts[-1]
    ip = _norm_ip(client)
    if isinstance(ip, ipaddress.IPv6Address):
        return str(ipaddress.IPv6Network((ip, 64), strict=False))
    return client