"""In-process token-bucket rate limiting (spec ยง136: tiers anonymous / authenticated / paid / internal). Key = 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. API keys are looked up by sha256 hash in `api_keys` (60 s in-process cache, negative results too); `last_used_at`/`request_count` are flushed lazily every `FLUSH_EVERY_S`. Admin requests (valid `X-CA-Admin-Token`) and health/docs paths bypass the limiter. `429` carries `Retry-After`; every limited response carries `X-RateLimit-Limit` / `X-RateLimit-Remaining` / `X-RateLimit-Tier`. """ from __future__ import annotations import asyncio import hashlib import logging import time from dataclasses import dataclass, field from datetime import UTC, datetime from typing import Any from fastapi import Request from companyatlas.api.common import AtlasJSONResponse, TTLCache, client_ip, is_admin_request from companyatlas.db import connection, execute, fetch_one, transaction log = logging.getLogger("companyatlas.api.ratelimit") TIER_LIMITS_PER_MIN: dict[str, int | None] = {"anonymous": 120, "authenticated": 600, "paid": 3000, "internal": None} KEY_CACHE_TTL_S = 60.0 FLUSH_EVERY_S = 30.0 BYPASS_PREFIXES = ("/health", "/ready", "/api/v1/health", "/api/v1/docs", "/api/v1/openapi.json") MAX_BUCKETS = 50_000 @dataclass(slots=True) class Bucket: tokens: float updated: float @dataclass class RateLimiter: limits: dict[str, int | None] = field(default_factory=lambda: dict(TIER_LIMITS_PER_MIN)) buckets: dict[str, Bucket] = field(default_factory=dict) def take(self, key: str, tier: str, now: float | None = None) -> tuple[bool, int, int, float]: """Consume one token. Returns (allowed, limit, remaining, retry_after_s).""" limit = self.limits.get(tier, self.limits["anonymous"]) if limit is None: return True, 0, 0, 0.0 now = time.monotonic() if now is None else now rate = limit / 60.0 b = self.buckets.get(key) if b is None: if len(self.buckets) >= MAX_BUCKETS: self._evict(now) b = self.buckets[key] = Bucket(tokens=float(limit), updated=now) else: b.tokens = min(float(limit), b.tokens + (now - b.updated) * rate) b.updated = now if b.tokens >= 1.0: b.tokens -= 1.0 return True, limit, int(b.tokens), 0.0 retry = (1.0 - b.tokens) / rate return False, limit, 0, retry def _evict(self, now: float) -> None: stale = [k for k, b in self.buckets.items() if now - b.updated > 120] for k in stale: self.buckets.pop(k, None) if len(self.buckets) >= MAX_BUCKETS: # still full: drop the oldest half for k in sorted(self.buckets, key=lambda k: self.buckets[k].updated)[: MAX_BUCKETS // 2]: self.buckets.pop(k, None) limiter = RateLimiter() _key_cache = TTLCache(max_items=10_000) _usage: dict[str, int] = {} _usage_lock = asyncio.Lock() _last_flush = time.monotonic() def hash_key(raw: str) -> str: return hashlib.sha256(raw.encode("utf-8")).hexdigest() async def resolve_api_key(raw: str | None) -> dict[str, Any] | None: """`{id, tier, name}` for a valid, non-revoked key; None otherwise (cached 60 s either way).""" if not raw or len(raw) < 16 or len(raw) > 200: return None h = hash_key(raw) hit = _key_cache.get(h) if hit is not None: return hit or None try: async with connection() as conn: row = await fetch_one(conn, "select id, tier, name from api_keys where key_hash = :h and revoked_at is null", h=h) except Exception: log.warning("api key lookup failed", exc_info=True) return None info = {"id": row["id"], "tier": row["tier"], "name": row["name"]} if row else {} _key_cache.set(h, info, KEY_CACHE_TTL_S) return info or None async def _write_usage(batch: dict[str, int]) -> None: if not batch: return try: async with transaction() as conn: for kid, n in batch.items(): await execute(conn, "update api_keys set last_used_at = :t, request_count = request_count + :n where id = :id", t=datetime.now(UTC), n=n, id=kid) except Exception: log.warning("api key usage flush failed", exc_info=True) async def _note_usage(key_id: str) -> None: global _last_flush async with _usage_lock: _usage[key_id] = _usage.get(key_id, 0) + 1 if time.monotonic() - _last_flush < FLUSH_EVERY_S: return batch, _last_flush = dict(_usage), time.monotonic() _usage.clear() await _write_usage(batch) async def flush_usage() -> None: """Force a usage flush (tests / shutdown).""" global _last_flush async with _usage_lock: batch, _last_flush = dict(_usage), time.monotonic() _usage.clear() await _write_usage(batch) def _is_local_server_side(request: Request) -> bool: if request.headers.get("x-forwarded-for") or request.headers.get("x-real-ip"): return False host = request.client.host if request.client else "" return host in ("127.0.0.1", "::1") async def rate_limit_middleware(request: Request, call_next): # type: ignore[no-untyped-def] path = request.url.path if request.method == "OPTIONS" or path.startswith(BYPASS_PREFIXES): return await call_next(request) if is_admin_request(request): response = await call_next(request) response.headers["x-ratelimit-tier"] = "admin" return response key_info = await resolve_api_key(request.headers.get("x-ca-api-key")) if key_info: tier, bucket_key = key_info["tier"], f"key:{key_info['id']}" asyncio.get_running_loop().create_task(_note_usage(key_info["id"])) elif _is_local_server_side(request): # Next.js server components fetch the loopback API without X-Forwarded-For: that is our own SSR, not a public client. # Browser traffic proxied through the Next rewrite (and through Caddy) always carries X-Forwarded-For and stays anonymous. tier, bucket_key = "internal", "internal:ssr" else: tier, bucket_key = "anonymous", f"ip:{client_ip(request)}" allowed, limit, remaining, retry = limiter.take(bucket_key, tier) if not allowed: return AtlasJSONResponse({"detail": "rate limit exceeded"}, status_code=429, headers={"retry-after": str(max(1, int(retry + 0.999))), "x-ratelimit-limit": str(limit), "x-ratelimit-remaining": "0", "x-ratelimit-tier": tier, "cache-control": "no-store"}) response = await call_next(request) response.headers["x-ratelimit-tier"] = tier if limit: response.headers["x-ratelimit-limit"] = str(limit) response.headers["x-ratelimit-remaining"] = str(remaining) return response __all__ = ["TIER_LIMITS_PER_MIN", "RateLimiter", "flush_usage", "hash_key", "limiter", "rate_limit_middleware", "resolve_api_key"]