app/core/rate_limit.py

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"},
        )