(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
2025-04-24 18:24:20 +00:00 · 2025-02-05 17:25:42 -08:00 · 2025-02-05 17:25:42 -08:00 · 6cef115bb0
commit 6cef115bb0
parent 88e7046165
3 changed files with 115 additions and 42 deletions
--- 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
--- 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"""
--- a/tests/proxy_security_tests/test_master_key_not_in_db.py
+++ 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."
+    )