DadaoTask/backend/app/middleware.py

"""中间件：请求上下文（request_id / 访问日志）与幂等键。

顺序：RequestContext → Idempotency → 路由。
request_id 必须最先建立，否则后面所有日志都没有 trace。
"""

from __future__ import annotations

import json
import time
import uuid
from typing import Any

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

from .config import get_settings
from .errors import ErrorCode
from .logging_setup import get_logger, request_id_ctx
from .security import decode_token
from .stores import dumps, get_store, idempotency_key

settings = get_settings()
logger = get_logger("app.access")

_IDEMPOTENT_METHODS = {"POST", "PUT", "PATCH", "DELETE"}
_MAX_CACHE_BYTES = 256 * 1024


class RequestContextMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        rid = request.headers.get("X-Request-Id") or uuid.uuid4().hex
        token = request_id_ctx.set(rid)
        started = time.perf_counter()
        status = 500
        try:
            response = await call_next(request)
            status = response.status_code
            response.headers["X-Request-Id"] = rid
            return response
        finally:
            elapsed_ms = round((time.perf_counter() - started) * 1000, 2)
            logger.info(
                "%s %s %s %sms",
                request.method,
                request.url.path,
                status,
                elapsed_ms,
                extra={
                    "extra_fields": {
                        "method": request.method,
                        "path": request.url.path,
                        "status": status,
                        "elapsed_ms": elapsed_ms,
                        "client": request.client.host if request.client else "-",
                    }
                },
            )
            request_id_ctx.reset(token)


class IdempotencyMiddleware(BaseHTTPMiddleware):
    """带 Idempotency-Key 的写请求，重复提交返回第一次的结果。

    限制（M0 已知）：同 key 并发请求存在竞态窗口，靠唯一键约束兜底留到业务层；
    只缓存 JSON 且 < 256KB 的响应；5xx 不缓存。
    """

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        if not settings.idempotency_enabled or request.method not in _IDEMPOTENT_METHODS:
            return await call_next(request)

        raw_key = request.headers.get("Idempotency-Key")
        if not raw_key:
            return await call_next(request)

        subject = self._subject(request)
        if subject is None:
            # 未认证请求不做幂等（登录接口另有验证码限流）
            return await call_next(request)

        store = get_store()
        cache_key = idempotency_key(subject, raw_key)
        try:
            cached = await store.get(cache_key)
        except Exception:  # noqa: BLE001 - 存储不可用不能拖垮业务
            logger.warning("幂等存储不可用，降级为直接执行", exc_info=True)
            return await call_next(request)

        if cached:
            payload = json.loads(cached)
            return JSONResponse(
                content=payload["body"],
                status_code=payload["status"],
                headers={"X-Idempotent-Replay": "1"},
            )

        response = await call_next(request)
        if response.status_code >= 500:
            return response

        body = b""
        async for chunk in response.body_iterator:  # type: ignore[attr-defined]
            body += chunk.encode() if isinstance(chunk, str) else chunk
        if len(body) <= _MAX_CACHE_BYTES and "application/json" in response.headers.get(
            "content-type", ""
        ):
            try:
                await store.set(
                    cache_key,
                    dumps({"status": response.status_code, "body": json.loads(body)}),
                    settings.idempotency_ttl_seconds,
                )
            except Exception:  # noqa: BLE001
                logger.warning("幂等结果写入失败", exc_info=True)

        return Response(
            content=body,
            status_code=response.status_code,
            headers=dict(response.headers),
            media_type=response.headers.get("content-type"),
        )

    @staticmethod
    def _subject(request: Request) -> str | None:
        header = request.headers.get("Authorization", "")
        if not header.lower().startswith("bearer "):
            return None
        try:
            return str(decode_token(header.split(" ", 1)[1].strip(), "access"))
        except Exception:  # noqa: BLE001 - 令牌问题交给后面的依赖统一报错
            return None


def error_body(request_id: str, code: ErrorCode, message: str, details: Any = None) -> dict:
    return {
        "code": int(code),
        "message": message,
        "data": details,
        "request_id": request_id,
    }