update token handling logic (#83)

* Update token store, refresh, whitelist

* Update the token handling logic when update a user's multi-login status

* Delete all tokens when the user delete

* Fix the user logout interface

* Fix the get refresh token interface

* Fix multi-point login judgement when creat token

* Update the refresh_token interface

* Fix redis prefix deletion exclusion

* Fix token deletion error when user deleted

* Update the token time base to datetime.now()

* Update limiter storage prefix to settings

* Fix user login time not updated to database

* Allowing a user to have multiple refresh tokens.

* Add code comment to user multi-login update method

* Remove refresh token get and create interface

* Add user update multipoint login delete refresh token

Author: wu-clan

backend/app/.env.example

@@ -17,4 +17,3 @@ APS_REDIS_PASSWORD=''
 APS_REDIS_DATABASE=1
 # Token
 TOKEN_SECRET_KEY='1VkVF75nsNABBjK_7-qz7GtzNy3AMvktc9TCPwKczCk'
-TOKEN_WHITE_LIST=[1]

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

@@ -1,22 +1,24 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from fastapi import APIRouter, Depends, Request
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, Request, Query
 from fastapi.security import OAuth2PasswordRequestForm
 from fastapi_limiter.depends import RateLimiter
 from starlette.background import BackgroundTasks
 
-from backend.app.common.jwt import DependsUser, get_token, jwt_decode, CurrentJwtAuth
+from backend.app.common.jwt import DependsUser, CurrentUser
 from backend.app.common.response.response_schema import response_base
-from backend.app.schemas.token import RefreshToken, LoginToken, SwaggerToken, RefreshTokenTime
+from backend.app.schemas.token import LoginToken, SwaggerToken, NewToken
 from backend.app.schemas.user import Auth
-from backend.app.services.user_service import UserService
+from backend.app.services.auth_service import AuthService
 
 router = APIRouter()
 
 
 @router.post('/swagger_login', summary='swagger 表单登录', description='form 格式登录，仅用于 swagger 文档调试接口')
 async def swagger_user_login(form_data: OAuth2PasswordRequestForm = Depends()) -> SwaggerToken:
-    token, user = await UserService().swagger_login(form_data)
+    token, user = await AuthService().swagger_login(form_data)
     return SwaggerToken(access_token=token, user=user)
 
 
@@ -27,7 +29,7 @@ async def swagger_user_login(form_data: OAuth2PasswordRequestForm = Depends()) -
     dependencies=[Depends(RateLimiter(times=5, minutes=15))],
 )
 async def user_login(request: Request, obj: Auth, background_tasks: BackgroundTasks):
