vlad
86cc3fa541
Align the project baseline with the latest admin interface styling and layout structure while documenting setup and usage updates in README.
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"
|