feat(auth): 新增ws登录认证、区分http和ws认证

core/routers/auth/services.py

@@ -5,11 +5,11 @@
 # @Author : xingc
 # @Desc :
 from datetime import datetime, timedelta
-from typing import Annotated
+from typing import Annotated, Optional
 
 import jwt
 import ujson
-from fastapi import HTTPException, status, Depends
+from fastapi import HTTPException, status, Depends, Request, WebSocket
 from fastapi.security import OAuth2PasswordBearer
 import redis.asyncio as aioredis
 from sqlalchemy.ext.asyncio import AsyncSession
@@ -47,11 +47,11 @@ def verify_access_token(token: str):
 
 
 async def get_current_user(
-        token: Annotated[str, Depends(oauth2_scheme)],
-        db_session: Annotated[AsyncSession, Depends(get_async_session)],
-        redis_conn: aioredis = Depends(get_cache)
+        token: str | None, db_session: AsyncSession, redis_conn: aioredis
 ) -> schema.LoginUser:
     """验证登录状态、获取的用户基础信息"""
+
+
     result = verify_access_token(token)
     user_id = result.get("sub")
     if user_id is None:
@@ -67,3 +67,28 @@ async def get_current_user(
         user = schema.LoginUser(**ujson.decode(user_str))
 
     return user
+
+
+async def get_current_user_form_http(
+        token: Annotated[str, Depends(oauth2_scheme)],
+        db_session: Annotated[AsyncSession, Depends(get_async_session)],
+        redis_conn: aioredis = Depends(get_cache),
+) -> schema.LoginUser:
+    """在http headers中提取token，并验证再获取当前用户"""
+    user = await get_current_user(token, db_session, redis_conn)
+    return user
+
+
+async def get_current_user_form_ws(
+        websocket: WebSocket,
+        db_session: Annotated[AsyncSession, Depends(get_async_session)] = None,
+        redis_conn: aioredis = Depends(get_cache),
+) -> schema.LoginUser:
+    """在ws headers中提取token，并验证再获取当前用户"""
+    auth_header = websocket.headers.get("authorization") or websocket.headers.get("Authorization")
+    if auth_header and auth_header.lower().startswith("bearer "):
+        token = auth_header.removeprefix("Bearer ").strip()
+    else:
+        token = ''
+    user = await get_current_user(token, db_session, redis_conn)
+    return user