-    access_token, refresh_token, access_expire, refresh_expire, user = await UserService().login(
+    access_token, refresh_token, access_expire, refresh_expire, user = await AuthService().login(
         request=request, obj=obj, background_tasks=background_tasks
     )
     data = LoginToken(
@@ -40,17 +42,14 @@ async def user_login(request: Request, obj: Auth, background_tasks: BackgroundTa
     return response_base.success(data=data)
 
 
-@router.post('/refresh_token', summary='刷新 token', dependencies=[DependsUser])
-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=user_id, custom_time=custom_time)
-    data = RefreshToken(refresh_token=refresh_token, refresh_token_expire_time=refresh_expire)
+@router.post('/new_token', summary='创建新 token', dependencies=[DependsUser])
+async def create_new_token(refresh_token: Annotated[str, Query(...)]):
+    access_token, access_expire = await AuthService.new_token(refresh_token)
+    data = NewToken(access_token=access_token, access_token_expire_time=access_expire)
     return response_base.success(data=data)
 
 
-@router.post('/logout', summary='用户登出', dependencies=[DependsUser])
-async def user_logout(jwt: CurrentJwtAuth):
-    user_id = jwt.get('sub')
-    await UserService.logout(user_id)
+@router.post('/logout', summary='用户登出')
+async def user_logout(request: Request, current_user: CurrentUser):
+    await AuthService.logout(request=request, current_user=current_user)
     return response_base.success()

backend/app/api/v1/user.py

@@ -2,7 +2,7 @@
 # -*- coding: utf-8 -*-
 from typing import Annotated
 
-from fastapi import APIRouter, Query
+from fastapi import APIRouter, Query, Request
 
 from backend.app.common.jwt import DependsUser, CurrentUser, DependsSuperUser
 from backend.app.common.pagination import paging_data, PageDepends
@@ -54,10 +54,10 @@ async def update_avatar(username: str, avatar: Avatar, current_user: CurrentUser
 
 @router.get('', summary='（模糊条件）分页获取所有用户', dependencies=[DependsUser, PageDepends])
 async def get_all_users(
-        db: CurrentSession,
-        username: Annotated[str | None, Query()] = None,
-        phone: Annotated[str | None, Query()] = None,
-        status: Annotated[int | None, Query()] = None,
+    db: CurrentSession,
+    username: Annotated[str | None, Query()] = None,
+    phone: Annotated[str | None, Query()] = None,
+    status: Annotated[int | None, Query()] = None,
 ):
     user_select = await UserService.get_select(username=username, phone=phone, status=status)
     page_data = await paging_data(db, user_select, GetAllUserInfo)
@@ -80,6 +80,14 @@ async def active_set(pk: int):
     return response_base.fail()
 
 
+@router.post('/{pk}/multi', summary='修改用户多点登录状态')
+async def multi_set(request: Request, pk: int, current_user: CurrentUser):
+    count = await UserService.update_multi_login(request=request, pk=pk, current_user=current_user)
+    if count > 0:
+        return response_base.success()
+    return response_base.fail()
+
+
 @router.delete('/{username}', summary='用户注销', description='用户注销 != 用户退出，注销之后用户将从数据库删除')
 async def delete_user(username: str, current_user: CurrentUser):
     count = await UserService.delete(username=username, current_user=current_user)

backend/app/common/jwt.py

@@ -16,7 +16,6 @@
 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')
 
@@ -53,46 +52,60 @@ async def create_access_token(sub: str, expires_delta: timedelta | None = None,
     :return:
     """
     if expires_delta:
-        expire = datetime.utcnow() + expires_delta
+        expire = datetime.now() + expires_delta
         expire_seconds = int(expires_delta.total_seconds())
     else:
-        expire = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
+        expire = datetime.now() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
         expire_seconds = settings.TOKEN_EXPIRE_SECONDS
+    multi_login = kwargs.pop('multi_login', None)
     to_encode = {'exp': expire, 'sub': sub, **kwargs}
     token = jwt.encode(to_encode, settings.TOKEN_SECRET_KEY, settings.TOKEN_ALGORITHM)
-    if sub not in settings.TOKEN_WHITE_LIST:
-        await redis_client.delete_prefix(f'{settings.TOKEN_REDIS_PREFIX}:{sub}:')
+    if multi_login is False:
+        prefix = f'{settings.TOKEN_REDIS_PREFIX}:{sub}:'
+        await redis_client.delete_prefix(prefix)
     key = f'{settings.TOKEN_REDIS_PREFIX}:{sub}:{token}'
     await redis_client.setex(key, expire_seconds, token)
     return token, expire
 
 
-async def create_refresh_token(
-    sub: str, expire_time: datetime | None = None, custom_expire_time: RefreshTokenTime | None = None, **kwargs
-) -> tuple[str, datetime]:
+async def create_refresh_token(sub: str, expire_time: datetime | None = None, **kwargs) -> tuple[str, datetime]:
     """
-    Generate encryption refresh token
+    Generate encryption refresh token, only used to create a new 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:
-        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())
+        expire = expire_time + timedelta(seconds=settings.TOKEN_REFRESH_EXPIRE_SECONDS)
+        expire_seconds = int((expire - datetime.now()).total_seconds())
     else:
-        expire = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
-        expire_seconds = settings.TOKEN_EXPIRE_SECONDS
+        expire = datetime.now() + timedelta(seconds=settings.TOKEN_REFRESH_EXPIRE_SECONDS)
+        expire_seconds = settings.TOKEN_REFRESH_EXPIRE_SECONDS
+    multi_login = kwargs.pop('multi_login', None)
     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, expire
+    refresh_token = jwt.encode(to_encode, settings.TOKEN_SECRET_KEY, settings.TOKEN_ALGORITHM)
+    if multi_login is False:
+        prefix = f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{sub}:'
+        await redis_client.delete_prefix(prefix)
+    key = f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{sub}:{refresh_token}'
+    await redis_client.setex(key, expire_seconds, refresh_token)
+    return refresh_token, expire
+
+
+async def create_new_token(sub: str, refresh_token: str, **kwargs) -> tuple[str, datetime]:
+    """
+    Generate new token
+
+    :param sub:
+    :param refresh_token:
+    :return:
+    """
+    redis_refresh_token = await redis_client.get(f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{sub}:{refresh_token}')
+    if not redis_refresh_token or redis_refresh_token != refresh_token:
+        raise TokenError(msg='refresh_token 已过期')
+    new_token, expire = await create_access_token(sub, **kwargs)
+    return new_token, expire
 
 
 def get_token(request: Request) -> str:
@@ -102,10 +115,10 @@ def get_token(request: Request) -> str:
     :return:
     """
     authorization = request.headers.get('Authorization')
-    scheme, param = get_authorization_scheme_param(authorization)
+    scheme, token = get_authorization_scheme_param(authorization)
     if not authorization or scheme.lower() != 'bearer':
         raise TokenError
-    return param
+    return token
 
 
 def jwt_decode(token: str) -> tuple[int, list[int]]:
@@ -118,12 +131,12 @@ def jwt_decode(token: str) -> tuple[int, list[int]]:
     try:
         payload = jwt.decode(token, settings.TOKEN_SECRET_KEY, algorithms=[settings.TOKEN_ALGORITHM])
         user_id = int(payload.get('sub'))
-        user_roles = list(payload.get('role_ids'))
-        if not user_id or not user_roles:
+        role_ids = list(payload.get('role_ids'))
+        if not user_id or not role_ids:
             raise TokenError
     except (jwt.JWTError, ValidationError, Exception):
         raise TokenError
-    return user_id, user_roles
+    return user_id, role_ids
 
 
 async def jwt_authentication(token: str = Depends(oauth2_schema)) -> dict[str, int]:
@@ -137,7 +150,7 @@ async def jwt_authentication(token: str = Depends(oauth2_schema)) -> dict[str, i
     key = f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{token}'
     token_verify = await redis_client.get(key)
     if not token_verify:
-        raise TokenError
+        raise TokenError(msg='token 已过期')
     return {'sub': user_id}
 
 

backend/app/common/redis.py

@@ -37,16 +37,26 @@ async def open(self):
             log.error('❌ 数据库 redis 连接异常 {}', e)
             sys.exit()
 
-    async def delete_prefix(self, key: str):
+    async def delete_prefix(self, prefix: str, exclude: str | list = None):
         """
         删除指定前缀的所有key
 
-        :param key:
+        :param prefix:
+        :param exclude:
         :return:
         """
-        keys = await self.keys(f'{key}*')
-        if keys:
-            await self.delete(*keys)
+        keys = []
+        async for key in self.scan_iter(match=f'{prefix}*'):
+            if isinstance(exclude, str):
+                if key != exclude:
+                    keys.append(key)
+            elif isinstance(exclude, list):
+                if key not in exclude:
+                    keys.append(key)
+            else:
+                keys.append(key)
+        for key in keys:
+            await self.delete(key)
 
 
 # 创建redis连接对象

backend/app/core/conf.py

@@ -30,7 +30,6 @@ class Settings(BaseSettings):
 
     # Env Token
     TOKEN_SECRET_KEY: str  # 密钥 secrets.token_urlsafe(32))
-    TOKEN_WHITE_LIST: list[str]  # 白名单用户ID，可多点登录
 
     # FastAPI
     API_V1_STR: str = '/v1'
@@ -58,6 +57,9 @@ def validator_api_url(cls, values):
     # Location Parse
     LOCATION_PARSE: Literal['online', 'offline', 'false'] = 'offline'
 
+    # Limiter
+    LIMITER_REDIS_PREFIX: str = 'fba_limiter'
+
     # MySQL
     DB_ECHO: bool = False
     DB_DATABASE: str = 'fba'
@@ -77,8 +79,10 @@ def validator_api_url(cls, values):
     # Token
     TOKEN_ALGORITHM: str = 'HS256'  # 算法
     TOKEN_EXPIRE_SECONDS: int = 60 * 60 * 24 * 1  # 过期时间，单位：秒
+    TOKEN_REFRESH_EXPIRE_SECONDS: int = 60 * 60 * 24 * 7  # 刷新过期时间，单位：秒
     TOKEN_URL_SWAGGER: str = f'{API_V1_STR}/auth/swagger_login'
     TOKEN_REDIS_PREFIX: str = 'fba_token'
+    TOKEN_REFRESH_REDIS_PREFIX: str = 'fba_refresh_token'
 
     # Log
     LOG_STDOUT_FILENAME: str = 'fba_access.log'

backend/app/core/registrar.py

@@ -28,7 +28,7 @@ async def register_init(app: FastAPI):
     # 连接 redis
     await redis_client.open()
     # 初始化 limiter
-    await FastAPILimiter.init(redis_client, prefix='fba_limiter')
+    await FastAPILimiter.init(redis_client, prefix=settings.LIMITER_REDIS_PREFIX)
     # 启动定时任务
     scheduler.start()
 

backend/app/crud/crud_user.py

@@ -24,6 +24,7 @@ async def get_by_username(self, db: AsyncSession, username: str) -> User | None:
 
     async def update_login_time(self, db: AsyncSession, username: str, login_time: datetime) -> int:
         user = await db.execute(update(self.model).where(self.model.username == username).values(last_login=login_time))
+        await db.commit()
         return user.rowcount
 
     async def create(self, db: AsyncSession, create: CreateUser) -> NoReturn:
@@ -91,6 +92,10 @@ async def get_active(self, db: AsyncSession, user_id: int) -> bool:
         user = await self.get(db, user_id)
         return user.is_active
 
+    async def get_multi_login(self, db: AsyncSession, user_id: int) -> bool:
+        user = await self.get(db, user_id)
+        return user.is_multi_login
+
     async def set_super(self, db: AsyncSession, user_id: int) -> int:
         super_status = await self.get_super(db, user_id)
         user = await db.execute(
@@ -105,6 +110,13 @@ async def set_active(self, db: AsyncSession, user_id: int) -> int:
         )
         return user.rowcount
 
+    async def set_multi_login(self, db: AsyncSession, user_id: int) -> int:
+        multi_login = await self.get_multi_login(db, user_id)
+        user = await db.execute(
+            update(self.model).where(self.model.id == user_id).values(is_multi_login=False if multi_login else True)
+        )
+        return user.rowcount
+
     async def get_role_ids(self, db: AsyncSession, user_id: int) -> list[int]:
         user = await db.execute(
             select(self.model).where(self.model.id == user_id).options(selectinload(self.model.roles))

backend/app/models/sys_user.py

@@ -23,6 +23,7 @@ class User(DataClassBase):
     email: Mapped[str] = mapped_column(String(50), unique=True, index=True, comment='邮箱')
     is_superuser: Mapped[bool] = mapped_column(default=False, comment='超级权限(0否 1是)')
     is_active: Mapped[bool] = mapped_column(default=True, comment='用户账号状态(0停用 1正常)')
+    is_multi_login: Mapped[bool] = mapped_column(default=False, comment='是否重复登陆(0否 1是)')
     avatar: Mapped[str | None] = mapped_column(String(255), default=None, comment='头像')
     phone: Mapped[str | None] = mapped_column(String(11), default=None, comment='手机号')
     time_joined: Mapped[datetime] = mapped_column(init=False, default=func.now(), comment='注册时间')

backend/app/schemas/token.py

@@ -1,9 +1,8 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from datetime import datetime, timedelta
+from datetime import datetime
 
-from pydantic import BaseModel, Field, validator
-from pydantic.datetime_parse import parse_datetime
+from pydantic import BaseModel
 
 from backend.app.schemas.user import GetUserInfoNoRelation
 
@@ -14,10 +13,13 @@ class SwaggerToken(BaseModel):
     user: GetUserInfoNoRelation
 
 
-class LoginToken(BaseModel):
+class AccessToken(BaseModel):
     access_token: str
     access_token_type: str = 'Bearer'
     access_token_expire_time: datetime
+
+
+class LoginToken(AccessToken):
     refresh_token: str
     refresh_token_type: str = 'Bearer'
     refresh_token_expire_time: datetime
@@ -30,20 +32,5 @@ class RefreshToken(BaseModel):
     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
+class NewToken(AccessToken):
+    pass