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)