Spaces:
Running
Running
| """ | |
| Per-UID rate limiter — design settled 2026-09-09 (see | |
| docs/blueprint/adr/0010-per-uid-in-memory-rate-limiting.md), built 2026-09-25. | |
| Keys by authenticated Firebase UID, not IP — every limited route already | |
| requires a verified token, and a UID can't be spoofed the way an IP/header | |
| can. Counters are in-memory, per-process: safe from client/devtools | |
| tampering by construction, sufficient for the current single-container | |
| Hugging Face Space deployment, and reset on redeploy. If this ever needs to | |
| survive across multiple replicas, move the counters into Firestore/Redis — | |
| the enforce_rate_limit() call site wouldn't need to change. | |
| """ | |
| import time | |
| import threading | |
| from collections import defaultdict | |
| from fastapi import Header, HTTPException | |
| from firebase_admin import firestore | |
| from auth import verify_firebase_token | |
| # (label, limit, period_seconds) per tier. Tiers are named by cost, not by | |
| # route, so multiple routes can share one (see routes/*.py's with_rate_limit | |
| # calls for which tier each route uses). | |
| TIERS: dict[str, list[tuple[str, int, int]]] = { | |
| "chat": [("minute", 10, 60), ("day", 200, 86400)], | |
| "generation": [("hour", 8, 3600), ("day", 30, 86400)], | |
| "transcription": [("day", 20, 86400)], | |
| } | |
| class _Window: | |
| __slots__ = ("count", "reset_at") | |
| def __init__(self) -> None: | |
| self.count = 0 | |
| self.reset_at = 0.0 | |
| _lock = threading.Lock() | |
| _counters: dict[tuple[str, str, str], _Window] = defaultdict(_Window) | |
| # Consecutive-violation count per UID, purely for the "log repeat offenders" | |
| # requirement — never used to make an enforcement decision itself. | |
| _violation_counts: dict[str, int] = defaultdict(int) | |
| def _log_violation(uid: str, tier: str, label: str, count: int, limit: int) -> None: | |
| _violation_counts[uid] += 1 | |
| print( | |
| f"[rate-limit] uid={uid} tier={tier} window={label} " | |
| f"count={count}/{limit} (consecutive violations: {_violation_counts[uid]})" | |
| ) | |
| # Best-effort — a Firestore write failing must never break enforcement itself. | |
| try: | |
| firestore.client().collection("rate_limit_violations").add({ | |
| "uid": uid, | |
| "tier": tier, | |
| "window": label, | |
| "count": count, | |
| "limit": limit, | |
| "consecutiveViolations": _violation_counts[uid], | |
| "at": firestore.SERVER_TIMESTAMP, | |
| }) | |
| except Exception: | |
| pass | |
| def enforce_rate_limit(uid: str, tier: str) -> None: | |
| """Raises HTTPException(429) if uid has exceeded any of tier's limits.""" | |
| if tier not in TIERS: | |
| raise ValueError(f"Unknown rate limit tier: {tier!r}") | |
| now = time.time() | |
| with _lock: | |
| for label, limit, period in TIERS[tier]: | |
| window = _counters[(uid, tier, label)] | |
| if now >= window.reset_at: | |
| window.count = 0 | |
| window.reset_at = now + period | |
| window.count += 1 | |
| if window.count > limit: | |
| retry_after = max(1, int(window.reset_at - now)) | |
| _log_violation(uid, tier, label, window.count, limit) | |
| raise HTTPException( | |
| status_code=429, | |
| detail=( | |
| f"You've hit the {label}ly limit for this action " | |
| f"({limit} per {label}). Try again in {retry_after} seconds." | |
| ), | |
| headers={"Retry-After": str(retry_after)}, | |
| ) | |
| def with_rate_limit(tier: str): | |
| """ | |
| FastAPI dependency: verifies the caller's token AND enforces `tier`'s rate | |
| limit in one pass. Use this in place of `Depends(verify_firebase_token)` | |
| on any route that should be limited — it calls verify_firebase_token | |
| itself rather than composing as a second dependency, so the token is only | |
| verified once per request. | |
| """ | |
| async def dependency(authorization: str = Header(...)) -> dict: | |
| user = await verify_firebase_token(authorization) | |
| enforce_rate_limit(user["uid"], tier) | |
| return user | |
| return dependency | |