spb/company-atlas

1"""In-process token-bucket rate limiting (spec §136: tiers anonymous / authenticated / paid / internal).23Key = client IP (first hop of `X-Forwarded-For`) for anonymous traffic, or the API key id when a valid `X-CA-API-Key` is presented.4API keys are looked up by sha256 hash in `api_keys` (60 s in-process cache, negative results too); `last_used_at`/`request_count`5are flushed lazily every `FLUSH_EVERY_S`. Admin requests (valid `X-CA-Admin-Token`) and health/docs paths bypass the limiter.6`429` carries `Retry-After`; every limited response carries `X-RateLimit-Limit` / `X-RateLimit-Remaining` / `X-RateLimit-Tier`.7"""8from __future__ import annotations910import asyncio11import hashlib12import logging13import time14from dataclasses import dataclass, field15from datetime import UTC, datetime16from typing import Any1718from fastapi import Request1920from companyatlas.api.common import AtlasJSONResponse, TTLCache, client_ip, is_admin_request21from companyatlas.db import connection, execute, fetch_one, transaction2223log = logging.getLogger("companyatlas.api.ratelimit")2425TIER_LIMITS_PER_MIN: dict[str, int | None] = {"anonymous": 120, "authenticated": 600, "paid": 3000, "internal": None}26KEY_CACHE_TTL_S = 60.027FLUSH_EVERY_S = 30.028BYPASS_PREFIXES = ("/health", "/ready", "/api/v1/health", "/api/v1/docs", "/api/v1/openapi.json")29MAX_BUCKETS = 50_000303132@dataclass(slots=True)33class Bucket:34    tokens: float35    updated: float363738@dataclass39class RateLimiter:40    limits: dict[str, int | None] = field(default_factory=lambda: dict(TIER_LIMITS_PER_MIN))41    buckets: dict[str, Bucket] = field(default_factory=dict)4243    def take(self, key: str, tier: str, now: float | None = None) -> tuple[bool, int, int, float]:44        """Consume one token. Returns (allowed, limit, remaining, retry_after_s)."""45        limit = self.limits.get(tier, self.limits["anonymous"])46        if limit is None:47            return True, 0, 0, 0.048        now = time.monotonic() if now is None else now49        rate = limit / 60.050        b = self.buckets.get(key)51        if b is None:52            if len(self.buckets) >= MAX_BUCKETS:53                self._evict(now)54            b = self.buckets[key] = Bucket(tokens=float(limit), updated=now)55        else:56            b.tokens = min(float(limit), b.tokens + (now - b.updated) * rate)57            b.updated = now58        if b.tokens >= 1.0:59            b.tokens -= 1.060            return True, limit, int(b.tokens), 0.061        retry = (1.0 - b.tokens) / rate62        return False, limit, 0, retry6364    def _evict(self, now: float) -> None:65        stale = [k for k, b in self.buckets.items() if now - b.updated > 120]66        for k in stale:67            self.buckets.pop(k, None)68        if len(self.buckets) >= MAX_BUCKETS:      # still full: drop the oldest half69            for k in sorted(self.buckets, key=lambda k: self.buckets[k].updated)[: MAX_BUCKETS // 2]:70                self.buckets.pop(k, None)717273limiter = RateLimiter()74_key_cache = TTLCache(max_items=10_000)75_usage: dict[str, int] = {}76_usage_lock = asyncio.Lock()77_last_flush = time.monotonic()787980def hash_key(raw: str) -> str:81    return hashlib.sha256(raw.encode("utf-8")).hexdigest()828384async def resolve_api_key(raw: str | None) -> dict[str, Any] | None:85    """`{id, tier, name}` for a valid, non-revoked key; None otherwise (cached 60 s either way)."""86    if not raw or len(raw) < 16 or len(raw) > 200:87        return None88    h = hash_key(raw)89    hit = _key_cache.get(h)90    if hit is not None:91        return hit or None92    try:93        async with connection() as conn:94            row = await fetch_one(conn, "select id, tier, name from api_keys where key_hash = :h and revoked_at is null", h=h)95    except Exception:96        log.warning("api key lookup failed", exc_info=True)97        return None98    info = {"id": row["id"], "tier": row["tier"], "name": row["name"]} if row else {}99    _key_cache.set(h, info, KEY_CACHE_TTL_S)100    return info or None101102103async def _write_usage(batch: dict[str, int]) -> None:104    if not batch:105        return106    try:107        async with transaction() as conn:108            for kid, n in batch.items():109                await execute(conn, "update api_keys set last_used_at = :t, request_count = request_count + :n where id = :id",110                              t=datetime.now(UTC), n=n, id=kid)111    except Exception:112        log.warning("api key usage flush failed", exc_info=True)113114115async def _note_usage(key_id: str) -> None:116    global _last_flush117    async with _usage_lock:118        _usage[key_id] = _usage.get(key_id, 0) + 1119        if time.monotonic() - _last_flush < FLUSH_EVERY_S:120            return121        batch, _last_flush = dict(_usage), time.monotonic()122        _usage.clear()123    await _write_usage(batch)124125126async def flush_usage() -> None:127    """Force a usage flush (tests / shutdown)."""128    global _last_flush129    async with _usage_lock:130        batch, _last_flush = dict(_usage), time.monotonic()131        _usage.clear()132    await _write_usage(batch)133134135def _is_local_server_side(request: Request) -> bool:136    if request.headers.get("x-forwarded-for") or request.headers.get("x-real-ip"):137        return False138    host = request.client.host if request.client else ""139    return host in ("127.0.0.1", "::1")140141142async def rate_limit_middleware(request: Request, call_next):  # type: ignore[no-untyped-def]143    path = request.url.path144    if request.method == "OPTIONS" or path.startswith(BYPASS_PREFIXES):145        return await call_next(request)146    if is_admin_request(request):147        response = await call_next(request)148        response.headers["x-ratelimit-tier"] = "admin"149        return response150    key_info = await resolve_api_key(request.headers.get("x-ca-api-key"))151    if key_info:152        tier, bucket_key = key_info["tier"], f"key:{key_info['id']}"153        asyncio.get_running_loop().create_task(_note_usage(key_info["id"]))154    elif _is_local_server_side(request):155        # Next.js server components fetch the loopback API without X-Forwarded-For: that is our own SSR, not a public client.156        # Browser traffic proxied through the Next rewrite (and through Caddy) always carries X-Forwarded-For and stays anonymous.157        tier, bucket_key = "internal", "internal:ssr"158    else:159        tier, bucket_key = "anonymous", f"ip:{client_ip(request)}"160    allowed, limit, remaining, retry = limiter.take(bucket_key, tier)161    if not allowed:162        return AtlasJSONResponse({"detail": "rate limit exceeded"}, status_code=429,163                                 headers={"retry-after": str(max(1, int(retry + 0.999))), "x-ratelimit-limit": str(limit),164                                          "x-ratelimit-remaining": "0", "x-ratelimit-tier": tier, "cache-control": "no-store"})165    response = await call_next(request)166    response.headers["x-ratelimit-tier"] = tier167    if limit:168        response.headers["x-ratelimit-limit"] = str(limit)169        response.headers["x-ratelimit-remaining"] = str(remaining)170    return response171172173__all__ = ["TIER_LIMITS_PER_MIN", "RateLimiter", "flush_usage", "hash_key", "limiter", "rate_limit_middleware", "resolve_api_key"]174