update token refresh expire time rule (#67)

* update token refresh expire time rule

* update token refresh time checks

Author: wu-clan

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

@@ -5,7 +5,7 @@
 
 from backend.app.common.jwt import DependsUser, get_token, jwt_decode, CurrentJwtAuth
 from backend.app.common.response.response_schema import response_base
-from backend.app.schemas.token import RefreshToken, LoginToken, SwaggerToken
+from backend.app.schemas.token import RefreshToken, LoginToken, SwaggerToken, RefreshTokenTime
 from backend.app.schemas.user import Auth
 from backend.app.services.user_service import UserService
 
@@ -32,10 +32,10 @@ async def user_login(obj: Auth):
 
 
 @router.post('/refresh_token', summary='刷新 token', dependencies=[DependsUser])
-async def get_refresh_token(request: Request):
+async def get_refresh_token(request: Request, custom_time: RefreshTokenTime):
     token = get_token(request)
     user_id, _ = jwt_decode(token)
-    refresh_token, refresh_expire = await UserService.refresh_token(user_id)
+    refresh_token, refresh_expire = await UserService.refresh_token(user_id, custom_time)
     data = RefreshToken(refresh_token=refresh_token, refresh_token_expire_time=refresh_expire)
     return response_base.success(data=data)
 

backend/app/common/exception/exception_handler.py

@@ -75,9 +75,9 @@ def validation_exception_handler(request: Request, exc: RequestValidationError):
                     _msg = error.get('msg')
                     errors_len = errors_len - 1
                     message += (
-                        f'{data.get(field, field)} {_msg}' + ', '
+                        f'{data.get(field, field) if field != "__root__" else ""} {_msg}' + ', '
                         if errors_len > 0
-                        else f'{data.get(field, field)} {_msg}' + '.'
+                        else f'{data.get(field, field) if field != "__root__" else ""} {_msg}' + '.'
                     )
             elif isinstance(raw_error.exc, json.JSONDecodeError):
                 message += 'json解析失败'

backend/app/common/jwt.py

@@ -16,6 +16,7 @@
 from backend.app.crud.crud_user import UserDao
 from backend.app.database.db_mysql import CurrentSession
 from backend.app.models import User
+from backend.app.schemas.token import RefreshTokenTime
 
 pwd_context = CryptContext(schemes=['bcrypt'], deprecated='auto')
 
@@ -66,26 +67,32 @@ async def create_access_token(sub: str, expires_delta: timedelta | None = None,
     return token, expire
 
 
-async def create_refresh_token(sub: str, expire_time: datetime | None = None, **kwargs) -> tuple[str, datetime]:
+async def create_refresh_token(
+    sub: str, expire_time: datetime | None = None, custom_expire_time: RefreshTokenTime | None = None, **kwargs
+) -> tuple[str, datetime]:
     """
     Generate encryption refresh token
 
     :param sub: The subject/userid of the JWT
     :param expire_time: expiry time
+    :param custom_expire_time: custom expiry time
     :return:
     """
     if expire_time:
-        expires = expire_time + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
-        expire_seconds = int((expires - datetime.utcnow()).total_seconds())
+        expire = expire_time + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
+        expire_seconds = int((expire - datetime.utcnow()).total_seconds())
+    elif custom_expire_time:
+        expire = custom_expire_time.expire_time
+        expire_seconds = int((expire - datetime.utcnow()).total_seconds())
     else:
-        expires = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
+        expire = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
         expire_seconds = settings.TOKEN_EXPIRE_SECONDS
-    to_encode = {'exp': expires, 'sub': sub, **kwargs}
+    to_encode = {'exp': expire, 'sub': sub, **kwargs}
     token = jwt.encode(to_encode, settings.TOKEN_SECRET_KEY, settings.TOKEN_ALGORITHM)
     # 刷新 token 时，保持旧 token 有效，不执行删除操作
     key = f'{settings.TOKEN_REDIS_PREFIX}:{sub}:{token}'
     await redis_client.setex(key, expire_seconds, token)
-    return token, expires
+    return token, expire
 
 
 def get_token(request: Request) -> str:

backend/app/common/response/response_schema.py

@@ -29,7 +29,7 @@ def test():
         @router.get('/test')
         def test() -> ResponseModel:
             return ResponseModel(data={'test': 'test'})
-    """
+    """  # noqa: E501
 
     code: int = 200
     msg: str = 'Success'
@@ -52,7 +52,8 @@ class ResponseBase:
         @router.get('/test')
         def test():
             return response_base.success(data={'test': 'test'})
-    """
+    """  # noqa: E501
+
     @staticmethod
     def __json_encoder(data: Any, exclude: _ExcludeData | None = None, **kwargs):
         custom_encoder = {datetime: lambda x: x.strftime('%Y-%m-%d %H:%M:%S')}

backend/app/schemas/token.py

@@ -1,8 +1,9 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from datetime import datetime
+from datetime import datetime, timedelta
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field, validator
+from pydantic.datetime_parse import parse_datetime
 
 from backend.app.schemas.user import GetUserInfoNoRelation
 
@@ -27,3 +28,22 @@ class RefreshToken(BaseModel):
     refresh_token: str
     refresh_token_type: str = 'Bearer'
     refresh_token_expire_time: datetime
+
+
+class RefreshTokenTime(BaseModel):
+    expire_time: datetime | None = Field(None, description='自定义刷新令牌过期时间')
+
+    @validator('expire_time', pre=True)
+    def validate_expire_time(cls, v):
+        if v is None:
+            return None
+        if not isinstance(v, str) or 'T' not in v:
+            raise ValueError('输入时间格式错误')
+        v = parse_datetime(v)
+        utcnow = datetime.utcnow()
+        no_tz_v = v.replace(tzinfo=None)
+        if no_tz_v < utcnow:
+            raise ValueError('输入时间小于当前时间')
+        if no_tz_v > utcnow + timedelta(days=7):
+            raise ValueError('输入时间大于当前时间上限 7 天')
+        return no_tz_v

backend/app/services/user_service.py

@@ -12,6 +12,7 @@
 from backend.app.crud.crud_user import UserDao
 from backend.app.database.db_mysql import async_db_session
 from backend.app.models import User
+from backend.app.schemas.token import RefreshTokenTime
 from backend.app.schemas.user import CreateUser, ResetPassword, UpdateUser, Avatar, Auth
 from backend.app.utils import re_verify
 
@@ -57,7 +58,7 @@ async def login(obj: Auth):
             return access_token, refresh_token, access_token_expire_time, refresh_token_expire_time, user
 
     @staticmethod
-    async def refresh_token(user_id: int):
+    async def refresh_token(user_id: int, custom_time: RefreshTokenTime):
         async with async_db_session() as db:
             current_user = await UserDao.get_user_by_id(db, user_id)
             if not current_user:
@@ -66,7 +67,7 @@ async def refresh_token(user_id: int):
                 raise errors.AuthorizationError(msg='用户已锁定, 获取失败')
             user_role_ids = await UserDao.get_user_role_ids(db, current_user.id)
             refresh_token, refresh_token_expire_time = await jwt.create_refresh_token(
-                str(current_user.id), role_ids=user_role_ids
+                str(current_user.id), custom_expire_time=custom_time, role_ids=user_role_ids
             )
             return refresh_token, refresh_token_expire_time