From ed9796e473f0d630df6b525ff6349f6239d4cb89 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 28 Aug 2023 07:00:45 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(api=5Fkey.py):=20change=20va?= =?UTF-8?q?riable=20name=20from=20'e'=20to=20'exc'=20for=20better=20readab?= =?UTF-8?q?ility=20and=20semantics=20=F0=9F=90=9B=20fix(endpoints.py):=20a?= =?UTF-8?q?dd=20'api=5Fkey'=20parameter=20to=20'process=5Fflow'=20function?= =?UTF-8?q?=20to=20fix=20missing=20dependency=20error=20=F0=9F=90=9B=20fix?= =?UTF-8?q?(endpoints.py):=20add=20exception=20chaining=20to=20HTTPExcepti?= =?UTF-8?q?on=20in=20'process=5Fflow'=20function=20for=20better=20error=20?= =?UTF-8?q?handling=20=F0=9F=90=9B=20fix(utils.py):=20change=20return=20ty?= =?UTF-8?q?pe=20of=20'api=5Fkey=5Fsecurity'=20function=20to=20Optional[Api?= =?UTF-8?q?Key]=20for=20better=20type=20hinting=20=F0=9F=90=9B=20fix(utils?= =?UTF-8?q?.py):=20change=20variable=20name=20from=20'e'=20to=20'exc'=20fo?= =?UTF-8?q?r=20better=20readability=20and=20semantics?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/api/v1/api_key.py | 4 ++-- src/backend/langflow/api/v1/endpoints.py | 21 ++++++++++++++++----- src/backend/langflow/services/auth/utils.py | 16 ++++++++++------ 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/src/backend/langflow/api/v1/api_key.py b/src/backend/langflow/api/v1/api_key.py index aede9c437..280f240e8 100644 --- a/src/backend/langflow/api/v1/api_key.py +++ b/src/backend/langflow/api/v1/api_key.py @@ -31,8 +31,8 @@ def get_api_keys_route( keys = get_api_keys(db, user_id) return ApiKeysResponse(total_count=len(keys), user_id=user_id, api_keys=keys) - except Exception as e: - raise HTTPException(status_code=400, detail=str(e)) + except Exception as exc: + raise HTTPException(status_code=400, detail=str(exc)) from exc @router.post("/", response_model=UnmaskedApiKeyRead) diff --git a/src/backend/langflow/api/v1/endpoints.py b/src/backend/langflow/api/v1/endpoints.py index 3d2416296..423217015 100644 --- a/src/backend/langflow/api/v1/endpoints.py +++ b/src/backend/langflow/api/v1/endpoints.py @@ -85,7 +85,6 @@ def get_all(current_user: User = Depends(get_current_active_user)): @router.post( "/process/{flow_id}", response_model=ProcessResponse, - dependencies=[Depends(api_key_security)], ) async def process_flow( session: Annotated[Session, Depends(get_session)], @@ -94,13 +93,21 @@ async def process_flow( tweaks: Optional[dict] = None, clear_cache: Annotated[bool, Body(embed=True)] = False, # noqa: F821 session_id: Annotated[Union[None, str], Body(embed=True)] = None, # noqa: F821 + api_key=Depends(api_key_security), ): """ Endpoint to process an input with a given flow_id. """ try: - flow = session.get(Flow, flow_id) + api_key_user = api_key.user + # Get the flow that matches the flow_id and belongs to the user + flow = ( + session.query(Flow) + .filter(Flow.id == flow_id) + .filter(Flow.user_id == api_key_user.id) + .first() + ) if flow is None: raise ValueError(f"Flow {flow_id} not found") @@ -120,14 +127,18 @@ async def process_flow( # StatementError('(builtins.ValueError) badly formed hexadecimal UUID string') if "badly formed hexadecimal UUID string" in str(exc): # This means the Flow ID is not a valid UUID which means it can't find the flow - raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(exc) + ) from exc except ValueError as exc: if f"Flow {flow_id} not found" in str(exc): - raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(exc)) + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, detail=str(exc) + ) from exc else: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(exc) - ) + ) from exc except Exception as e: # Log stack trace logger.exception(e) diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index 0582ee382..8377b26cb 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -2,8 +2,9 @@ from datetime import datetime, timedelta, timezone from fastapi import Depends, HTTPException, Security, status from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer from jose import JWTError, jwt -from typing import Annotated, Coroutine +from typing import Annotated, Coroutine, Optional from uuid import UUID +from langflow.services.database.models.api_key.api_key import ApiKey from langflow.services.database.models.api_key.crud import check_key from langflow.services.database.models.user.user import User from langflow.services.database.models.user.crud import ( @@ -31,8 +32,9 @@ async def api_key_security( query_param: str = Security(api_key_query), header_param: str = Security(api_key_header), db: Session = Depends(get_session), -): +) -> Optional[ApiKey]: settings_manager = get_settings_manager() + result = None if settings_manager.auth_settings.AUTO_LOGIN: return settings_manager.auth_settings.API_KEY_SECRET_KEY @@ -42,12 +44,14 @@ async def api_key_security( detail="An API key must be passed as query or header", ) - elif query_param and check_key(db, query_param): - return query_param + elif query_param: + result = check_key(db, query_param) - elif header_param and check_key(db, header_param): - return header_param + else: + result = check_key(db, header_param) + if result: + return result else: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN,