from __future__ import annotations import ipaddress import json import structlog import time from typing import Optional, Tuple from uuid import uuid4 import redis.asyncio as aioredis from redis.asyncio.retry import Retry from redis.backoff import ExponentialBackoff from redis.exceptions import ( BusyLoadingError, ConnectionError as RedisConnectionError, NoScriptError, RedisError, TimeoutError as RedisTimeoutError, ) from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint from starlette.requests import Request from starlette.responses import Response from app.core.config import settings logger = structlog.get_logger(__name__) _EXEMPT_PATHS: frozenset[str] = frozenset({ "/api/v1/health", "/api/v1/ready", "/api/v1/ping", "/docs", "/redoc", "/openapi.json", "/metrics", }) _TRUSTED_PROXY_NETWORKS: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = ( ipaddress.ip_network("10.0.0.0/8"), ipaddress.ip_network("172.16.0.0/12"), ipaddress.ip_network("192.168.0.0/16"), ipaddress.ip_network("127.0.0.0/8"), ipaddress.ip_network("fc00::/7"), ipaddress.ip_network("::1/128"), ) _RATE_LIMIT_SCRIPT = """ local key = KEYS[1] local now = tonumber(ARGV[1]) local window = tonumber(ARGV[2]) local limit = tonumber(ARGV[3]) local unique_id = ARGV[4] local cutoff = now - window redis.call('zremrangebyscore', key, '-inf', cutoff) local count = redis.call('zcard', key) if count < limit then redis.call('zadd', key, now, unique_id) redis.call('expire', key, window + 1) return {0, limit - count - 1} end redis.call('expire', key, window + 1) return {1, 0} """ def _is_trusted_proxy(host: str) -> bool: try: addr = ipaddress.ip_address(host) return any(addr in net for net in _TRUSTED_PROXY_NETWORKS) except ValueError: return False def _extract_client_ip(request: Request) -> str: direct_host = request.client.host if request.client else None if not direct_host or not _is_trusted_proxy(direct_host): return direct_host or "unknown" forwarded_for = request.headers.get("x-forwarded-for", "") if forwarded_for: # Trust the rightmost non-proxy IP; that's the closest hop to us. for candidate in (ip.strip() for ip in reversed(forwarded_for.split(","))): if candidate and not _is_trusted_proxy(candidate): return candidate real_ip = request.headers.get("x-real-ip", "").strip() if real_ip and not _is_trusted_proxy(real_ip): return real_ip return direct_host class RateLimitMiddleware(BaseHTTPMiddleware): def __init__(self, app, **kwargs) -> None: super().__init__(app, **kwargs) self._pool: Optional[aioredis.ConnectionPool] = None self._client: Optional[aioredis.Redis] = None self._script_sha: Optional[str] = None def _build_pool(self) -> aioredis.ConnectionPool: retry = Retry( ExponentialBackoff(cap=0.5, base=0.1), retries=2, supported_errors=(BusyLoadingError, RedisConnectionError, RedisTimeoutError), ) return aioredis.ConnectionPool.from_url( settings.REDIS_URL, encoding="utf-8", decode_responses=True, socket_connect_timeout=1, socket_timeout=1, socket_keepalive=True, retry=retry, retry_on_timeout=True, max_connections=max(settings.WORKERS * 2, 10), health_check_interval=30, ) async def _get_client(self) -> aioredis.Redis: if self._pool is None: self._pool = self._build_pool() if self._client is None: self._client = aioredis.Redis(connection_pool=self._pool) return self._client async def _load_script(self, client: aioredis.Redis) -> str: if self._script_sha is None: self._script_sha = await client.script_load(_RATE_LIMIT_SCRIPT) return self._script_sha async def _check_rate_limit(self, client_ip: str) -> Tuple[bool, int]: now_ms = int(time.time() * 1000) window_ms = settings.RATE_LIMIT_WINDOW * 1000 key = f"rl:{client_ip}" unique_id = f"{now_ms}:{uuid4().hex}" client = await self._get_client() try: sha = await self._load_script(client) result = await client.evalsha( sha, 1, key, str(now_ms), str(window_ms), str(settings.RATE_LIMIT), unique_id, ) except NoScriptError: self._script_sha = None sha = await self._load_script(client) result = await client.evalsha( sha, 1, key, str(now_ms), str(window_ms), str(settings.RATE_LIMIT), unique_id, ) return bool(result[0]), int(result[1]) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: if request.url.path in _EXEMPT_PATHS: return await call_next(request) if not settings.rate_limit_redis_configured: return await call_next(request) client_ip = _extract_client_ip(request) try: is_limited, remaining = await self._check_rate_limit(client_ip) except (RedisConnectionError, RedisTimeoutError) as exc: logger.error("redis_unavailable", error=str(exc)) return self._service_unavailable_response() except RedisError as exc: logger.error("redis_error", error=str(exc)) return self._service_unavailable_response() if is_limited: request_id = getattr(request.state, "request_id", str(uuid4())) return self._rate_limit_response(request_id) response = await call_next(request) response.headers["X-RateLimit-Limit"] = str(settings.RATE_LIMIT) response.headers["X-RateLimit-Remaining"] = str(remaining) response.headers["X-RateLimit-Window"] = str(settings.RATE_LIMIT_WINDOW) return response def _rate_limit_response(self, request_id: str) -> Response: return Response( content=json.dumps({ "error": "Too many requests. Please slow down.", "request_id": request_id, }), status_code=429, media_type="application/json", headers={ "Retry-After": str(settings.RATE_LIMIT_WINDOW), "X-RateLimit-Limit": str(settings.RATE_LIMIT), "X-RateLimit-Remaining": "0", "X-RateLimit-Window": str(settings.RATE_LIMIT_WINDOW), }, ) def _service_unavailable_response(self) -> Response: return Response( content=json.dumps( {"error": "Service temporarily unavailable. Please retry shortly."} ), status_code=503, media_type="application/json", headers={"Retry-After": "5"}, )