118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
"""用户中心客户端：用请求携带的 token 换取当前用户，结果短时缓存"""
|
|
|
|
import hashlib
|
|
import time
|
|
from functools import cache
|
|
|
|
import httpx
|
|
from loguru import logger
|
|
|
|
from app.config import Settings, settings
|
|
from app.core.exceptions import BizException
|
|
from app.schemas.user import CurrentUser
|
|
|
|
# 用户中心 401 时提示调用方用 Bearer token 重新认证
|
|
BEARER_CHALLENGE = {"WWW-Authenticate": "Bearer"}
|
|
|
|
|
|
class UserCenterClient:
|
|
def __init__(
|
|
self,
|
|
config: Settings | None = None,
|
|
transport: httpx.AsyncBaseTransport | None = None,
|
|
):
|
|
self._config = config if config is not None else settings
|
|
self._transport = transport
|
|
self._client: httpx.AsyncClient | None = None
|
|
# token 摘要 -> (过期时刻, 用户)；不在内存中保存明文 token
|
|
self._cache: dict[str, tuple[float, CurrentUser]] = {}
|
|
|
|
def _get_client(self) -> httpx.AsyncClient:
|
|
if self._client is None:
|
|
self._client = httpx.AsyncClient(
|
|
base_url=self._config.user_center_url,
|
|
timeout=self._config.user_center_timeout_seconds,
|
|
transport=self._transport,
|
|
)
|
|
return self._client
|
|
|
|
def _cached(self, key: str) -> CurrentUser | None:
|
|
entry = self._cache.get(key)
|
|
if entry is None:
|
|
return None
|
|
if entry[0] <= time.monotonic():
|
|
del self._cache[key]
|
|
return None
|
|
return entry[1]
|
|
|
|
def _remember(self, key: str, user: CurrentUser) -> None:
|
|
ttl = self._config.auth_cache_ttl_seconds
|
|
if not ttl:
|
|
return
|
|
now = time.monotonic()
|
|
# 写入时顺带清掉过期项，缓存规模不超过 TTL 内的活跃 token 数
|
|
for stale in [k for k, (expires, _) in self._cache.items() if expires <= now]:
|
|
del self._cache[stale]
|
|
self._cache[key] = (now + ttl, user)
|
|
|
|
async def fetch_current_user(self, token: str) -> CurrentUser:
|
|
key = hashlib.sha256(token.encode()).hexdigest()
|
|
user = self._cached(key)
|
|
if user is not None:
|
|
return user
|
|
|
|
try:
|
|
response = await self._get_client().get(
|
|
"/api/users/me", headers={"Authorization": f"Bearer {token}"}
|
|
)
|
|
except httpx.HTTPError as error:
|
|
logger.opt(exception=error).warning("用户中心请求失败")
|
|
raise BizException(503, "用户中心不可用") from error
|
|
|
|
if response.status_code >= 500:
|
|
logger.warning("用户中心返回 HTTP {}", response.status_code)
|
|
raise BizException(503, "用户中心不可用")
|
|
try:
|
|
body = response.json()
|
|
except ValueError:
|
|
body = None
|
|
if not isinstance(body, dict):
|
|
body = {}
|
|
if response.status_code >= 400:
|
|
# 用户中心与本服务的错误响应同构，直接透传状态码和提示
|
|
raise BizException(
|
|
response.status_code,
|
|
body.get("msg") or "认证失败",
|
|
BEARER_CHALLENGE if response.status_code == 401 else None,
|
|
)
|
|
try:
|
|
user = CurrentUser.model_validate(body["data"])
|
|
except (KeyError, TypeError, ValueError) as error:
|
|
logger.warning("用户中心返回的数据格式无效")
|
|
raise BizException(503, "用户中心不可用") from error
|
|
|
|
self._remember(key, user)
|
|
return user
|
|
|
|
async def aclose(self) -> None:
|
|
if self._client is not None:
|
|
await self._client.aclose()
|
|
self._client = None
|
|
|
|
|
|
@cache
|
|
def _cached_client() -> UserCenterClient:
|
|
return UserCenterClient()
|
|
|
|
|
|
def get_user_center() -> UserCenterClient:
|
|
"""共享入口；适用于单个服务事件循环"""
|
|
return _cached_client()
|
|
|
|
|
|
async def close_user_center() -> None:
|
|
"""在应用关闭阶段调用"""
|
|
if _cached_client.cache_info().currsize:
|
|
await _cached_client().aclose()
|
|
_cached_client.cache_clear()
|