From 6cef115bb07ae4adf665fbcac7de0a6ce8c76788 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 5 Feb 2025 17:25:42 -0800 Subject: [PATCH] (Security fix) - remove code block that inserts master key hash into DB (#8268) * remove code block upserting master key hash to db * run test to check if key upserted into db * run ci/cd again * litellm_proxy_security_tests * litellm_proxy_security_tests * run prisma entrypoint * ci/cd run again * fix test master key not in db --- .circleci/config.yml | 60 ++++++++++++++++++- litellm/proxy/proxy_server.py | 41 ------------- .../test_master_key_not_in_db.py | 56 +++++++++++++++++ 3 files changed, 115 insertions(+), 42 deletions(-) create mode 100644 tests/proxy_security_tests/test_master_key_not_in_db.py diff --git a/.circleci/config.yml b/.circleci/config.yml index c23aa6027c1..16a869e1a18 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -415,6 +415,56 @@ jobs: paths: - litellm_router_coverage.xml - litellm_router_coverage + litellm_proxy_security_tests: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + steps: + - checkout + - run: + name: Show git commit hash + command: | + echo "Git commit hash: $CIRCLE_SHA1" + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-asyncio==0.21.1" + pip install "pytest-cov==5.0.0" + - run: + name: Run prisma ./docker/entrypoint.sh + command: | + set +e + chmod +x docker/entrypoint.sh + ./docker/entrypoint.sh + set -e + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest tests/proxy_security_tests --cov=litellm --cov-report=xml -vv -x -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml litellm_proxy_security_tests_coverage.xml + mv .coverage litellm_proxy_security_tests_coverage + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - litellm_proxy_security_tests_coverage.xml + - litellm_proxy_security_tests_coverage litellm_proxy_unit_testing: # Runs all tests with the "proxy", "key", "jwt" filenames docker: - image: cimg/python:3.11 @@ -1788,7 +1838,7 @@ jobs: python -m venv venv . venv/bin/activate pip install coverage - coverage combine llm_translation_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage + coverage combine llm_translation_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage auth_ui_unit_tests_coverage langfuse_coverage caching_coverage litellm_proxy_unit_tests_coverage image_gen_coverage pass_through_unit_tests_coverage batches_coverage litellm_proxy_security_tests_coverage coverage xml - codecov/upload: file: ./coverage.xml @@ -2045,6 +2095,12 @@ workflows: only: - main - /litellm_.*/ + - litellm_proxy_security_tests: + filters: + branches: + only: + - main + - /litellm_.*/ - litellm_assistants_api_testing: filters: branches: @@ -2158,6 +2214,7 @@ workflows: - litellm_router_testing - caching_unit_tests - litellm_proxy_unit_testing + - litellm_proxy_security_tests - langfuse_logging_unit_tests - local_testing - litellm_assistants_api_testing @@ -2219,6 +2276,7 @@ workflows: - db_migration_disable_update_check - e2e_ui_testing - litellm_proxy_unit_testing + - litellm_proxy_security_tests - installing_litellm_on_python - installing_litellm_on_python_3_13 - proxy_logging_guardrails_model_info_tests diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1cee4bf11ab..893f011dd5d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -515,14 +515,6 @@ async def proxy_startup_event(app: FastAPI): prompt_injection_detection_obj.update_environment(router=llm_router) verbose_proxy_logger.debug("prisma_client: %s", prisma_client) - if prisma_client is not None and master_key is not None: - ProxyStartupEvent._add_master_key_hash_to_db( - master_key=master_key, - prisma_client=prisma_client, - litellm_proxy_admin_name=litellm_proxy_admin_name, - general_settings=general_settings, - ) - if prisma_client is not None and litellm.max_budget > 0: ProxyStartupEvent._add_proxy_budget_to_db( litellm_proxy_budget_name=litellm_proxy_admin_name @@ -3205,39 +3197,6 @@ class ProxyStartupEvent: litellm_jwtauth=litellm_jwtauth, ) - @classmethod - def _add_master_key_hash_to_db( - cls, - master_key: str, - prisma_client: PrismaClient, - litellm_proxy_admin_name: str, - general_settings: dict, - ): - """Adds master key hash to db for cost tracking""" - if os.getenv("PROXY_ADMIN_ID", None) is not None: - litellm_proxy_admin_name = os.getenv( - "PROXY_ADMIN_ID", litellm_proxy_admin_name - ) - if general_settings.get("disable_adding_master_key_hash_to_db") is True: - verbose_proxy_logger.info("Skipping writing master key hash to db") - else: - # add master key to db - # add 'admin' user to db. Fixes https://github.com/BerriAI/litellm/issues/6206 - task_1 = generate_key_helper_fn( - request_type="user", - duration=None, - models=[], - aliases={}, - config={}, - spend=0, - token=master_key, - user_id=litellm_proxy_admin_name, - user_role=LitellmUserRoles.PROXY_ADMIN, - query_type="update_data", - update_key_values={"user_role": LitellmUserRoles.PROXY_ADMIN}, - ) - asyncio.create_task(task_1) - @classmethod def _add_proxy_budget_to_db(cls, litellm_proxy_budget_name: str): """Adds a global proxy budget to db""" diff --git a/tests/proxy_security_tests/test_master_key_not_in_db.py b/tests/proxy_security_tests/test_master_key_not_in_db.py new file mode 100644 index 00000000000..e563b735a21 --- /dev/null +++ b/tests/proxy_security_tests/test_master_key_not_in_db.py @@ -0,0 +1,56 @@ +import os +import pytest +from fastapi.testclient import TestClient +from litellm.proxy.proxy_server import app, ProxyLogging +from litellm.caching import DualCache + +TEST_DB_ENV_VAR_NAME = "MASTER_KEY_CHECK_DB_URL" + + +@pytest.fixture(autouse=True) +def override_env_settings(monkeypatch): + # Set environment variables only for tests using-monkeypatch (function scope by default). + monkeypatch.setenv("DATABASE_URL", os.environ[TEST_DB_ENV_VAR_NAME]) + monkeypatch.setenv("LITELLM_MASTER_KEY", "sk-1234") + monkeypatch.setenv("LITELLM_LOG", "DEBUG") + + +@pytest.fixture(scope="module") +def test_client(): + """ + This fixture starts up the test client which triggers FastAPI's startup events. + Prisma will connect to the DB using the provided DATABASE_URL. + """ + with TestClient(app) as client: + yield client + + +@pytest.mark.asyncio +async def test_master_key_not_inserted(test_client): + """ + This test ensures that when the app starts (or when you hit the /health endpoint + to trigger startup logic), no unexpected write occurs in the DB. + """ + # Hit an endpoint (like /health) that triggers any startup tasks. + response = test_client.get("/health/liveliness") + assert response.status_code == 200 + + from litellm.proxy.utils import PrismaClient + + prisma_client = PrismaClient( + database_url=os.environ[TEST_DB_ENV_VAR_NAME], + proxy_logging_obj=ProxyLogging( + user_api_key_cache=DualCache(), premium_user=True + ), + ) + + # Connect directly to the test database to inspect the data. + await prisma_client.connect() + result = await prisma_client.db.litellm_verificationtoken.find_many() + print(result) + + # The expectation is that no token (or unintended record) is added on startup. + assert len(result) == 0, ( + "SECURITY ALERT SECURITY ALERT SECURITY ALERT: Expected no record in the litellm_verificationtoken table. On startup - the master key should NOT be Inserted into the DB." + "We have found keys in the DB. This is unexpected and should not happen." + )