diff --git a/core/routers/auth/services.py b/core/routers/auth/services.py index ce19338..16c17f3 100644 --- a/core/routers/auth/services.py +++ b/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 diff --git a/core/routers/card/routes.py b/core/routers/card/routes.py index cedf939..0ecb830 100644 --- a/core/routers/card/routes.py +++ b/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) ): # 删除关联关系 diff --git a/core/routers/file/routes.py b/core/routers/file/routes.py index 684730d..2ec6a0c 100644 --- a/core/routers/file/routes.py +++ b/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: diff --git a/core/routers/user/routes.py b/core/routers/user/routes.py index 42af3b3..dbf37bd 100644 --- a/core/routers/user/routes.py +++ b/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)