From a6c51ba3de332ce163b619bada5dcbcecdf38624 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 19 Sep 2026 17:53:17 -0700 Subject: [PATCH] fix(proxy): never treat plaintext that base64-decodes to nothing as a ciphertext during the master key migration A string such as "*" or "..." has no base64 characters, so it decoded to no bytes and read as an empty plaintext under any key. The migration would have counted it and overwritten it with a ciphertext of the empty string. Also read from the writer database instead of a read replica, report a database error during the migration instead of crashing the boot, skip columns the connected schema lacks across every schema on the search path, cap the JSON walk depth for the recursion detector, and move the boot wiring into one tested function. --- litellm/proxy/auth/master_key_boot_check.py | 6 +- .../common_utils/encrypt_decrypt_utils.py | 28 ++-- litellm/proxy/db/master_key_migration.py | 66 +++++++++- litellm/proxy/proxy_server.py | 20 ++- .../code_coverage_tests/recursive_detector.py | 1 + .../test_encrypt_decrypt_utils.py | 19 ++- .../proxy/db/test_master_key_migration.py | 122 +++++++++++++++++- 7 files changed, 226 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/auth/master_key_boot_check.py b/litellm/proxy/auth/master_key_boot_check.py index c8c2ea3539c..4b91cbc676d 100644 --- a/litellm/proxy/auth/master_key_boot_check.py +++ b/litellm/proxy/auth/master_key_boot_check.py @@ -285,7 +285,7 @@ def _migration_lead(migration: StoredSecretsMigration) -> str: else f"holds {migration.encrypted_value_count} value(s)" ) return ( - f"Your database {found} encrypted with this master key, which encrypts stored\n" - f"credentials while {SALT_KEY_ENV_VAR} is not set. Replacing the key alone makes them unreadable, so also tell\n" - "the proxy which key to migrate from:" + f"Your database {found} encrypted with this master key,\n" + f"which encrypts stored credentials while {SALT_KEY_ENV_VAR} is not set. Replacing the key alone makes them\n" + "unreadable, so also tell the proxy which key to migrate from:" ) diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index e655d51b31e..3584aaaf833 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -119,29 +119,31 @@ def encrypt_value_helper(value: str, new_encryption_key: str | None = None): raise e +def _legacy_ciphertext_bytes(value: str) -> bytes: + # Try URL-safe base64 decoding first (new format) + # Fall back to standard base64 decoding for backwards compatibility (old format) + try: + return base64.urlsafe_b64decode(value) + except Exception: + return base64.b64decode(value) + + def _decrypt_with_signing_key(value: str, signing_key: str) -> str: # Versioned AES-256-GCM values are detected before any base64 decode. # The prefix is the algorithm tag the legacy nacl format never carried. if value.startswith(_V2_GCM_PREFIX): return _decrypt_aes_gcm(value=value, signing_key=signing_key) - # Try URL-safe base64 decoding first (new format) - # Fall back to standard base64 decoding for backwards compatibility (old format) - try: - decoded_b64 = base64.urlsafe_b64decode(value) - except Exception: - # If URL-safe decoding fails, try standard base64 decoding for backwards compatibility - decoded_b64 = base64.b64decode(value) - - return decrypt_value(value=decoded_b64, signing_key=signing_key) + return decrypt_value(value=_legacy_ciphertext_bytes(value), signing_key=signing_key) def decrypt_if_encrypted_with(value: str, signing_key: str) -> str | None: - """None unless value is a ciphertext under signing_key. Both ciphers are authenticated, so a wrong key never passes.""" - if not value: - return None + """None unless value is a ciphertext under signing_key.""" try: - return _decrypt_with_signing_key(value=value, signing_key=signing_key) + # base64 decoding skips characters outside its alphabet, so "" and "*" decode to no bytes, + # which decrypt_value reads as an empty plaintext under any key. + decodes_to_nothing: Final = not value.startswith(_V2_GCM_PREFIX) and not _legacy_ciphertext_bytes(value) + return None if decodes_to_nothing else _decrypt_with_signing_key(value=value, signing_key=signing_key) except Exception: # noqa: BLE001 # base64, nacl and AES-GCM each raise their own "not a ciphertext" type return None diff --git a/litellm/proxy/db/master_key_migration.py b/litellm/proxy/db/master_key_migration.py index 452298516bd..78a00cf9ce9 100644 --- a/litellm/proxy/db/master_key_migration.py +++ b/litellm/proxy/db/master_key_migration.py @@ -7,6 +7,7 @@ from typing import Final from pydantic import JsonValue, TypeAdapter from typing_extensions import assert_never +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy.auth.master_key_boot_check import MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR from litellm.proxy.common_utils.callback_utils import CALLBACK_VAR_ENCRYPTED_PREFIX from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper @@ -50,17 +51,19 @@ _PRIMARY_KEY: Final = TypeAdapter(str) ReplaceCiphertext = Callable[[str], str | None] -def replace_ciphertexts(value: JsonValue, replacement_for: ReplaceCiphertext) -> tuple[JsonValue, int]: +def replace_ciphertexts(value: JsonValue, replacement_for: ReplaceCiphertext, depth: int = 0) -> tuple[JsonValue, int]: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return value, 0 match value: case str(): marker: Final = CALLBACK_VAR_ENCRYPTED_PREFIX if value.startswith(CALLBACK_VAR_ENCRYPTED_PREFIX) else "" replacement: Final = replacement_for(value.removeprefix(marker)) return (value, 0) if replacement is None else (marker + replacement, 1) case list(): - items: Final = tuple(replace_ciphertexts(item, replacement_for) for item in value) + items: Final = tuple(replace_ciphertexts(item, replacement_for, depth + 1) for item in value) return [item for item, _ in items], sum(count for _, count in items) case dict(): - fields: Final = {key: replace_ciphertexts(item, replacement_for) for key, item in value.items()} + fields: Final = {key: replace_ciphertexts(item, replacement_for, depth + 1) for key, item in value.items()} return {key: item for key, (item, _) in fields.items()}, sum(count for _, count in fields.values()) case _: return value, 0 @@ -109,7 +112,8 @@ async def _secret_columns_in(database: SupportsRawQueries) -> tuple[_SecretColum existing: Final = frozenset( (row["table_name"], row["column_name"]) for row in await database.query_raw( - "SELECT table_name, column_name FROM information_schema.columns WHERE table_schema = current_schema()" + "SELECT table_name, column_name FROM information_schema.columns " + "WHERE table_schema = ANY (current_schemas(false))" ) ) return tuple( @@ -168,7 +172,31 @@ class Migrated: remaining: int -MigrationOutcome = NothingToMigrate | Migrated +@dataclass(frozen=True, slots=True) +class MigrationFailed: + error: str + + +MigrationOutcome = NothingToMigrate | Migrated | MigrationFailed + + +async def migrate_if_requested( + *, + environ: Mapping[str, str], + master_key: str | None, + connected_database: Callable[[], SupportsRawQueries | None], + log: Callable[[str], None], +) -> MigrationOutcome | None: + previous_master_key: Final = environ.get(MIGRATE_FROM_MASTER_KEY_ENV_VAR) + if previous_master_key is None or master_key is None: + return None + return await migrate_from_previous_master_key( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=SALT_KEY_ENV_VAR in environ, + database=connected_database(), + log=log, + ) async def migrate_from_previous_master_key( @@ -179,7 +207,7 @@ async def migrate_from_previous_master_key( database: SupportsRawQueries | None, log: Callable[[str], None], ) -> MigrationOutcome: - outcome: Final = await _migrate( + outcome: Final = await _migrate_or_failure( previous_master_key=previous_master_key, master_key=master_key, salt_key_is_set=salt_key_is_set, @@ -190,6 +218,26 @@ async def migrate_from_previous_master_key( return outcome +async def _migrate_or_failure( + *, + previous_master_key: str, + master_key: str, + salt_key_is_set: bool, + database: SupportsRawQueries | None, + log: Callable[[str], None], +) -> MigrationOutcome: + try: + return await _migrate( + previous_master_key=previous_master_key, + master_key=master_key, + salt_key_is_set=salt_key_is_set, + database=database, + log=log, + ) + except Exception as error: # noqa: BLE001 # the proxy tolerates a database outage at boot, so the migration must too + return MigrationFailed(error=f"{type(error).__name__}: {error}"[:300]) + + async def _migrate( *, previous_master_key: str, @@ -244,5 +292,11 @@ def describe_outcome(outcome: MigrationOutcome) -> str: f"because they changed during the migration. Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart " "the proxy to migrate them." ) + case MigrationFailed(error=error): + return ( + f"Could not migrate stored values from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key ({error}). Values " + "still encrypted with the previous key cannot be read until the migration succeeds. Keep " + f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart the proxy once the database is reachable." + ) case _: assert_never(outcome) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5eff79193f1..0b15ea902ac 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -344,7 +344,6 @@ from litellm.proxy.auth.login_throttle import ( ) from litellm.proxy.auth.master_key_boot_check import ( MASTER_KEY_ENV_VAR, - MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR, UNSAFE_PROXY_OVERRIDE_ENV_VAR, announce_on_stderr_at_exit, @@ -491,7 +490,7 @@ from litellm.proxy.db.gateway_request_tracking import ( ) from litellm.proxy.db.master_key_migration import ( count_values_encrypted_with_or_none, - migrate_from_previous_master_key, + migrate_if_requested, ) from litellm.proxy.db.proxy_worker_heartbeat import ( PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS, @@ -1143,7 +1142,7 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries: database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj ) await client.connect() - return client.db + return client.writer_db @asynccontextmanager @@ -1263,15 +1262,12 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: user_api_key_cache=user_api_key_cache, ) - previous_master_key: Final = os.getenv(MIGRATE_FROM_MASTER_KEY_ENV_VAR) - if previous_master_key is not None and master_key is not None: - await migrate_from_previous_master_key( - previous_master_key=previous_master_key, - master_key=master_key, - salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None, - database=None if prisma_client is None else prisma_client.db, - log=verbose_proxy_logger.warning, - ) + await migrate_if_requested( + environ=os.environ, + master_key=master_key, + connected_database=lambda: None if prisma_client is None else prisma_client.writer_db, + log=verbose_proxy_logger.warning, + ) if prisma_client is not None: diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 3c6a6a58820..071191183df 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -55,6 +55,7 @@ IGNORE_FUNCTIONS = [ "apply_json_merge_patch", # max depth set (_MAX_MERGE_DEPTH=64); fails closed by raising ValueError at the cap. "_filter_argument_value", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by blocking the tool call at the cap. "_redact_scanned_content", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); fails closed by returning "[REDACTED]" at the cap. + "replace_ciphertexts", # max depth set (DEFAULT_MAX_RECURSE_DEPTH); walks stored JSON, which has no cycles, and leaves values below the cap untouched. "_iter_fallback_targets", # max depth set (2 * ROUTER_MAX_FALLBACKS); fails closed by raising ValueError at the cap. "_mergeable_branch", # max depth set (_MAX_SCHEMA_FLATTEN_DEPTH=32) plus a seen_refs cycle guard; passes the schema through untouched at the cap. "json_string_leaves", # max depth set (MAX_STRUCTURED_CONTENT_SCAN_DEPTH); fails closed by raising at the cap so nothing goes unscanned. diff --git a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py index fe9659f4fd7..9c07242bd23 100644 --- a/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_encrypt_decrypt_utils.py @@ -202,7 +202,24 @@ def test_explicit_key_decrypt_reads_only_values_written_under_that_key(monkeypat assert decrypt_value_helper(written_with_previous_key, key="t", exception_type="debug") is None -@pytest.mark.parametrize("not_a_ciphertext", ["", "gpt-5.4-mini", "https://example.invalid/v1", "v2:gcm:", "aGVsbG8="]) +@pytest.mark.parametrize( + "not_a_ciphertext", + [ + "", + "gpt-5.4-mini", + "https://example.invalid/v1", + "v2:gcm:", + "aGVsbG8=", + "*", + "-", + "_", + "...", + " ", + "{}", + "[]", + "=", + ], +) def test_explicit_key_decrypt_rejects_values_that_are_not_ciphertexts(not_a_ciphertext: str): assert decrypt_if_encrypted_with(not_a_ciphertext, "sk-1234") is None diff --git a/tests/test_litellm/proxy/db/test_master_key_migration.py b/tests/test_litellm/proxy/db/test_master_key_migration.py index 31732adc666..4cc9a652969 100644 --- a/tests/test_litellm/proxy/db/test_master_key_migration.py +++ b/tests/test_litellm/proxy/db/test_master_key_migration.py @@ -1,19 +1,24 @@ import json import re from collections.abc import Mapping, Sequence +from functools import reduce import pytest +from pydantic import JsonValue +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.proxy import proxy_server from litellm.proxy.auth.master_key_boot_check import MIGRATE_FROM_MASTER_KEY_ENV_VAR, SALT_KEY_ENV_VAR from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_if_encrypted_with, encrypt_value_helper from litellm.proxy.db.master_key_migration import ( _SECRET_COLUMNS, Migrated, + MigrationFailed, NothingToMigrate, count_values_encrypted_with, describe_outcome, migrate_from_previous_master_key, + migrate_if_requested, reencrypt_stored_values, replace_ciphertexts, ) @@ -43,7 +48,7 @@ class _FakeDatabase: async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: if "information_schema.columns" in query: - assert "table_schema = current_schema()" in query + assert "table_schema = ANY (current_schemas(false))" in query return [ {"table_name": secret_column.table, "column_name": secret_column.column} for secret_column in _SECRET_COLUMNS @@ -201,6 +206,26 @@ async def test_only_rows_holding_values_under_the_previous_key_are_written(): ] +@pytest.mark.asyncio +async def test_plaintext_that_base64_decodes_to_nothing_is_neither_counted_nor_rewritten(): + settings = {"allowed_routes": ["*"], "ui_name": "-", "separator": "...", "blank": " ", "shape": "{}"} + tables: Tables = { + "LiteLLM_Config": [{"param_name": "general_settings", "param_value": dict(settings)}], + "LiteLLM_VerificationToken": [ + {"token": "hashed", "metadata": {"notes": "...", "secret": "litellm_enc::" + _encrypted("callback-secret")}} + ], + } + database = _FakeDatabase(tables) + + found = await count_values_encrypted_with(database, PREVIOUS_KEY) + migrated = await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY) + + assert found == migrated == 1 + assert tables["LiteLLM_Config"][0]["param_value"] == settings + assert tables["LiteLLM_VerificationToken"][0]["metadata"]["notes"] == "..." + assert database.writes == [("LiteLLM_VerificationToken", "metadata", "hashed")] + + @pytest.mark.asyncio async def test_schema_without_some_of_the_tables_is_migrated_for_the_tables_it_has(): missing = frozenset({"LiteLLM_MCPUserCredentials", "LiteLLM_SSOIdentityAssertion"}) @@ -244,6 +269,22 @@ def test_replacing_ciphertexts_keeps_structure_markers_and_non_strings(): assert value["swap"] == ["old", {"nested": "litellm_enc::old"}] +def _nested(levels: int, leaf: str) -> JsonValue: + return reduce(lambda inner, _: [inner], range(levels), leaf) + + +@pytest.mark.parametrize("levels_past_the_cap, replaced_count", [(0, 1), (1, 0), (50, 0)]) +def test_walk_stops_at_the_recursion_cap_and_leaves_deeper_values_as_they_were( + levels_past_the_cap: int, replaced_count: int +): + value = _nested(DEFAULT_MAX_RECURSE_DEPTH + levels_past_the_cap, "old") + + replaced, count = replace_ciphertexts(value, lambda text: "new") + + assert count == replaced_count + assert replaced == _nested(DEFAULT_MAX_RECURSE_DEPTH + levels_past_the_cap, "new" if replaced_count else "old") + + async def _run( database: _FakeDatabase | _DatabaseThatMustNotBeTouched | None, *, @@ -396,3 +437,82 @@ async def test_encrypted_empty_string_is_migrated_like_any_other_value(): assert migrated == 1 assert decrypt_if_encrypted_with(str(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"]), NEW_KEY) == "" + + +@pytest.mark.asyncio +async def test_database_error_during_the_migration_is_reported_instead_of_crashing_the_boot(): + class _DatabaseIsDown(_FakeDatabase): + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: + raise ConnectionError("Can't reach database server") + + outcome, logged = await _run(_DatabaseIsDown(_seeded_tables())) + + assert outcome == MigrationFailed(error="ConnectionError: Can't reach database server") + assert len(logged) == 1 + assert "ConnectionError: Can't reach database server" in logged[0] + assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[0] + assert "ou may now delete" not in logged[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("previous_master_key", [PREVIOUS_KEY, ""]) +async def test_boot_migrates_from_the_environment_variable_to_the_running_master_key(previous_master_key: str): + tables: Tables = { + "LiteLLM_CredentialsTable": [ + { + "credential_id": "cred-1", + "credential_values": { + "api_key": _encrypted("sk-provider", previous_master_key) + if previous_master_key + else _encrypted_with_empty_key("sk-provider") + }, + } + ] + } + logged: list[str] = [] + + outcome = await migrate_if_requested( + environ={MIGRATE_FROM_MASTER_KEY_ENV_VAR: previous_master_key}, + master_key=NEW_KEY, + connected_database=lambda: _FakeDatabase(tables), + log=logged.append, + ) + + assert outcome == Migrated(migrated=1, remaining=0) + stored = tables["LiteLLM_CredentialsTable"][0]["credential_values"] + assert isinstance(stored, dict) + assert decrypt_if_encrypted_with(stored["api_key"], NEW_KEY) == "sk-provider" + assert "Done re-encrypting 1 stored value(s)" in logged[-1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "environ, master_key, outcome", + [ + ({}, NEW_KEY, None), + ({MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY}, None, None), + ( + {MIGRATE_FROM_MASTER_KEY_ENV_VAR: PREVIOUS_KEY, SALT_KEY_ENV_VAR: "a-salt-key"}, + NEW_KEY, + NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES, + ), + ], + ids=["variable-not-set", "no-master-key", "salt-key-set"], +) +async def test_boot_leaves_the_database_alone_unless_a_migration_was_requested_and_can_apply( + environ: dict[str, str], master_key: str | None, outcome: NothingToMigrate | None +): + logged: list[str] = [] + database_handles_taken: list[str] = [] + + def connected_database() -> _DatabaseThatMustNotBeTouched: + database_handles_taken.append("taken") + return _DatabaseThatMustNotBeTouched() + + result = await migrate_if_requested( + environ=environ, master_key=master_key, connected_database=connected_database, log=logged.append + ) + + assert result is outcome + assert len(database_handles_taken) == (0 if outcome is None else 1) + assert len(logged) == (0 if outcome is None else 1)