From b6774cf3d843fafed1232693595dc3be30780103 Mon Sep 17 00:00:00 2001 From: Cristhian Zanforlin Lousa <72977554+Cristhianzl@users.noreply.github.com> Date: Thu, 25 Jul 2024 19:05:38 -0300 Subject: [PATCH] bugfix: langflow application losing store api-key on refresh page/backend (#2960) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * changing api key journey * refactor(api_key.py): remove unnecessary code related to config_dir and secret_key_path refactor(login.py): refactor setting api_key cookie to use user's store_api_key refactor(variable/service.py): re-encrypt stored value if secret_key changes to ensure validity * 📝 (api_key.py): Remove unused imports and clean up code for better readability 📝 (login.py): Remove unused imports and clean up code for better readability 📝 (auth/utils.py): Remove unused imports and clean up code for better readability 📝 (store/service.py): Remove unused imports and clean up code for better readability * 🔧 (utils.py): replace hashing logic with random key generation for key length less than 32 bytes to ensure key length is always 32 bytes 📝 (utils.py): update comments for clarity and accuracy regarding key generation and encryption process --------- Co-authored-by: Gabriel Luiz Freitas Almeida --- src/backend/base/langflow/api/v1/api_key.py | 17 ++++++++++++- src/backend/base/langflow/api/v1/login.py | 28 +++++++++++++++++++++ src/frontend/src/contexts/authContext.tsx | 1 - 3 files changed, 44 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/api/v1/api_key.py b/src/backend/base/langflow/api/v1/api_key.py index f8dab0dc1..fcf7f2786 100644 --- a/src/backend/base/langflow/api/v1/api_key.py +++ b/src/backend/base/langflow/api/v1/api_key.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING from uuid import UUID -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends, HTTPException, Response from sqlmodel import Session from langflow.api.v1.schemas import ApiKeyCreateRequest, ApiKeysResponse @@ -62,17 +62,32 @@ def delete_api_key_route( @router.post("/store") def save_store_api_key( api_key_request: ApiKeyCreateRequest, + response: Response, current_user: User = Depends(auth_utils.get_current_active_user), db: Session = Depends(get_session), settings_service=Depends(get_settings_service), ): + auth_settings = settings_service.auth_settings + try: api_key = api_key_request.api_key + # Encrypt the API key encrypted = auth_utils.encrypt_api_key(api_key, settings_service=settings_service) current_user.store_api_key = encrypted db.add(current_user) db.commit() + + response.set_cookie( + "apikey_tkn_lflw", + encrypted, + httponly=auth_settings.ACCESS_HTTPONLY, + samesite=auth_settings.ACCESS_SAME_SITE, + secure=auth_settings.ACCESS_SECURE, + expires=None, # Set to None to make it a session cookie + domain=auth_settings.COOKIE_DOMAIN, + ) + return {"detail": "API Key saved"} except Exception as e: raise HTTPException(status_code=400, detail=str(e)) from e diff --git a/src/backend/base/langflow/api/v1/login.py b/src/backend/base/langflow/api/v1/login.py index 9ba570f3c..92370cbce 100644 --- a/src/backend/base/langflow/api/v1/login.py +++ b/src/backend/base/langflow/api/v1/login.py @@ -1,5 +1,6 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.security import OAuth2PasswordRequestForm +from langflow.services.database.models.user.crud import get_user_by_id from sqlmodel import Session from langflow.api.v1.schemas import Token @@ -57,6 +58,15 @@ async def login_to_get_access_token( expires=auth_settings.ACCESS_TOKEN_EXPIRE_SECONDS, domain=auth_settings.COOKIE_DOMAIN, ) + response.set_cookie( + "apikey_tkn_lflw", + str(user.store_api_key), + httponly=auth_settings.ACCESS_HTTPONLY, + samesite=auth_settings.ACCESS_SAME_SITE, + secure=auth_settings.ACCESS_SECURE, + expires=None, # Set to None to make it a session cookie + domain=auth_settings.COOKIE_DOMAIN, + ) variable_service.initialize_user_variables(user.id, db) # Create default folder for user if it doesn't exist create_default_folder_if_it_doesnt_exist(db, user.id) @@ -74,6 +84,7 @@ async def auto_login( response: Response, db: Session = Depends(get_session), settings_service=Depends(get_settings_service) ): auth_settings = settings_service.auth_settings + if settings_service.auth_settings.AUTO_LOGIN: user_id, tokens = create_user_longterm_token(db) response.set_cookie( @@ -86,6 +97,22 @@ async def auto_login( domain=auth_settings.COOKIE_DOMAIN, ) + user = get_user_by_id(db, user_id) + + if user: + if user.store_api_key is None: + user.store_api_key = "" + + response.set_cookie( + "apikey_tkn_lflw", + str(user.store_api_key), # Ensure it's a string + httponly=auth_settings.ACCESS_HTTPONLY, + samesite=auth_settings.ACCESS_SAME_SITE, + secure=auth_settings.ACCESS_SECURE, + expires=None, # Set to None to make it a session cookie + domain=auth_settings.COOKIE_DOMAIN, + ) + return tokens raise HTTPException( @@ -140,4 +167,5 @@ async def refresh_token( async def logout(response: Response): response.delete_cookie("refresh_token_lf") response.delete_cookie("access_token_lf") + response.delete_cookie("apikey_tkn_lflw") return {"message": "Logout successful"} diff --git a/src/frontend/src/contexts/authContext.tsx b/src/frontend/src/contexts/authContext.tsx index b1d415aca..3521fe8b6 100644 --- a/src/frontend/src/contexts/authContext.tsx +++ b/src/frontend/src/contexts/authContext.tsx @@ -117,7 +117,6 @@ export function AuthProvider({ children }): React.ReactElement { } function storeApiKey(apikey: string) { - cookies.set(LANGFLOW_API_TOKEN, apikey, { path: "/" }); setApiKey(apikey); }