Spaces:
Running
Running
File size: 4,070 Bytes
889e86b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """
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
|