Updated refresh token storage logic (#403)

Author: wu-clan

backend/app/admin/api/v1/auth/auth.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, Query, Request
+from fastapi import APIRouter, Depends, Request, Response
 from fastapi.security import HTTPBasicCredentials
 from fastapi_limiter.depends import RateLimiter
 from starlette.background import BackgroundTasks
@@ -28,18 +28,20 @@ async def swagger_login(obj: Annotated[HTTPBasicCredentials, Depends()]) -> GetS
     description='json 格式登录, 仅支持在第三方api工具调试, 例如: postman',
     dependencies=[Depends(RateLimiter(times=5, minutes=1))],
 )
-async def user_login(request: Request, obj: AuthLoginParam, background_tasks: BackgroundTasks) -> ResponseModel:
-    data = await auth_service.login(request=request, obj=obj, background_tasks=background_tasks)
+async def user_login(
+    request: Request, response: Response, obj: AuthLoginParam, background_tasks: BackgroundTasks
+) -> ResponseModel:
+    data = await auth_service.login(request=request, response=response, obj=obj, background_tasks=background_tasks)
     return response_base.success(data=data)
 
 
 @router.post('/token/new', summary='创建新 token', dependencies=[DependsJwtAuth])
-async def create_new_token(request: Request, refresh_token: Annotated[str, Query(...)]) -> ResponseModel:
-    data = await auth_service.new_token(request=request, refresh_token=refresh_token)
+async def create_new_token(request: Request, response: Response) -> ResponseModel:
+    data = await auth_service.new_token(request=request, response=response)
     return response_base.success(data=data)
 
 
 @router.post('/logout', summary='用户登出', dependencies=[DependsJwtAuth])
-async def user_logout(request: Request) -> ResponseModel:
-    await auth_service.logout(request=request)
+async def user_logout(request: Request, response: Response) -> ResponseModel:
+    await auth_service.logout(request=request, response=response)
     return response_base.success()

backend/app/admin/schema/token.py

@@ -19,10 +19,8 @@ class AccessTokenBase(SchemaBase):
 
 
 class GetNewToken(AccessTokenBase):
-    refresh_token: str
-    refresh_token_type: str = 'Bearer'
-    refresh_token_expire_time: datetime
+    pass
 
 
-class GetLoginToken(GetNewToken):
+class GetLoginToken(AccessTokenBase):
     user: GetUserInfoNoRelationDetail

backend/app/admin/service/auth_service.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from fastapi import Request
+from fastapi import Request, Response
 from fastapi.security import HTTPBasicCredentials
 from starlette.background import BackgroundTask, BackgroundTasks
 
@@ -38,12 +38,14 @@ async def swagger_login(*, obj: HTTPBasicCredentials) -> tuple[str, User]:
                 raise errors.AuthorizationError(msg='用户名或密码有误')
             elif not current_user.status:
                 raise errors.AuthorizationError(msg='用户已被锁定, 请联系统管理员')
-            access_token, _ = await create_access_token(str(current_user.id), multi_login=current_user.is_multi_login)
+            access_token = await create_access_token(str(current_user.id), current_user.is_multi_login)
             await user_dao.update_login_time(db, obj.username)
-            return access_token, current_user
+            return access_token.access_token, current_user
 
     @staticmethod
-    async def login(*, request: Request, obj: AuthLoginParam, background_tasks: BackgroundTasks) -> GetLoginToken:
+    async def login(
+        *, request: Request, response: Response, obj: AuthLoginParam, background_tasks: BackgroundTasks
+    ) -> GetLoginToken:
         async with async_db_session.begin() as db:
             try:
                 current_user = await user_dao.get_by_username(db, obj.username)
