From 77323d9b25dd25a0ad9b87a7701e15f45579c3e8 Mon Sep 17 00:00:00 2001 From: Anuraag Jain Date: Mon, 1 Jan 2024 10:55:50 +0200 Subject: [PATCH] refac: remove the verify_token and use get-current user for auth+user --- backend/apps/web/main.py | 22 +++------------------- backend/apps/web/routers/auths.py | 13 +++---------- backend/apps/web/routers/chats.py | 1 + backend/apps/web/routers/modelfiles.py | 8 +++----- backend/utils/utils.py | 9 ++------- 5 files changed, 12 insertions(+), 41 deletions(-) diff --git a/backend/apps/web/main.py b/backend/apps/web/main.py index b19a7c1f..b3d90381 100644 --- a/backend/apps/web/main.py +++ b/backend/apps/web/main.py @@ -3,7 +3,6 @@ from fastapi.routing import APIRoute from fastapi.middleware.cors import CORSMiddleware from apps.web.routers import auths, users, chats, modelfiles, utils from config import WEBUI_VERSION, WEBUI_AUTH -from utils.utils import verify_auth_token app = FastAPI() @@ -19,24 +18,9 @@ app.add_middleware( app.include_router(auths.router, prefix="/auths", tags=["auths"]) -app.include_router( - users.router, - prefix="/users", - tags=["users"], - dependencies=[Depends(verify_auth_token)], -) -app.include_router( - chats.router, - prefix="/chats", - tags=["chats"], - dependencies=[Depends(verify_auth_token)], -) -app.include_router( - modelfiles.router, - prefix="/modelfiles", - tags=["modelfiles"], - dependencies=[Depends(verify_auth_token)], -) +app.include_router(users.router, prefix="/users", tags=["users"]) +app.include_router(chats.router, prefix="/chats", tags=["chats"]) +app.include_router(modelfiles.router, prefix="/modelfiles", tags=["modelfiles"]) app.include_router(utils.router, prefix="/utils", tags=["utils"]) diff --git a/backend/apps/web/routers/auths.py b/backend/apps/web/routers/auths.py index 7ff645dc..24e9f426 100644 --- a/backend/apps/web/routers/auths.py +++ b/backend/apps/web/routers/auths.py @@ -19,12 +19,7 @@ from apps.web.models.auths import ( from apps.web.models.users import Users -from utils.utils import ( - get_password_hash, - get_current_user, - create_token, - verify_auth_token, -) +from utils.utils import get_password_hash, get_current_user, create_token from utils.misc import get_gravatar_url from constants import ERROR_MESSAGES @@ -36,7 +31,7 @@ router = APIRouter() ############################ -@router.get("/", response_model=UserResponse, dependencies=[Depends(verify_auth_token)]) +@router.get("/", response_model=UserResponse) async def get_session_user(user=Depends(get_current_user)): return { "id": user.id, @@ -52,9 +47,7 @@ async def get_session_user(user=Depends(get_current_user)): ############################ -@router.post( - "/update/password", response_model=bool, dependencies=[Depends(verify_auth_token)] -) +@router.post("/update/password", response_model=bool) async def update_password( form_data: UpdatePasswordForm, session_user=Depends(get_current_user) ): diff --git a/backend/apps/web/routers/chats.py b/backend/apps/web/routers/chats.py index 794eb7f7..0eec4553 100644 --- a/backend/apps/web/routers/chats.py +++ b/backend/apps/web/routers/chats.py @@ -108,6 +108,7 @@ async def delete_chat_by_id(id: str, user=Depends(get_current_user)): result = Chats.delete_chat_by_id_and_user_id(id, user.id) return result + ############################ # DeleteAllChats ############################ diff --git a/backend/apps/web/routers/modelfiles.py b/backend/apps/web/routers/modelfiles.py index c54ef4a2..841d534d 100644 --- a/backend/apps/web/routers/modelfiles.py +++ b/backend/apps/web/routers/modelfiles.py @@ -5,8 +5,6 @@ from typing import List, Union, Optional from fastapi import APIRouter from pydantic import BaseModel import json - -from apps.web.models.users import Users from apps.web.models.modelfiles import ( Modelfiles, ModelfileForm, @@ -15,7 +13,7 @@ from apps.web.models.modelfiles import ( ModelfileResponse, ) -from utils.utils import bearer_scheme, get_current_user +from utils.utils import get_current_user from constants import ERROR_MESSAGES router = APIRouter() @@ -26,7 +24,7 @@ router = APIRouter() @router.get("/", response_model=List[ModelfileResponse]) -async def get_modelfiles(skip: int = 0, limit: int = 50, cred=Depends(bearer_scheme)): +async def get_modelfiles(skip: int = 0, limit: int = 50, user=Depends(get_current_user)): return Modelfiles.get_modelfiles(skip, limit) @@ -67,7 +65,7 @@ async def create_new_modelfile( @router.post("/", response_model=Optional[ModelfileResponse]) -async def get_modelfile_by_tag_name(form_data: ModelfileTagNameForm): +async def get_modelfile_by_tag_name(form_data: ModelfileTagNameForm, user=Depends(get_current_user)): modelfile = Modelfiles.get_modelfile_by_tag_name(form_data.tag_name) if modelfile: diff --git a/backend/utils/utils.py b/backend/utils/utils.py index c8991152..f98644f5 100644 --- a/backend/utils/utils.py +++ b/backend/utils/utils.py @@ -55,7 +55,7 @@ def extract_token_from_auth_header(auth_header: str): return auth_header[len("Bearer ") :] -def verify_auth_token(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): +def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): data = decode_token(auth_token.credentials) if data != None and "email" in data: user = Users.get_user_by_email(data["email"]) @@ -64,14 +64,9 @@ def verify_auth_token(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBea status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.INVALID_TOKEN, ) - return + return user else: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail=ERROR_MESSAGES.UNAUTHORIZED, ) - - -def get_current_user(auth_token: HTTPAuthorizationCredentials = Depends(HTTPBearer())): - data = decode_token(auth_token.credentials) - return Users.get_user_by_email(data["email"])