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.
This commit is contained in:
ryan-crabbe-berri 2026-09-19 17:53:17 -07:00
parent 38d776bd2b
commit a6c51ba3de
7 changed files with 226 additions and 36 deletions

View file

@ -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:"
)

View file

@ -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

View file

@ -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)

View file

@ -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:

View file

@ -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.

View file

@ -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

View file

@ -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)