@@ -61,14 +63,8 @@ async def login(*, request: Request, obj: AuthLoginParam, background_tasks: Back
                 if captcha_code.lower() != obj.captcha.lower():
                     raise errors.CustomError(error=CustomErrorCode.CAPTCHA_ERROR)
                 current_user_id = current_user.id
-                access_token, access_token_expire_time = await create_access_token(
-                    str(current_user_id), multi_login=current_user.is_multi_login
-                )
-                refresh_token, refresh_token_expire_time = await create_refresh_token(
-                    sub=str(current_user_id),
-                    expire_time=access_token_expire_time,
-                    multi_login=current_user.is_multi_login,
-                )
+                access_token = await create_access_token(str(current_user_id), current_user.is_multi_login)
+                refresh_token = await create_refresh_token(str(current_user_id), current_user.is_multi_login)
             except errors.NotFoundError as e:
                 raise errors.NotFoundError(msg=e.msg)
             except (errors.AuthorizationError, errors.CustomError) as e:
@@ -102,19 +98,29 @@ async def login(*, request: Request, obj: AuthLoginParam, background_tasks: Back
                 )
                 await redis_client.delete(f'{admin_settings.CAPTCHA_LOGIN_REDIS_PREFIX}:{request.state.ip}')
                 await user_dao.update_login_time(db, obj.username)
+                response.set_cookie(
+                    settings.COOKIE_REFRESH_TOKEN_KEY,
+                    refresh_token.refresh_token,
+                    settings.COOKIE_REFRESH_TOKEN_EXPIRE_SECONDS,
+                    refresh_token.refresh_token_expire_time,
+                )
                 await db.refresh(current_user)
                 data = GetLoginToken(
-                    access_token=access_token,
-                    refresh_token=refresh_token,
-                    access_token_expire_time=access_token_expire_time,
-                    refresh_token_expire_time=refresh_token_expire_time,
+                    access_token=access_token.access_token,
+                    access_token_expire_time=access_token.access_token_expire_time,
                     user=current_user,  # type: ignore
                 )
                 return data
 
     @staticmethod
-    async def new_token(*, request: Request, refresh_token: str) -> GetNewToken:
-        user_id = jwt_decode(refresh_token)
+    async def new_token(*, request: Request, response: Response) -> GetNewToken:
+        refresh_token = request.cookies.get(settings.COOKIE_REFRESH_TOKEN_KEY)
+        if not refresh_token:
+            raise errors.TokenError(msg='Refresh Token 丢失，请重新登录')
+        try:
+            user_id = jwt_decode(refresh_token)
+        except Exception:
+            raise errors.TokenError(msg='Refresh Token 无效')
         if request.user.id != user_id:
             raise errors.TokenError(msg='Refresh Token 无效')
         async with async_db_session() as db:
@@ -130,23 +136,34 @@ async def new_token(*, request: Request, refresh_token: str) -> GetNewToken:
                 refresh_token=refresh_token,
                 multi_login=current_user.is_multi_login,
             )
+            response.set_cookie(
+                settings.COOKIE_REFRESH_TOKEN_KEY,
+                new_token.new_refresh_token,
+                settings.COOKIE_REFRESH_TOKEN_EXPIRE_SECONDS,
+                new_token.new_refresh_token_expire_time,
+            )
             data = GetNewToken(
                 access_token=new_token.new_access_token,
                 access_token_expire_time=new_token.new_access_token_expire_time,
-                refresh_token=new_token.new_refresh_token,
-                refresh_token_expire_time=new_token.new_refresh_token_expire_time,
             )
             return data
 
     @staticmethod
-    async def logout(*, request: Request) -> None:
+    async def logout(*, request: Request, response: Response) -> None:
         token = await get_token(request)
+        refresh_token = request.cookies.get(settings.COOKIE_REFRESH_TOKEN_KEY)
+        response.delete_cookie(settings.COOKIE_REFRESH_TOKEN_KEY)
         if request.user.is_multi_login:
             key = f'{settings.TOKEN_REDIS_PREFIX}:{request.user.id}:{token}'
             await redis_client.delete(key)
+            if refresh_token:
+                key = f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{request.user.id}:{refresh_token}'
+                await redis_client.delete(key)
         else:
             key_prefix = f'{settings.TOKEN_REDIS_PREFIX}:{request.user.id}:'
             await redis_client.delete_prefix(key_prefix)
+            key_prefix = f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{request.user.id}:'
+            await redis_client.delete_prefix(key_prefix)
 
 
 auth_service = AuthService()

