""" FastAPI 依赖注入：从 JWT 提取当前用户信息。 用法： @router.get("/protected") async def protected_endpoint(user: CurrentUser = Depends(get_current_user)): print(user.user_id, user.site_id) # 允许 pending 用户（受限令牌）访问 @router.get("/apply") async def apply_endpoint(user: CurrentUser = Depends(get_current_user_or_limited)): if user.limited: ... # 受限逻辑 """ from __future__ import annotations import time from dataclasses import dataclass, field from datetime import datetime from fastapi import Depends, HTTPException, status from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from jose import JWTError from app.auth.jwt import decode_access_token from app.trace.context import SpanType, TraceSpan, get_current_trace from app.trace.decorators import truncate_token # Bearer token 提取器 _bearer_scheme = HTTPBearer(auto_error=True) # ── 鉴权失败原因分类常量 ── AUTH_EXPIRED = "AUTH_EXPIRED" AUTH_INVALID = "AUTH_INVALID" AUTH_MALFORMED = "AUTH_MALFORMED" AUTH_LIMITED = "AUTH_LIMITED" AUTH_FORBIDDEN = "AUTH_FORBIDDEN" def _record_auth_span( *, token: str, success: bool, user_id: int | None = None, site_id: int | None = None, roles: list[str] | None = None, user_status: str = "", failure_reason: str = "", detail: str = "", duration_ms: float = 0.0, ) -> None: """向当前 TraceContext 添加 AUTH span（无 trace 时静默跳过）。""" ctx = get_current_trace() if ctx is None: return token_prefix = truncate_token(token) if success: desc_zh = f"JWT 鉴权通过：user_id={user_id}, site_id={site_id}, roles={roles}" desc_en = f"JWT auth passed: user_id={user_id}, site_id={site_id}, roles={roles}" result_summary = "approved" else: desc_zh = f"JWT 鉴权失败：{failure_reason} — {detail}" desc_en = f"JWT auth failed: {failure_reason} — {detail}" result_summary = failure_reason extra: dict = {} if failure_reason: extra["failure_reason"] = failure_reason ctx.add_span(TraceSpan( span_type=SpanType.AUTH, module="auth.dependencies", function="get_current_user", description_zh=desc_zh, description_en=desc_en, params={"token_prefix": token_prefix}, result_summary=result_summary, duration_ms=duration_ms, timestamp=datetime.now().isoformat(), extra=extra, )) # 鉴权成功时将 user_id / site_id 写入 TraceContext if success and user_id is not None: ctx.user_id = user_id if site_id is not None: ctx.site_id = site_id def _classify_jwt_error(exc: JWTError) -> str: """根据 JWTError 消息分类失败原因。""" msg = str(exc).lower() if "expired" in msg or "exp" in msg: return AUTH_EXPIRED return AUTH_INVALID @dataclass(frozen=True) class CurrentUser: """从 JWT 解析出的当前用户上下文。""" user_id: int site_id: int = 0 roles: list[str] = field(default_factory=list) status: str = "pending" limited: bool = False async def get_current_user( credentials: HTTPAuthorizationCredentials = Depends(_bearer_scheme), ) -> CurrentUser: """ FastAPI 依赖：从 Authorization header 提取 JWT，验证后返回用户信息。 要求完整令牌（非 limited），失败时抛出 401。 """ token = credentials.credentials start = time.perf_counter() try: payload = decode_access_token(token) except JWTError as exc: elapsed = (time.perf_counter() - start) * 1000 reason = _classify_jwt_error(exc) _record_auth_span( token=token, success=False, failure_reason=reason, detail="无效的令牌", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的令牌", headers={"WWW-Authenticate": "Bearer"}, ) # 受限令牌不允许通过此依赖 if payload.get("limited"): elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=False, failure_reason=AUTH_LIMITED, detail="受限令牌无法访问此端点", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="受限令牌无法访问此端点", headers={"WWW-Authenticate": "Bearer"}, ) user_id_raw = payload.get("sub") site_id = payload.get("site_id") if user_id_raw is None or site_id is None: elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=False, failure_reason=AUTH_MALFORMED, detail="令牌缺少必要字段", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="令牌缺少必要字段", headers={"WWW-Authenticate": "Bearer"}, ) try: user_id = int(user_id_raw) except (TypeError, ValueError): elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=False, failure_reason=AUTH_MALFORMED, detail="令牌中 user_id 格式无效", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="令牌中 user_id 格式无效", headers={"WWW-Authenticate": "Bearer"}, ) roles = payload.get("roles", []) elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=True, user_id=user_id, site_id=site_id, roles=roles, user_status="approved", duration_ms=elapsed, ) return CurrentUser( user_id=user_id, site_id=site_id, roles=roles, status="approved", limited=False, ) async def get_current_user_or_limited( credentials: HTTPAuthorizationCredentials = Depends(_bearer_scheme), ) -> CurrentUser: """ FastAPI 依赖：允许 pending 用户（受限令牌）访问。 - 受限令牌（limited=True）：返回 CurrentUser(limited=True, roles=[], status="pending") - 完整令牌：正常返回 CurrentUser """ token = credentials.credentials start = time.perf_counter() try: payload = decode_access_token(token) except JWTError as exc: elapsed = (time.perf_counter() - start) * 1000 reason = _classify_jwt_error(exc) _record_auth_span( token=token, success=False, failure_reason=reason, detail="无效的令牌", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="无效的令牌", headers={"WWW-Authenticate": "Bearer"}, ) user_id_raw = payload.get("sub") if user_id_raw is None: elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=False, failure_reason=AUTH_MALFORMED, detail="令牌缺少必要字段", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="令牌缺少必要字段", headers={"WWW-Authenticate": "Bearer"}, ) try: user_id = int(user_id_raw) except (TypeError, ValueError): elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=False, failure_reason=AUTH_MALFORMED, detail="令牌中 user_id 格式无效", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="令牌中 user_id 格式无效", headers={"WWW-Authenticate": "Bearer"}, ) # 受限令牌：pending 用户 if payload.get("limited"): elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=True, user_id=user_id, site_id=0, roles=[], user_status="pending", duration_ms=elapsed, ) return CurrentUser( user_id=user_id, site_id=0, roles=[], status="pending", limited=True, ) # 完整令牌：要求 site_id site_id = payload.get("site_id") if site_id is None: elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=False, failure_reason=AUTH_MALFORMED, detail="令牌缺少必要字段", duration_ms=elapsed, ) raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="令牌缺少必要字段", headers={"WWW-Authenticate": "Bearer"}, ) roles = payload.get("roles", []) elapsed = (time.perf_counter() - start) * 1000 _record_auth_span( token=token, success=True, user_id=user_id, site_id=site_id, roles=roles, user_status="approved", duration_ms=elapsed, ) return CurrentUser( user_id=user_id, site_id=site_id, roles=roles, status="approved", limited=False, )