app/middleware.py

"""HTTP 中间件：速率限制等。

只对 /v1/* 业务路径限流，admin / health / docs 不限流。
超限返回 429 + Retry-After 头。
"""
from __future__ import annotations

import json
from typing import Awaitable, Callable

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from .services import rate_limit_store


def _extract_access_key_from_request(request: Request) -> str:
    auth = request.headers.get("authorization")
    if auth:
        a = auth.strip()
        if a.lower().startswith("bearer "):
            return a[7:].strip()
        if a.lower().startswith("basic "):
            return a[6:].strip()
        return a
    return request.headers.get("x-xtc-access-key") or request.headers.get("x-xtc-access-token") or ""


class RateLimitMiddleware(BaseHTTPMiddleware):
    """per-key / per-IP 速率限制。"""

    async def dispatch(
        self,
        request: Request,
        call_next: Callable[[Request], Awaitable[Response]],
    ) -> Response:
        path = request.url.path
        if not rate_limit_store.is_rate_limited_path(path):
            return await call_next(request)

        # OPTIONS 预检直接放行
        if request.method == "OPTIONS":
            return await call_next(request)

        access_key = _extract_access_key_from_request(request)
        client_ip = rate_limit_store.get_client_ip(request)

        allowed, reason, retry_after = rate_limit_store.check(
            access_key=access_key or None,
            client_ip=client_ip,
        )
        if not allowed:
            body = {
                "ok": False,
                "error": {
                    "code": "rate_limited",
                    "message": f"Rate limit exceeded ({reason}). Retry after {retry_after}s.",
                    "status": 429,
                    "retryable": True,
                    "hint": f"请等待 {retry_after} 秒后重试",
                    "retry_after": retry_after,
                    "scope": reason,
                },
            }
            return JSONResponse(
                status_code=429,
                content=body,
                headers={
                    "Retry-After": str(retry_after),
                    "X-RateLimit-Scope": reason or "",
                    "Access-Control-Allow-Origin": "*",
                    "Access-Control-Allow-Methods": "*",
                    "Access-Control-Allow-Headers": "*",
                },
            )

        # 把限流信息附加到响应头（便于客户端观察）
        response = await call_next(request)
        try:
            response.headers["X-RateLimit-Window"] = "60"
        except Exception:
            pass
        return response