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