Optimize token detection and caching logic (#677)

This commit is contained in:
Wu Clan
2025-06-21 20:18:13 +08:00
committed by GitHub
parent 8638c26db1
commit bac41a46f8
7 changed files with 105 additions and 85 deletions
+9 -7
View File
@@ -34,13 +34,6 @@ class RequestCallNext:
response: Response
@dataclasses.dataclass
class NewToken:
new_access_token: str
new_access_token_expire_time: datetime
session_uuid: str
@dataclasses.dataclass
class AccessToken:
access_token: str
@@ -54,6 +47,15 @@ class RefreshToken:
refresh_token_expire_time: datetime
@dataclasses.dataclass
class NewToken:
new_access_token: str
new_access_token_expire_time: datetime
new_refresh_token: str
new_refresh_token_expire_time: datetime
session_uuid: str
@dataclasses.dataclass
class TokenPayload:
id: int
+42 -25
View File
@@ -60,11 +60,7 @@ def jwt_encode(payload: dict[str, Any]) -> str:
:param payload: 载荷
:return:
"""
return jwt.encode(
payload,
settings.TOKEN_SECRET_KEY,
settings.TOKEN_ALGORITHM,
)
return jwt.encode(payload, settings.TOKEN_SECRET_KEY, settings.TOKEN_ALGORITHM)
def jwt_decode(token: str) -> TokenPayload:
@@ -75,20 +71,27 @@ def jwt_decode(token: str) -> TokenPayload:
:return:
"""
try:
payload = jwt.decode(token, settings.TOKEN_SECRET_KEY, algorithms=[settings.TOKEN_ALGORITHM])
session_uuid = payload.get('session_uuid') or 'debug'
payload = jwt.decode(
token,
settings.TOKEN_SECRET_KEY,
algorithms=[settings.TOKEN_ALGORITHM],
options={'verify_exp': True},
)
session_uuid = payload.get('session_uuid')
user_id = payload.get('sub')
expire_time = payload.get('exp')
if not user_id:
expire = payload.get('exp')
if not session_uuid or not user_id or not expire:
raise errors.TokenError(msg='Token 无效')
except ExpiredSignatureError:
raise errors.TokenError(msg='Token 已过期')
except (JWTError, Exception):
raise errors.TokenError(msg='Token 无效')
return TokenPayload(id=int(user_id), session_uuid=session_uuid, expire_time=expire_time)
return TokenPayload(
id=int(user_id), session_uuid=session_uuid, expire_time=timezone.from_datetime(timezone.to_utc(expire))
)
async def create_access_token(user_id: str, multi_login: bool, **kwargs) -> AccessToken:
async def create_access_token(user_id: int, multi_login: bool, **kwargs) -> AccessToken:
"""
生成加密 token
@@ -101,8 +104,8 @@ async def create_access_token(user_id: str, multi_login: bool, **kwargs) -> Acce
session_uuid = str(uuid4())
access_token = jwt_encode({
'session_uuid': session_uuid,
'exp': expire,
'sub': user_id,
'exp': timezone.to_utc(expire).timestamp(),
'sub': str(user_id),
})
if not multi_login:
@@ -117,7 +120,7 @@ async def create_access_token(user_id: str, multi_login: bool, **kwargs) -> Acce
# Token 附加信息单独存储
if kwargs:
await redis_client.setex(
f'{settings.TOKEN_EXTRA_INFO_REDIS_PREFIX}:{session_uuid}',
f'{settings.TOKEN_EXTRA_INFO_REDIS_PREFIX}:{user_id}:{session_uuid}',
settings.TOKEN_EXPIRE_SECONDS,
json.dumps(kwargs, ensure_ascii=False),
)
@@ -125,51 +128,65 @@ async def create_access_token(user_id: str, multi_login: bool, **kwargs) -> Acce
return AccessToken(access_token=access_token, access_token_expire_time=expire, session_uuid=session_uuid)
async def create_refresh_token(user_id: str, multi_login: bool) -> RefreshToken:
async def create_refresh_token(session_uuid: str, user_id: int, multi_login: bool) -> RefreshToken:
"""
生成加密刷新 token,仅用于创建新的 token
:param session_uuid: 会话 UUID
:param user_id: 用户 ID
:param multi_login: 是否允许多端登录
:return:
"""
expire = timezone.now() + timedelta(seconds=settings.TOKEN_REFRESH_EXPIRE_SECONDS)
refresh_token = jwt_encode({'exp': expire, 'sub': user_id})
refresh_token = jwt_encode({
'session_uuid': session_uuid,
'exp': timezone.to_utc(expire).timestamp(),
'sub': str(user_id),
})
if not multi_login:
key_prefix = f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}'
await redis_client.delete_prefix(key_prefix)
await redis_client.delete_prefix(f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}')
await redis_client.setex(
f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{refresh_token}',
f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{session_uuid}',
settings.TOKEN_REFRESH_EXPIRE_SECONDS,
refresh_token,
)
return RefreshToken(refresh_token=refresh_token, refresh_token_expire_time=expire)
async def create_new_token(user_id: str, refresh_token: str, multi_login: bool, **kwargs) -> NewToken:
async def create_new_token(
refresh_token: str, session_uuid: str, user_id: int, multi_login: bool, **kwargs
) -> NewToken:
"""
生成新的 token
:param user_id: 用户 ID
:param refresh_token: 刷新 token
:param session_uuid: 会话 UUID
:param user_id: 用户 ID
:param multi_login: 是否允许多端登录
:param kwargs: token 附加信息
:return:
"""
redis_refresh_token = await redis_client.get(f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{refresh_token}')
redis_refresh_token = await redis_client.get(f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{session_uuid}')
if not redis_refresh_token or redis_refresh_token != refresh_token:
raise errors.TokenError(msg='Refresh Token 已过期,请重新登录')
await redis_client.delete(f'{settings.TOKEN_REFRESH_REDIS_PREFIX}:{user_id}:{session_uuid}')
await redis_client.delete(f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{session_uuid}')
new_access_token = await create_access_token(user_id, multi_login, **kwargs)
new_refresh_token = await create_refresh_token(new_access_token.session_uuid, user_id, multi_login)
return NewToken(
new_access_token=new_access_token.access_token,
new_access_token_expire_time=new_access_token.access_token_expire_time,
new_refresh_token=new_refresh_token.refresh_token,
new_refresh_token_expire_time=new_refresh_token.refresh_token_expire_time,
session_uuid=new_access_token.session_uuid,
)
async def revoke_token(user_id: str, session_uuid: str) -> None:
async def revoke_token(user_id: int, session_uuid: str) -> None:
"""
撤销 token
@@ -177,8 +194,8 @@ async def revoke_token(user_id: str, session_uuid: str) -> None:
:param session_uuid: 会话 ID
:return:
"""
token_key = f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{session_uuid}'
await redis_client.delete(token_key)
await redis_client.delete(f'{settings.TOKEN_REDIS_PREFIX}:{user_id}:{session_uuid}')
await redis_client.delete(f'{settings.TOKEN_EXTRA_INFO_REDIS_PREFIX}:{user_id}:{session_uuid}')
def get_token(request: Request) -> str: