75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
"""Redis-first rate limiter with in-memory fallback."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
from fastapi import HTTPException, Request, status
|
|
from redis import Redis
|
|
from redis.exceptions import RedisError
|
|
|
|
from app.core.config import settings
|
|
|
|
_buckets: dict[str, list[datetime]] = defaultdict(list)
|
|
_redis_client: Redis | None = None
|
|
|
|
|
|
def get_redis_client() -> Redis | None:
|
|
global _redis_client
|
|
if _redis_client is not None:
|
|
return _redis_client
|
|
try:
|
|
_redis_client = Redis.from_url(settings.redis_url, decode_responses=True)
|
|
_redis_client.ping()
|
|
return _redis_client
|
|
except RedisError:
|
|
_redis_client = None
|
|
return None
|
|
|
|
|
|
def check_rate_limit(key: str, limit: int, window_seconds: int) -> None:
|
|
if not settings.enable_rate_limit:
|
|
return
|
|
redis_client = get_redis_client()
|
|
if redis_client is not None:
|
|
redis_key = f"rl:{key}"
|
|
try:
|
|
current = redis_client.incr(redis_key)
|
|
if current == 1:
|
|
redis_client.expire(redis_key, window_seconds)
|
|
if current > limit:
|
|
ttl = redis_client.ttl(redis_key)
|
|
retry_after = ttl if ttl and ttl > 0 else window_seconds
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="RATE_LIMIT_EXCEEDED",
|
|
headers={"Retry-After": str(retry_after)},
|
|
)
|
|
return
|
|
except RedisError:
|
|
pass
|
|
|
|
now = datetime.now(UTC)
|
|
cutoff = now - timedelta(seconds=window_seconds)
|
|
timestamps = [moment for moment in _buckets[key] if moment > cutoff]
|
|
if len(timestamps) >= limit:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_429_TOO_MANY_REQUESTS,
|
|
detail="RATE_LIMIT_EXCEEDED",
|
|
headers={"Retry-After": str(window_seconds)},
|
|
)
|
|
timestamps.append(now)
|
|
_buckets[key] = timestamps
|
|
|
|
|
|
def client_ip(request: Request) -> str:
|
|
trusted_proxy_ips = {item.strip() for item in settings.trusted_proxy_ips.split(",") if item.strip()}
|
|
forwarded = request.headers.get("X-Forwarded-For")
|
|
request_ip = request.client.host if request.client else ""
|
|
if forwarded and request_ip in trusted_proxy_ips:
|
|
return forwarded.split(",")[0].strip()
|
|
if request.client:
|
|
return request_ip
|
|
return "unknown"
|