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)],
+        token: str | None, db_session: AsyncSession, redis_conn: aioredis
-        db_session: Annotated[AsyncSession, Depends(get_async_session)],
-        redis_conn: aioredis = Depends(get_cache)
 ) -> 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

core/routers/card/routes.py

@@ -19,7 +19,7 @@ from core.db.crud import CRUD
 from core.db.engine import get_async_session, get_cache
 from core.db.models import Card, SearchCardCache, UserCollection
 from core.routers.auth.schema import LoginUser
-from core.routers.auth.services import get_current_user
+from core.routers.auth.services import get_current_user_form_http
 from core.utils.helper import generate_keyname, check_picture_for_empty
 
 from . import schema
@@ -85,7 +85,7 @@ async def get_card_search(
     ],
 )
 async def get_me_all_card(
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         page: int = Query(default=1, gt=0),
         page_size: int = Query(default=15, gt=0, le=200),
         category: Literal['all', 'anime', 'sports', 'other', 'displayed'] = Query(default='all'),
@@ -143,7 +143,7 @@ async def get_me_all_card(
 )
 async def get_visit_user_all_card(
         user_id: int,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         page: int = Query(default=1, gt=0),
         page_size: int = Query(default=15, gt=0, le=200),
         db_session: AsyncSession = Depends(get_async_session)
@@ -186,7 +186,7 @@ async def get_visit_user_all_card(
 @card_router.post('/create', summary='新增卡片', response_model=base_schema.BaseResponse[schema.CardOut])
 async def created_card(
         payload: schema.CardIn,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     card_crud = CRUD(db_session, Card)
@@ -264,7 +264,7 @@ async def created_card(
 async def update_card_price_info(
         card_id: int,
         payload: schema.UpdateCard,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session),
 ):
     db_crud = CRUD(db_session, UserCollection)
@@ -279,7 +279,7 @@ async def update_card_price_info(
 async def display_card_status(
         card_id: int,
         action: Literal['displayed', 'undisplayed'],
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     db_crud = CRUD(db_session, UserCollection)
@@ -295,7 +295,7 @@ async def display_card_status(
 @card_router.delete('/delete/', summary='卡片删除（批量）', response_model=base_schema.BaseResponse)
 async def delete_card(
         payload: schema.DeleteCardIn,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     # 删除关联关系

core/routers/file/routes.py

@@ -17,7 +17,7 @@ from PIL import Image
 from config import settings
 from core.base.schema import BaseResponse
 from core.routers.auth.schema import LoginUser
-from core.routers.auth.services import get_current_user
+from core.routers.auth.services import get_current_user_form_http
 
 from . import schema
 
@@ -29,7 +29,7 @@ file_router = APIRouter(prefix='/file', tags=['文件存储模块'])
 async def upload_file(
         upload_type: Literal['card', 'chat', 'avatar'],
         file: UploadFile = File(...),
-        login_user: LoginUser = Depends(get_current_user)
+        login_user: LoginUser = Depends(get_current_user_form_http)
 ):
     # 验证文件类型
     if file.content_type not in settings.UPLOADED_ALLOWED_MIME_TYPES:

core/routers/user/routes.py

@@ -15,7 +15,7 @@ from core.db.crud import CRUD
 from core.db.engine import get_async_session
 from core.db.models import User, UserFollows
 from core.routers.auth.schema import LoginUser
-from core.routers.auth.services import get_current_user
+from core.routers.auth.services import get_current_user_form_http
 
 from . import schema
 
@@ -24,7 +24,7 @@ user_router = APIRouter(prefix='/user', tags=['用户模块'])
 
 @user_router.get('/me/info', summary='用户信息（我的）', response_model=base_schema.BaseResponse[schema.User])
 async def get_me_info(
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     db_crud = CRUD(db_session, User)
@@ -42,7 +42,7 @@ async def get_me_info(
 )
 async def get_visit_user_info(
         user_id: int,
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     if user_id == login_user.user_id:
@@ -65,7 +65,7 @@ async def get_visit_user_info(
     ]
 )
 async def get_all_followers(
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         page: int = Query(default=1, gt=0),
         page_size: int = Query(default=15, gt=0, le=200),
         db_session: AsyncSession = Depends(get_async_session)
@@ -104,11 +104,11 @@ async def get_all_followers(
 async def follow_user(
         user_id: int,
         action: Literal['follower', 'unfollower'],
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     if user_id == login_user.user_id:
-        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail='无法操作')
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail='无法操作')
 
     db_crud = CRUD(db_session, User)
     follower_user: User | None = await db_crud.get(filters={'user_id': user_id})
@@ -143,7 +143,7 @@ async def follow_user(
 async def search_user(
         nickname: str = Query(...),
         page_size: int = Query(default=30, gt=0, le=30),
-        login_user: LoginUser = Depends(get_current_user),
+        login_user: LoginUser = Depends(get_current_user_form_http),
         db_session: AsyncSession = Depends(get_async_session)
 ):
     db_crud = CRUD(db_session, User)