From a39cdb93ee06a766453277e24185b41fd9c6f66e Mon Sep 17 00:00:00 2001 From: Maryam Abdoli Date: Thu, 2 Nov 2023 10:26:10 -0400 Subject: [PATCH 1/5] modify get_current_user to accept api_key as authentication method --- src/backend/langflow/services/auth/utils.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index f88a1cd12..54f5733b8 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -4,6 +4,7 @@ from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer from jose import JWTError, jwt from typing import Annotated, Coroutine, Optional, Union from uuid import UUID +from starlette.requests import Request 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 @@ -69,6 +70,25 @@ async def api_key_security( async def get_current_user( + request: Request, + query_param: str = Security(api_key_query), + header_param: str = Security(api_key_header), + db: Session = Depends(get_session), +) -> User: + try: + token = oauth2_login(request) + return await get_current_user_by_jwt(token, db) + except HTTPException: + user = await api_key_security(query_param, header_param, db) + if user: + return user + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="Invalid or missing API key", + ) + + +async def get_current_user_by_jwt( token: Annotated[str, Depends(oauth2_login)], db: Session = Depends(get_session), ) -> User: From 5686a6fd6374bb4d25c98af7e3e353e8850f1313 Mon Sep 17 00:00:00 2001 From: Maryam Abdoli Date: Thu, 2 Nov 2023 16:35:00 -0400 Subject: [PATCH 2/5] pass the oauth2 token to the fastapi Security --- src/backend/langflow/services/auth/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index 54f5733b8..3d014970d 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -70,13 +70,12 @@ async def api_key_security( async def get_current_user( - request: Request, + token: str = Security(oauth2_login), query_param: str = Security(api_key_query), header_param: str = Security(api_key_header), db: Session = Depends(get_session), ) -> User: try: - token = oauth2_login(request) return await get_current_user_by_jwt(token, db) except HTTPException: user = await api_key_security(query_param, header_param, db) From 83425b09cd289d15161799768597018652b96c10 Mon Sep 17 00:00:00 2001 From: Maryam Abdoli Date: Fri, 3 Nov 2023 09:29:24 -0400 Subject: [PATCH 3/5] handle all exceptions during user retrieving --- src/backend/langflow/services/auth/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index 3d014970d..9881b09f2 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -4,7 +4,6 @@ from fastapi.security import APIKeyHeader, APIKeyQuery, OAuth2PasswordBearer from jose import JWTError, jwt from typing import Annotated, Coroutine, Optional, Union from uuid import UUID -from starlette.requests import Request 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 @@ -77,7 +76,9 @@ async def get_current_user( ) -> User: try: return await get_current_user_by_jwt(token, db) - except HTTPException: + except HTTPException as exc: + if not query_param and not header_param: + raise exc user = await api_key_security(query_param, header_param, db) if user: return user @@ -85,6 +86,11 @@ async def get_current_user( status_code=status.HTTP_403_FORBIDDEN, detail="Invalid or missing API key", ) + except Exception as exc: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail=f"Internal server error: {exc}", + ) async def get_current_user_by_jwt( From ef949838e26e97166605f976ac079c85433e86c8 Mon Sep 17 00:00:00 2001 From: Maryam Abdoli Date: Fri, 3 Nov 2023 13:01:28 -0400 Subject: [PATCH 4/5] fix lints and also the bug in the get_current_user --- src/backend/langflow/api/v1/chat.py | 7 +++---- src/backend/langflow/services/auth/utils.py | 17 ++++++++--------- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index 0f277fbb7..0647ef71f 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -2,7 +2,6 @@ from fastapi import ( APIRouter, Depends, HTTPException, - Query, WebSocket, WebSocketException, status, @@ -11,8 +10,9 @@ from fastapi.responses import StreamingResponse from langflow.api.utils import build_input_keys_response from langflow.api.v1.schemas import BuildStatus, BuiltResponse, InitResponse, StreamData +from langflow.services.database.models.user.user import User from langflow.graph.graph.base import Graph -from langflow.services.auth.utils import get_current_active_user, get_current_user +from langflow.services.auth.utils import get_current_active_user from langflow.services.cache.utils import update_build_status from loguru import logger from langflow.services.getters import get_chat_service, get_session, get_cache_service @@ -28,14 +28,13 @@ router = APIRouter(tags=["Chat"]) async def chat( client_id: str, websocket: WebSocket, - token: str = Query(...), db: Session = Depends(get_session), chat_service: "ChatService" = Depends(get_chat_service), + user: User = Depends(get_current_active_user), ): """Websocket endpoint for chat.""" try: await websocket.accept() - user = await get_current_user(token, db) if not user: await websocket.close( code=status.WS_1008_POLICY_VIOLATION, reason="Unauthorized" diff --git a/src/backend/langflow/services/auth/utils.py b/src/backend/langflow/services/auth/utils.py index 9881b09f2..7cc91c117 100644 --- a/src/backend/langflow/services/auth/utils.py +++ b/src/backend/langflow/services/auth/utils.py @@ -15,7 +15,7 @@ from langflow.services.database.models.user.crud import ( from langflow.services.getters import get_session, get_settings_service from sqlmodel import Session -oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login") +oauth2_login = OAuth2PasswordBearer(tokenUrl="api/v1/login", auto_error=False) API_KEY_NAME = "x-api-key" @@ -74,23 +74,22 @@ async def get_current_user( header_param: str = Security(api_key_header), db: Session = Depends(get_session), ) -> User: - try: + if token: return await get_current_user_by_jwt(token, db) - except HTTPException as exc: + else: if not query_param and not header_param: - raise exc + raise HTTPException( + status_code=status.HTTP_403_FORBIDDEN, + detail="An API key must be passed as query or header", + ) user = await api_key_security(query_param, header_param, db) if user: return user + raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail="Invalid or missing API key", ) - except Exception as exc: - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=f"Internal server error: {exc}", - ) async def get_current_user_by_jwt( From b5b66ef077112687be9e2aa81c1341c6a73c368d Mon Sep 17 00:00:00 2001 From: Maryam Abdoli Date: Fri, 3 Nov 2023 18:29:51 -0400 Subject: [PATCH 5/5] fix tests by reverting the chat endpoint dependecies --- src/backend/langflow/api/v1/chat.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py index 0647ef71f..e1e588d8f 100644 --- a/src/backend/langflow/api/v1/chat.py +++ b/src/backend/langflow/api/v1/chat.py @@ -2,6 +2,7 @@ from fastapi import ( APIRouter, Depends, HTTPException, + Query, WebSocket, WebSocketException, status, @@ -10,9 +11,11 @@ from fastapi.responses import StreamingResponse from langflow.api.utils import build_input_keys_response from langflow.api.v1.schemas import BuildStatus, BuiltResponse, InitResponse, StreamData -from langflow.services.database.models.user.user import User from langflow.graph.graph.base import Graph -from langflow.services.auth.utils import get_current_active_user +from langflow.services.auth.utils import ( + get_current_active_user, + get_current_user_by_jwt, +) from langflow.services.cache.utils import update_build_status from loguru import logger from langflow.services.getters import get_chat_service, get_session, get_cache_service @@ -28,12 +31,13 @@ router = APIRouter(tags=["Chat"]) async def chat( client_id: str, websocket: WebSocket, + token: str = Query(...), db: Session = Depends(get_session), chat_service: "ChatService" = Depends(get_chat_service), - user: User = Depends(get_current_active_user), ): """Websocket endpoint for chat.""" try: + user = await get_current_user_by_jwt(token, db) await websocket.accept() if not user: await websocket.close(