71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
"""request_id 关联中间件（ARCH §9.3 / CLAUDE.md）。
|
||
|
||
每请求生成/透传 request_id，绑入 structlog contextvars，并回写响应头，
|
||
错误信封也带上同一 id，便于端到端 grep。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import uuid
|
||
from collections.abc import Awaitable, Callable
|
||
|
||
import structlog
|
||
from starlette.requests import Request
|
||
from starlette.responses import JSONResponse, Response
|
||
from ww_shared import ErrorBody, ErrorCode, ErrorEnvelope
|
||
|
||
REQUEST_ID_HEADER = "x-request-id"
|
||
|
||
# 请求体大小上限（2 MiB）：超过即 413，防超大 body 撑爆内存/成本（CR-H9）。
|
||
MAX_BODY_BYTES = 2 * 1024 * 1024
|
||
|
||
# 放行安全字符集（字母数字 . _ -，1–128 位）：覆盖 uuid4().hex、常见 trace id 与客户端
|
||
# 自定义短 id；含空白/控制字符/换行/注入序列/超长的一律丢弃改生成新 uuid——防止把未经
|
||
# 校验的客户端值原样写进日志/响应头（日志注入防护）。
|
||
_REQUEST_ID_PATTERN = re.compile(r"^[A-Za-z0-9._-]{1,128}$")
|
||
|
||
|
||
def _sanitize_request_id(raw: str | None) -> str:
|
||
if raw is not None and _REQUEST_ID_PATTERN.match(raw):
|
||
return raw
|
||
return uuid.uuid4().hex
|
||
|
||
|
||
async def request_id_middleware(
|
||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||
) -> Response:
|
||
request_id = _sanitize_request_id(request.headers.get(REQUEST_ID_HEADER))
|
||
structlog.contextvars.clear_contextvars()
|
||
structlog.contextvars.bind_contextvars(request_id=request_id)
|
||
request.state.request_id = request_id
|
||
response = await call_next(request)
|
||
response.headers[REQUEST_ID_HEADER] = request_id
|
||
return response
|
||
|
||
|
||
async def body_size_limit_middleware(
|
||
request: Request, call_next: Callable[[Request], Awaitable[Response]]
|
||
) -> Response:
|
||
"""按 `Content-Length` 拦截超大请求体，超限直接 413（CR-H9）。
|
||
|
||
用户中间件跑在 FastAPI 的 AppError 处理器**之外**，故必须自建错误信封响应（不能靠抛
|
||
AppError）。分块编码（无 `Content-Length`）无法在此拦截——单用户原型可接受的限制。
|
||
"""
|
||
content_length = request.headers.get("content-length")
|
||
if content_length is not None:
|
||
try:
|
||
size = int(content_length)
|
||
except ValueError:
|
||
size = 0 # 非法头防御性视作 0，交由下游正常解析/报错。
|
||
if size > MAX_BODY_BYTES:
|
||
envelope = ErrorEnvelope(
|
||
error=ErrorBody(
|
||
code=ErrorCode.PAYLOAD_TOO_LARGE,
|
||
message="请求体过大",
|
||
request_id=_sanitize_request_id(request.headers.get(REQUEST_ID_HEADER)),
|
||
)
|
||
)
|
||
return JSONResponse(status_code=413, content=envelope.model_dump())
|
||
return await call_next(request)
|