Update the JWT for easier scaling (#1011)

* Update the JWT for easier scaling

* Fix comment

Author: wu-clan

backend/common/security/jwt.py

@@ -217,6 +217,30 @@ async def get_current_user(db: AsyncSession, pk: int) -> User:
     return user
 
 
+async def get_jwt_user(user_id: int) -> GetUserInfoWithRelationDetail:
+    """
+    获取 JWT 用户
+
+    :param user_id:
+    :return:
+    """
+    cache_user = await redis_client.get(f'{settings.JWT_USER_REDIS_PREFIX}:{user_id}')
+    if not cache_user:
+        async with async_db_session() as db:
+            current_user = await get_current_user(db, user_id)
+            user = GetUserInfoWithRelationDetail.model_validate(current_user)
+            await redis_client.setex(
+                f'{settings.JWT_USER_REDIS_PREFIX}:{user_id}',
+                settings.TOKEN_EXPIRE_SECONDS,
+                user.model_dump_json(),
+            )
+    else:
+        # TODO: 在恰当的时机，应替换为使用 model_validate_json
+        # https://docs.pydantic.dev/latest/concepts/json/#partial-json-parsing
+        user = GetUserInfoWithRelationDetail.model_validate(from_json(cache_user, allow_partial=True))
+    return user
+
+
 def superuser_verify(request: Request, _token: str = DependsJwtAuth) -> bool:
     """
     验证当前用户超级管理员权限
@@ -247,21 +271,7 @@ async def jwt_authentication(token: str) -> GetUserInfoWithRelationDetail:
     if token != redis_token:
         raise errors.TokenError(msg='Token 已失效')
 
-    cache_user = await redis_client.get(f'{settings.JWT_USER_REDIS_PREFIX}:{user_id}')
-    if not cache_user:
-        async with async_db_session() as db:
-            current_user = await get_current_user(db, user_id)
-            user = GetUserInfoWithRelationDetail.model_validate(current_user)
-            await redis_client.setex(
-                f'{settings.JWT_USER_REDIS_PREFIX}:{user_id}',
-                settings.TOKEN_EXPIRE_SECONDS,
-                user.model_dump_json(),
-            )
-    else:
-        # TODO: 在恰当的时机，应替换为使用 model_validate_json
-        # https://docs.pydantic.dev/latest/concepts/json/#partial-json-parsing
-        user = GetUserInfoWithRelationDetail.model_validate(from_json(cache_user, allow_partial=True))
-    return user
+    return await get_jwt_user(user_id)
 
 
 # 超级管理员鉴权依赖注入

backend/middleware/jwt_auth_middleware.py

@@ -2,7 +2,8 @@
 
 from fastapi import Request, Response
 from fastapi.security.utils import get_authorization_scheme_param
-from starlette.authentication import AuthCredentials, AuthenticationBackend, AuthenticationError
+from starlette.authentication import AuthCredentials, AuthenticationBackend
+from starlette.authentication import AuthenticationError as StarletteAuthenticationError
 from starlette.requests import HTTPConnection
 
 from backend.app.admin.schema.user import GetUserInfoWithRelationDetail
@@ -13,7 +14,7 @@
 from backend.utils.serializers import MsgSpecJSONResponse
 
 
-class _AuthenticationError(AuthenticationError):
+class AuthenticationError(StarletteAuthenticationError):
     """重写内部认证错误类"""
 
     def __init__(
@@ -40,7 +41,7 @@ class JwtAuthMiddleware(AuthenticationBackend):
     """JWT 认证中间件"""
 
     @staticmethod
-    def auth_exception_handler(conn: HTTPConnection, exc: _AuthenticationError) -> Response:
+    def auth_exception_handler(conn: HTTPConnection, exc: AuthenticationError) -> Response:
         """
         覆盖内部认证错误处理
 
@@ -50,15 +51,16 @@ def auth_exception_handler(conn: HTTPConnection, exc: _AuthenticationError) -> R
         """
         return MsgSpecJSONResponse(content={'code': exc.code, 'msg': exc.msg, 'data': None}, status_code=exc.code)
 
-    async def authenticate(self, request: Request) -> tuple[AuthCredentials, GetUserInfoWithRelationDetail] | None:
+    @staticmethod
+    def extract_token(request: Request) -> str | None:
         """
-        认证请求
+        从请求中提取 Bearer Token
 
         :param request: FastAPI 请求对象
         :return:
         """
-        token = request.headers.get('Authorization')
-        if not token:
+        authorization = request.headers.get('Authorization')
+        if not authorization:
             return None
 
         path = request.url.path
@@ -68,17 +70,30 @@ async def authenticate(self, request: Request) -> tuple[AuthCredentials, GetUser
             if pattern.match(path):
                 return None
 
-        scheme, token = get_authorization_scheme_param(token)
+        scheme, token = get_authorization_scheme_param(authorization)
         if scheme.lower() != 'bearer':
             return None
 
+        return token
+
+    async def authenticate(self, request: Request) -> tuple[AuthCredentials, GetUserInfoWithRelationDetail] | None:
+        """
+        认证请求
+
+        :param request: FastAPI 请求对象
+        :return:
+        """
+        token = self.extract_token(request)
+        if token is None:
+            return None
+
         try:
             user = await jwt_authentication(token)
         except TokenError as exc:
-            raise _AuthenticationError(code=exc.code, msg=exc.detail, headers=exc.headers)
+            raise AuthenticationError(code=exc.code, msg=exc.detail, headers=exc.headers)
         except Exception as e:
             log.exception(f'JWT 授权异常：{e}')
-            raise _AuthenticationError(code=getattr(e, 'code', 500), msg=getattr(e, 'msg', 'Internal Server Error'))
+            raise AuthenticationError(code=getattr(e, 'code', 500), msg=getattr(e, 'msg', 'Internal Server Error'))
 
         # 请注意，此返回使用非标准模式，所以在认证通过时，将丢失某些标准特性
         # 标准返回模式请查看：https://www.starlette.io/authentication/