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

This commit is contained in:
xingc
2025-04-16 14:42:34 +08:00
parent 94ddc6b926
commit 42e6b66984
4 changed files with 46 additions and 21 deletions

View File

@@ -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

View File

@@ -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)
):
# 删除关联关系

View File

@@ -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:

View File

@@ -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)