add token refreshing mechanism (#62)

* add token refreshing mechanism

* update token_expires to token_expire_time

* Fix implicit type conversion exception catch

Author: wu-clan

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

@@ -1,36 +1,47 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Request
 from fastapi.security import OAuth2PasswordRequestForm
 
-from backend.app.common.jwt import DependsUser, JwtAuthentication
-from backend.app.common.redis import redis_client
+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 Token
+from backend.app.schemas.token import RefreshToken, LoginToken, SwaggerToken
 from backend.app.schemas.user import Auth
 from backend.app.services.user_service import UserService
 
 router = APIRouter()
 
 
 @router.post('/swagger_login', summary='swagger 表单登录', description='form 格式登录，仅用于 swagger 文档调试接口')
-async def swagger_user_login(form_data: OAuth2PasswordRequestForm = Depends()) -> Token:
+async def swagger_user_login(form_data: OAuth2PasswordRequestForm = Depends()) -> SwaggerToken:
     token, user = await UserService.swagger_login(form_data)
-    return Token(access_token=token, user=user)
+    return SwaggerToken(access_token=token, user=user)
 
 
 @router.post('/login', summary='用户登录', description='json 格式登录, 仅支持在第三方api工具调试接口, 例如: postman')
 async def user_login(obj: Auth):
-    token, user = await UserService.login(obj)
-    # TODO: token 存储
-    data = Token(access_token=token, user=user)
+    access_token, refresh_token, access_expire, refresh_expire, user = await UserService.login(obj)
+    data = LoginToken(
+        access_token=access_token,
+        refresh_token=refresh_token,
+        access_token_expire_time=access_expire,
+        refresh_token_expire_time=refresh_expire,
+        user=user,
+    )
+    return response_base.success(data=data)
+
+
+@router.post('/refresh_token', summary='刷新 token', dependencies=[DependsUser])
+async def get_refresh_token(request: Request):
+    token = get_token(request)
+    user_id, _ = jwt_decode(token)
+    refresh_token, refresh_expire = await UserService.refresh_token(user_id)
+    data = RefreshToken(refresh_token=refresh_token, refresh_token_expire_time=refresh_expire)
     return response_base.success(data=data)
 
 
 @router.post('/logout', summary='用户登出', dependencies=[DependsUser])
-async def user_logout(jwt: JwtAuthentication):
-    user_id = jwt.get('payload').get('sub')
-    token = jwt.get('token')
-    key = f'token:{user_id}:{token}'
-    await redis_client.delete(key)
+async def user_logout(jwt: CurrentJwtAuth):
+    user_id = jwt.get('sub')
+    await UserService.logout(user_id)
     return response_base.success()

backend/app/common/casbin_rbac.py

@@ -22,6 +22,7 @@ def get_casbin_enforcer(self) -> casbin.Enforcer:
 
         :return:
         """
+        # TODO: https://github.com/pycasbin/async-sqlalchemy-adapter/issues/4
         adapter = casbin_sqlalchemy_adapter.Adapter(self._CASBIN_DATABASE_URL, db_class=CasbinRule)
 
         enforcer = casbin.Enforcer(RBAC_MODEL_CONF, adapter)

backend/app/common/jwt.py

@@ -1,10 +1,10 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 from datetime import datetime, timedelta
-from typing import Any
 
-from fastapi import Depends
+from fastapi import Depends, Request
 from fastapi.security import OAuth2PasswordBearer
+from fastapi.security.utils import get_authorization_scheme_param
 from jose import jwt
 from passlib.context import CryptContext
 from pydantic import ValidationError
@@ -43,7 +43,7 @@ def password_verify(plain_password: str, hashed_password: str) -> bool:
     return pwd_context.verify(plain_password, hashed_password)
 
 
-async def create_access_token(sub: int | Any, expires_delta: timedelta | None = None, **kwargs) -> str:
+async def create_access_token(sub: str, expires_delta: timedelta | None = None, **kwargs) -> tuple[str, datetime]:
     """
     Generate encryption token
 
@@ -52,41 +52,86 @@ async def create_access_token(sub: int | Any, expires_delta: timedelta | None =
     :return:
     """
     if expires_delta:
-        expires = datetime.utcnow() + expires_delta
-        expire_seconds = expires_delta.total_seconds()
+        expire = datetime.utcnow() + expires_delta
+        expire_seconds = int(expires_delta.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': str(sub), **kwargs}
+    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(f'token:{sub}:*')
-    key = f'token:{sub}:{token}'
+        await redis_client.delete_prefix(f'{settings.TOKEN_REDIS_PREFIX}:{sub}:')
+    key = f'{settings.TOKEN_REDIS_PREFIX}:{sub}:{token}'
     await redis_client.setex(key, expire_seconds, token)
-    return token
+    return token, expire
 
 
-async def jwt_authentication(token: str = Depends(oauth2_schema)):
+async def create_refresh_token(sub: str, expire_time: datetime | None = None, **kwargs) -> tuple[str, datetime]:
     """
-    JWT authentication
+    Generate encryption refresh token
+
+    :param sub: The subject/userid of the JWT
+    :param expire_time: expiry time
+    :return:
+    """
+    if expire_time:
+        expires = expire_time + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
+        expire_seconds = int((expires - datetime.utcnow()).total_seconds())
+    else:
+        expires = datetime.utcnow() + timedelta(seconds=settings.TOKEN_EXPIRE_SECONDS)
+        expire_seconds = settings.TOKEN_EXPIRE_SECONDS
+    to_encode = {'exp': expires, '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
+
+
+def get_token(request: Request) -> str:
+    """
+    Get token for request header
+
+    :return:
+    """
+    authorization = request.headers.get('Authorization')
+    scheme, param = get_authorization_scheme_param(authorization)
+    if not authorization or scheme.lower() != 'bearer':
+        raise TokenError
+    return param
+
+
+def jwt_decode(token: str) -> tuple[int, list[int]]:
+    """
+    Decode token
 
     :param token:
     :return:
     """
     try:
         payload = jwt.decode(token, settings.TOKEN_SECRET_KEY, algorithms=[settings.TOKEN_ALGORITHM])
-        user_id = payload.get('sub')
-        user_role = payload.get('role_ids')
-        if not user_id or not user_role:
-            raise TokenError
-        # 验证token是否有效
-        key = f'token:{user_id}:{token}'
-        valid_token = await redis_client.get(key)
-        if not valid_token:
+        user_id = int(payload.get('sub'))
+        user_roles = list(payload.get('role_ids'))
+        if not user_id or not user_roles:
             raise TokenError
-        return {'payload': payload, 'token': token}
-    except (jwt.JWTError, ValidationError):
+    except (jwt.JWTError, ValidationError, Exception):
+        raise TokenError
+    return user_id, user_roles
+
+
+async def jwt_authentication(token: str = Depends(oauth2_schema)) -> dict[str, int]:
+    """
+    JWT authentication
+
+    :param token:
+    :return:
+    """
+    user_id, _ = jwt_decode(token)
+    key = f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{token}'
+    token_verify = await redis_client.get(key)
+    if not token_verify:
         raise TokenError
+    return {'sub': user_id}
 
 
 async def get_current_user(db: CurrentSession, data: dict = Depends(jwt_authentication)) -> User:
@@ -97,7 +142,7 @@ async def get_current_user(db: CurrentSession, data: dict = Depends(jwt_authenti
     :param data:
     :return:
     """
-    user_id = data.get('payload').get('sub')
+    user_id = data.get('sub')
     user = await UserDao.get_user_with_relation(db, user_id=user_id)
     if not user:
         raise TokenError
@@ -121,7 +166,7 @@ async def get_current_is_superuser(user: User = Depends(get_current_user)):
 CurrentUser = Annotated[User, Depends(get_current_user)]
 CurrentSuperUser = Annotated[bool, Depends(get_current_is_superuser)]
 # Token dependency injection
-JwtAuthentication = Annotated[dict, Depends(jwt_authentication)]
+CurrentJwtAuth = Annotated[dict, Depends(jwt_authentication)]
 # Permission dependency injection
 DependsUser = Depends(get_current_user)
 DependsSuperUser = Depends(get_current_is_superuser)

backend/app/common/redis.py

@@ -37,6 +37,17 @@ async def open(self):
             log.error('❌ 数据库 redis 连接异常 {}', e)
             sys.exit()
 
+    async def delete_prefix(self, key: str):
+        """
+        删除指定前缀的所有key
+
+        :param key:
+        :return:
+        """
+        keys = await self.keys(f'{key}*')
+        if keys:
+            await self.delete(*keys)
+
 
 # 创建redis连接对象
 redis_client = RedisCli()

backend/app/core/conf.py

@@ -73,6 +73,7 @@ def validator_api_url(cls, values):
     TOKEN_ALGORITHM: str = 'HS256'  # 算法
     TOKEN_EXPIRE_SECONDS: int = 60 * 60 * 24 * 1  # 过期时间，单位：秒
     TOKEN_URL_SWAGGER: str = '/v1/auth/users/swagger_login'
+    TOKEN_REDIS_PREFIX: str = 'fba_token'
 
     # Log
     LOG_FILE_NAME: str = 'fba.log'

backend/app/crud/crud_user.py

@@ -55,7 +55,7 @@ async def update_avatar(self, db: AsyncSession, current_user: User, avatar: Avat
     async def delete_user(self, db: AsyncSession, user_id: int) -> int:
         return await self.delete(db, user_id)
 
-    async def check_email(self, db: AsyncSession, email: str) -> User:
+    async def check_email(self, db: AsyncSession, email: str) -> User | None:
         mail = await db.execute(select(self.model).where(self.model.email == email))
         return mail.scalars().first()
 
@@ -101,7 +101,9 @@ async def get_user_role_ids(self, db: AsyncSession, user_id: int) -> list[int]:
         roles_id = [role.id for role in user.scalars().first().roles]
         return roles_id
 
-    async def get_user_with_relation(self, db: AsyncSession, *, user_id: int = None, username: str = None) -> User:
+    async def get_user_with_relation(
+        self, db: AsyncSession, *, user_id: int = None, username: str = None
+    ) -> User | None:
         where = 'condition'
         if user_id:
             where = 'self.model.id == user_id'

backend/app/schemas/token.py

@@ -1,11 +1,29 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
+from datetime import datetime
+
 from pydantic import BaseModel
 
 from backend.app.schemas.user import GetUserInfoNoRelation
 
 
-class Token(BaseModel):
+class SwaggerToken(BaseModel):
     access_token: str
     token_type: str = 'Bearer'
     user: GetUserInfoNoRelation
+
+
+class LoginToken(BaseModel):
+    access_token: str
+    access_token_type: str = 'Bearer'
+    access_token_expire_time: datetime
+    refresh_token: str
+    refresh_token_type: str = 'Bearer'
+    refresh_token_expire_time: datetime
+    user: GetUserInfoNoRelation
+
+
+class RefreshToken(BaseModel):
+    refresh_token: str
+    refresh_token_type: str = 'Bearer'
+    refresh_token_expire_time: datetime

backend/app/services/user_service.py

@@ -5,6 +5,8 @@
 
 from backend.app.common import jwt
 from backend.app.common.exception import errors
+from backend.app.common.redis import redis_client
+from backend.app.core.conf import settings
 from backend.app.crud.crud_dept import DeptDao
 from backend.app.crud.crud_role import RoleDao
 from backend.app.crud.crud_user import UserDao
@@ -20,36 +22,59 @@ async def swagger_login(form_data: OAuth2PasswordRequestForm):
         async with async_db_session() as db:
             current_user = await UserDao.get_user_by_username(db, form_data.username)
             if not current_user:
-                raise errors.NotFoundError(msg='用户名不存在')
+                raise errors.NotFoundError(msg='用户不存在')
             elif not jwt.password_verify(form_data.password, current_user.password):
                 raise errors.AuthorizationError(msg='密码错误')
             elif not current_user.is_active:
-                raise errors.AuthorizationError(msg='该用户已被锁定，无法登录')
+                raise errors.AuthorizationError(msg='用户已锁定, 登陆失败')
             # 更新登陆时间
             await UserDao.update_user_login_time(db, form_data.username)
             # 查询用户角色
             user_role_ids = await UserDao.get_user_role_ids(db, current_user.id)
             # 获取最新用户信息
             user = await UserDao.get_user_by_id(db, current_user.id)
             # 创建token
-            access_token = await jwt.create_access_token(user.id, role_ids=user_role_ids)
+            access_token, _ = await jwt.create_access_token(str(user.id), role_ids=user_role_ids)
             return access_token, user
 
     @staticmethod
     async def login(obj: Auth):
         async with async_db_session() as db:
             current_user = await UserDao.get_user_by_username(db, obj.username)
             if not current_user:
-                raise errors.NotFoundError(msg='用户名不存在')
+                raise errors.NotFoundError(msg='用户不存在')
             elif not jwt.password_verify(obj.password, current_user.password):
                 raise errors.AuthorizationError(msg='密码错误')
             elif not current_user.is_active:
-                raise errors.AuthorizationError(msg='该用户已被锁定，无法登录')
+                raise errors.AuthorizationError(msg='用户已锁定, 登陆失败')
             await UserDao.update_user_login_time(db, obj.username)
             user_role_ids = await UserDao.get_user_role_ids(db, current_user.id)
             user = await UserDao.get_user_by_id(db, current_user.id)
-            access_token = await jwt.create_access_token(user.id, role_ids=user_role_ids)
-            return access_token, user
+            access_token, access_token_expire_time = await jwt.create_access_token(str(user.id), role_ids=user_role_ids)
+            refresh_token, refresh_token_expire_time = await jwt.create_refresh_token(
+                str(user.id), access_token_expire_time, role_ids=user_role_ids
+            )
+            return access_token, refresh_token, access_token_expire_time, refresh_token_expire_time, user
+
+    @staticmethod
+    async def refresh_token(user_id: int):
+        async with async_db_session() as db:
+            current_user = await UserDao.get_user_by_id(db, user_id)
+            if not current_user:
+                raise errors.NotFoundError(msg='用户不存在')
+            elif not current_user.is_active:
+                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
+            )
+            return refresh_token, refresh_token_expire_time
+
+    @staticmethod
+    async def logout(user_id: int):
+        key = f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:'
+        await redis_client.delete_prefix(key)
+        return
 
     @staticmethod
     async def register(obj: CreateUser):