backend/app/admin/service/oauth2_service.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 from fast_captcha import text_captcha
-from fastapi import BackgroundTasks, Request
+from fastapi import BackgroundTasks, Request, Response
 
 from backend.app.admin.conf import admin_settings
 from backend.app.admin.crud.crud_user import user_dao
@@ -13,6 +13,7 @@
 from backend.common.enums import LoginLogStatusType, UserSocialType
 from backend.common.exception.errors import AuthorizationError
 from backend.common.security import jwt
+from backend.core.conf import settings
 from backend.database.db_mysql import async_db_session
 from backend.database.db_redis import redis_client
 from backend.utils.timezone import timezone
@@ -21,7 +22,7 @@
 class OAuth2Service:
     @staticmethod
     async def create_with_login(
-        *, request: Request, background_tasks: BackgroundTasks, user: dict, social: UserSocialType
+        *, request: Request, response: Response, background_tasks: BackgroundTasks, user: dict, social: UserSocialType
     ) -> GetLoginToken | None:
         async with async_db_session.begin() as db:
             # 获取 OAuth2 平台用户信息
@@ -54,12 +55,8 @@ async def create_with_login(
                 new_user_social = CreateUserSocialParam(source=social.value, uid=str(_id), user_id=sys_user.id)
                 await user_social_dao.create(db, new_user_social)
             # 创建 token
-            access_token, access_token_expire_time = await jwt.create_access_token(
-                str(sys_user.id), multi_login=sys_user.is_multi_login
-            )
-            refresh_token, refresh_token_expire_time = await jwt.create_refresh_token(
-                str(sys_user.id), access_token_expire_time, multi_login=sys_user.is_multi_login
-            )
+            access_token = await jwt.create_access_token(str(sys_user.id), sys_user.is_multi_login)
+            refresh_token = await jwt.create_refresh_token(str(sys_user.id), multi_login=sys_user.is_multi_login)
             await user_dao.update_login_time(db, sys_user.username)
             await db.refresh(sys_user)
             login_log = dict(
@@ -72,11 +69,15 @@ async def create_with_login(
             )
             background_tasks.add_task(LoginLogService.create, **login_log)
             await redis_client.delete(f'{admin_settings.CAPTCHA_LOGIN_REDIS_PREFIX}:{request.state.ip}')
+            response.set_cookie(
+                settings.COOKIE_REFRESH_TOKEN_KEY,
+                refresh_token.refresh_token,
+                settings.COOKIE_REFRESH_TOKEN_EXPIRE_SECONDS,
+                refresh_token.refresh_token_expire_time,
+            )
             data = GetLoginToken(
-                access_token=access_token,
-                refresh_token=refresh_token,
-                access_token_expire_time=access_token_expire_time,
-                refresh_token_expire_time=refresh_token_expire_time,
+                access_token=access_token.access_token,
+                access_token_expire_time=access_token.access_token_expire_time,
                 user=sys_user,  # type: ignore
             )
             return data

backend/common/dataclasses.py

@@ -26,7 +26,7 @@ class UserAgentInfo:
 
 
 @dataclasses.dataclass
-class RequestCallNextReturn:
+class RequestCallNext:
     code: str
     msg: str
     status: StatusType
@@ -35,8 +35,20 @@ class RequestCallNextReturn:
 
 
 @dataclasses.dataclass
-class NewTokenReturn:
+class NewToken:
     new_access_token: str
-    new_refresh_token: str
     new_access_token_expire_time: datetime
+    new_refresh_token: str
     new_refresh_token_expire_time: datetime
+
+
+@dataclasses.dataclass
+class AccessToken:
+    access_token: str
+    access_token_expire_time: datetime
+
+
+@dataclasses.dataclass
+class RefreshToken:
+    refresh_token: str
+    refresh_token_expire_time: datetime

backend/common/response/response_schema.py

@@ -57,7 +57,7 @@ class ResponseBase:
 
         @router.get('/test')
         def test() -> ResponseModel:
-            return await response_base.success(data={'test': 'test'})
+            return response_base.success(data={'test': 'test'})
     """
 
     @staticmethod