"""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"