mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
feat(proxy): re-encrypt stored secrets at boot from LITELLM_MIGRATE_FROM_MASTER_KEY so an unsafe key can be replaced while the proxy refuses to start
Rotating through POST /key/regenerate needs a running proxy, which a refused boot does not have. The refusal now counts the stored values that decrypt under the unsafe key. When there are none it only asks for a new key. When there are some it also asks for LITELLM_MIGRATE_FROM_MASTER_KEY, and the next boot with a safe key re-encrypts them and logs that the variable can be deleted. Leaving the variable set afterwards is a no-op with one notice.
This commit is contained in:
ryan-crabbe-berri
parent
fdd614d759
commit
38d776bd2b
9 changed files with 1032 additions and 72 deletions
|
|
@ -1,7 +1,7 @@
|
|||
import atexit
|
||||
import sys
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
from enum import Enum
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
|
@ -15,6 +15,7 @@ UNSAFE_PROXY_OVERRIDE_ENV_VAR: Final = "LITELLM_DANGEROUSLY_ALLOW_UNSAFE_PROXY"
|
|||
MASTER_KEY_SETTING: Final = "master_key"
|
||||
MASTER_KEY_ENV_VAR: Final = "LITELLM_MASTER_KEY"
|
||||
SALT_KEY_ENV_VAR: Final = "LITELLM_SALT_KEY"
|
||||
MIGRATE_FROM_MASTER_KEY_ENV_VAR: Final = "LITELLM_MIGRATE_FROM_MASTER_KEY"
|
||||
PUBLICLY_KNOWN_MASTER_KEYS: Final = frozenset({"sk-1234"})
|
||||
ROTATION_DOCS_URL: Final = "https://docs.litellm.ai/docs/proxy/master_key_rotations#proxy-refuses-to-start"
|
||||
_NEW_MASTER_KEY: Final = "sk-$(openssl rand -hex 32)"
|
||||
|
|
@ -51,12 +52,18 @@ class UnsafeMasterKeyAllowed:
|
|||
reason: UnsafeMasterKeyReason
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class StoredSecretsMigration:
|
||||
from_master_key: str
|
||||
encrypted_value_count: int | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UnsafeMasterKeyRefused:
|
||||
reason: UnsafeMasterKeyReason
|
||||
source: MasterKeySource
|
||||
environment_variable_is_set: bool
|
||||
stored_credentials_need_rotation: bool
|
||||
migration: StoredSecretsMigration | None
|
||||
|
||||
|
||||
MasterKeyBootVerdict = SafeMasterKey | UnsafeMasterKeyAllowed | UnsafeMasterKeyRefused
|
||||
|
|
@ -90,12 +97,23 @@ def master_key_boot_verdict(
|
|||
else EnvironmentSource()
|
||||
),
|
||||
environment_variable_is_set=environment_master_key is not None,
|
||||
stored_credentials_need_rotation=(
|
||||
reason is UnsafeMasterKeyReason.PUBLICLY_KNOWN and not salt_key_is_set and database_is_configured
|
||||
migration=(
|
||||
StoredSecretsMigration(from_master_key=master_key, encrypted_value_count=None)
|
||||
if master_key is not None and not salt_key_is_set and database_is_configured
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def with_stored_secrets_counted(
|
||||
verdict: MasterKeyBootVerdict, count_values_encrypted_with: Callable[[str], Awaitable[int | None]]
|
||||
) -> MasterKeyBootVerdict:
|
||||
if not isinstance(verdict, UnsafeMasterKeyRefused) or verdict.migration is None:
|
||||
return verdict
|
||||
count: Final = await count_values_encrypted_with(verdict.migration.from_master_key)
|
||||
return replace(verdict, migration=None if count == 0 else replace(verdict.migration, encrypted_value_count=count))
|
||||
|
||||
|
||||
def enforce_master_key_boot_verdict(verdict: MasterKeyBootVerdict, announce: Callable[[str], object]) -> None:
|
||||
match verdict:
|
||||
case SafeMasterKey():
|
||||
|
|
@ -129,7 +147,7 @@ def render_refusal(refusal: UnsafeMasterKeyRefused) -> str:
|
|||
return "\n\n".join(
|
||||
(
|
||||
f"LiteLLM proxy refused to start: {_REFUSAL_HEADLINE[refusal.reason]}\n{_source_line(refusal)}",
|
||||
_ROTATE_INSTEAD_OF_REPLACING if refusal.stored_credentials_need_rotation else _fix_steps(refusal),
|
||||
_fix_steps(refusal),
|
||||
_OVERRIDE_HINT,
|
||||
)
|
||||
)
|
||||
|
|
@ -168,12 +186,9 @@ _REPLACE_EXPORTED_KEY_STEP: Final = (
|
|||
" already exported in the environment wins over .env."
|
||||
)
|
||||
|
||||
_ROTATE_INSTEAD_OF_REPLACING: Final = (
|
||||
f"Credentials stored in your database are encrypted with this master key because {SALT_KEY_ENV_VAR} is not\n"
|
||||
"set, so replacing the key makes them undecryptable. Rotate it by following this guide, which re-encrypts them:\n"
|
||||
f" {ROTATION_DOCS_URL}\n"
|
||||
"Generate the new key for it with (save it only once the guide says to):\n"
|
||||
f" {PRINT_NEW_MASTER_KEY_COMMAND}"
|
||||
_RESTART_TO_MIGRATE_STEP: Final = (
|
||||
"Start the proxy again. It re-encrypts the stored values with the new key, then logs that\n"
|
||||
f" {MIGRATE_FROM_MASTER_KEY_ENV_VAR} can be removed. Details: {ROTATION_DOCS_URL}"
|
||||
)
|
||||
|
||||
_OVERRIDE_HINT: Final = (
|
||||
|
|
@ -218,15 +233,59 @@ def _source_line(refusal: UnsafeMasterKeyRefused) -> str:
|
|||
|
||||
|
||||
def _fix_steps(refusal: UnsafeMasterKeyRefused) -> str:
|
||||
set_key_step: Final = _REPLACE_EXPORTED_KEY_STEP if refusal.environment_variable_is_set else _SAVE_KEY_STEP
|
||||
match refusal.source:
|
||||
case ConfigFileSource() as source:
|
||||
steps: Final = (*_config_steps(refusal.source), *_key_steps(refusal))
|
||||
numbered: Final = "\n".join(f"{number}. {step}" for number, step in enumerate(steps, start=1))
|
||||
return numbered if refusal.migration is None else f"{_migration_lead(refusal.migration)}\n{numbered}"
|
||||
|
||||
|
||||
def _config_steps(source: MasterKeySource) -> tuple[str, ...]:
|
||||
match source:
|
||||
case ConfigFileSource():
|
||||
return (
|
||||
f"1. Make sure {_config_label(source)} reads the key from the environment:\n"
|
||||
f" general_settings:\n {MASTER_KEY_SETTING}: os.environ/{MASTER_KEY_ENV_VAR}\n"
|
||||
f"2. {set_key_step}"
|
||||
f"Make sure {_config_label(source)} reads the key from the environment:\n"
|
||||
f" general_settings:\n {MASTER_KEY_SETTING}: os.environ/{MASTER_KEY_ENV_VAR}",
|
||||
)
|
||||
case EnvironmentSource():
|
||||
return f"1. {set_key_step}"
|
||||
return ()
|
||||
case _:
|
||||
assert_never(refusal.source)
|
||||
assert_never(source)
|
||||
|
||||
|
||||
def _key_steps(refusal: UnsafeMasterKeyRefused) -> tuple[str, ...]:
|
||||
if refusal.migration is None:
|
||||
return (_REPLACE_EXPORTED_KEY_STEP if refusal.environment_variable_is_set else _SAVE_KEY_STEP,)
|
||||
if refusal.environment_variable_is_set:
|
||||
return (
|
||||
f"Set the key to migrate from next to {MASTER_KEY_ENV_VAR}, wherever that is set (a shell export, your\n"
|
||||
" container or deployment environment, or .env):\n"
|
||||
f" {_migrate_from_assignment(refusal.migration)}",
|
||||
_REPLACE_EXPORTED_KEY_STEP,
|
||||
_RESTART_TO_MIGRATE_STEP,
|
||||
)
|
||||
return (
|
||||
"Save the key to migrate from and a newly generated key to .env:\n"
|
||||
f" echo '{_migrate_from_assignment(refusal.migration)}' | tee -a .env\n"
|
||||
f" {GENERATE_MASTER_KEY_COMMAND}\n"
|
||||
" Not using a .env file (docker run, Kubernetes, pip install)? Pass the same two values as\n"
|
||||
" environment variables instead.",
|
||||
_RESTART_TO_MIGRATE_STEP,
|
||||
)
|
||||
|
||||
|
||||
def _migrate_from_assignment(migration: StoredSecretsMigration) -> str:
|
||||
key: Final = migration.from_master_key
|
||||
value: Final = key if key == key.strip() else f'"{key}"'
|
||||
return f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}={value}"
|
||||
|
||||
|
||||
def _migration_lead(migration: StoredSecretsMigration) -> str:
|
||||
found: Final = (
|
||||
"could not be checked for values"
|
||||
if migration.encrypted_value_count is None
|
||||
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:"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -44,7 +44,8 @@ _EXTRA_SENSITIVE_CALLBACK_KEYS: Final = {"gcs_path_service_account"}
|
|||
# Sentinel prefix on encrypted callback_var values. Lets us detect
|
||||
# already-encrypted input cheaply (no decrypt-attempt round trip) and
|
||||
# avoid double-encrypting if `LITELLM_SALT_KEY` is rotated between writes.
|
||||
_CALLBACK_VAR_ENCRYPTED_PREFIX: Final = "litellm_enc::"
|
||||
CALLBACK_VAR_ENCRYPTED_PREFIX: Final = "litellm_enc::"
|
||||
_CALLBACK_VAR_ENCRYPTED_PREFIX: Final = CALLBACK_VAR_ENCRYPTED_PREFIX
|
||||
# Metadata slots that hold operator-configured callback and secret-manager setup
|
||||
# (and therefore integration credentials). Resolved from UserAPIKeyAuth during
|
||||
# pre-call setup, never read back off the copies stamped into request metadata.
|
||||
|
|
|
|||
|
|
@ -119,6 +119,33 @@ def encrypt_value_helper(value: str, new_encryption_key: str | None = None):
|
|||
raise e
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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
|
||||
try:
|
||||
return _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
|
||||
|
||||
|
||||
def decrypt_value_helper(
|
||||
value: str,
|
||||
key: str, # this is just for debug purposes, showing the k,v pair that's invalid. not a signing key.
|
||||
|
|
@ -129,21 +156,7 @@ def decrypt_value_helper(
|
|||
|
||||
try:
|
||||
if isinstance(value, 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=cast(str, 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)
|
||||
|
||||
value = decrypt_value(value=decoded_b64, signing_key=signing_key)
|
||||
return value
|
||||
return _decrypt_with_signing_key(value=value, signing_key=cast(str, signing_key))
|
||||
|
||||
# if it's not str - do not decrypt it, return the value
|
||||
return value
|
||||
|
|
|
|||
248
litellm/proxy/db/master_key_migration.py
Normal file
248
litellm/proxy/db/master_key_migration.py
Normal file
|
|
@ -0,0 +1,248 @@
|
|||
import json
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
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
|
||||
from litellm.proxy.db.create_views import SupportsRawQueries
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SecretColumn:
|
||||
table: str
|
||||
primary_key: str
|
||||
column: str
|
||||
is_json: bool = True
|
||||
only_rows_with_marked_ciphertexts: bool = False
|
||||
|
||||
|
||||
_SECRET_COLUMNS: Final = (
|
||||
_SecretColumn("LiteLLM_ProxyModelTable", "model_id", "litellm_params"),
|
||||
_SecretColumn("LiteLLM_CredentialsTable", "credential_id", "credential_values"),
|
||||
_SecretColumn("LiteLLM_Config", "param_name", "param_value"),
|
||||
_SecretColumn("LiteLLM_SSOConfig", "id", "sso_settings"),
|
||||
_SecretColumn("LiteLLM_CacheConfig", "id", "cache_settings"),
|
||||
_SecretColumn("LiteLLM_ConfigOverrides", "config_type", "config_value"),
|
||||
_SecretColumn("LiteLLM_MCPServerTable", "server_id", "credentials"),
|
||||
_SecretColumn("LiteLLM_MCPServerTable", "server_id", "static_headers"),
|
||||
_SecretColumn("LiteLLM_MCPServerTable", "server_id", "env_vars"),
|
||||
_SecretColumn("LiteLLM_MCPServerTable", "server_id", "env"),
|
||||
_SecretColumn("LiteLLM_MCPServerOAuthClient", "server_id", "credentials"),
|
||||
_SecretColumn("LiteLLM_MCPUserCredentials", "id", "credential_b64", is_json=False),
|
||||
_SecretColumn("LiteLLM_MCPUserEnvVars", "id", "values_b64", is_json=False),
|
||||
_SecretColumn("LiteLLM_SSOIdentityAssertion", "user_id", "assertion_b64", is_json=False),
|
||||
_SecretColumn("LiteLLM_TeamTable", "team_id", "metadata", only_rows_with_marked_ciphertexts=True),
|
||||
_SecretColumn("LiteLLM_VerificationToken", "token", "metadata", only_rows_with_marked_ciphertexts=True),
|
||||
_SecretColumn("LiteLLM_UserTable", "user_id", "metadata", only_rows_with_marked_ciphertexts=True),
|
||||
_SecretColumn("LiteLLM_DeletedTeamTable", "id", "metadata", only_rows_with_marked_ciphertexts=True),
|
||||
_SecretColumn("LiteLLM_DeletedVerificationToken", "id", "metadata", only_rows_with_marked_ciphertexts=True),
|
||||
)
|
||||
|
||||
_STORED_VALUE: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_PRIMARY_KEY: Final = TypeAdapter(str)
|
||||
|
||||
ReplaceCiphertext = Callable[[str], str | None]
|
||||
|
||||
|
||||
def replace_ciphertexts(value: JsonValue, replacement_for: ReplaceCiphertext) -> tuple[JsonValue, int]:
|
||||
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)
|
||||
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()}
|
||||
return {key: item for key, (item, _) in fields.items()}, sum(count for _, count in fields.values())
|
||||
case _:
|
||||
return value, 0
|
||||
|
||||
|
||||
async def count_values_encrypted_with(database: SupportsRawQueries, signing_key: str) -> int:
|
||||
def keep(value: str) -> str | None:
|
||||
return None if decrypt_if_encrypted_with(value, signing_key) is None else value
|
||||
|
||||
return sum(
|
||||
[
|
||||
replace_ciphertexts(_STORED_VALUE.validate_python(row[secret_column.column]), keep)[1]
|
||||
for secret_column in await _secret_columns_in(database)
|
||||
for row in await _rows_of(database, secret_column)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
async def count_values_encrypted_with_or_none(
|
||||
connect: Callable[[], Awaitable[SupportsRawQueries]], signing_key: str
|
||||
) -> int | None:
|
||||
try:
|
||||
return await count_values_encrypted_with(await connect(), signing_key)
|
||||
except Exception: # noqa: BLE001 # an unreadable database must not replace the boot refusal with a traceback
|
||||
return None
|
||||
|
||||
|
||||
async def reencrypt_stored_values(database: SupportsRawQueries, *, from_key: str, to_key: str) -> int:
|
||||
def reencrypted(value: str) -> str | None:
|
||||
plaintext: Final = decrypt_if_encrypted_with(value, from_key)
|
||||
return None if plaintext is None else _CIPHERTEXT.validate_python(encrypt_value_helper(plaintext, to_key))
|
||||
|
||||
return sum(
|
||||
[
|
||||
await _reencrypt_row(database, secret_column, row, reencrypted)
|
||||
for secret_column in await _secret_columns_in(database)
|
||||
for row in await _rows_of(database, secret_column)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
_CIPHERTEXT: Final = TypeAdapter(str)
|
||||
|
||||
|
||||
async def _secret_columns_in(database: SupportsRawQueries) -> tuple[_SecretColumn, ...]:
|
||||
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()"
|
||||
)
|
||||
)
|
||||
return tuple(
|
||||
secret_column for secret_column in _SECRET_COLUMNS if (secret_column.table, secret_column.column) in existing
|
||||
)
|
||||
|
||||
|
||||
async def _rows_of(database: SupportsRawQueries, secret_column: _SecretColumn) -> tuple[Mapping[str, object], ...]:
|
||||
marked_only: Final = (
|
||||
f" AND \"{secret_column.column}\"::text LIKE '%{CALLBACK_VAR_ENCRYPTED_PREFIX}%'"
|
||||
if secret_column.only_rows_with_marked_ciphertexts
|
||||
else ""
|
||||
)
|
||||
return tuple(
|
||||
await database.query_raw(
|
||||
f'SELECT "{secret_column.primary_key}", "{secret_column.column}" FROM "{secret_column.table}" '
|
||||
f'WHERE "{secret_column.column}" IS NOT NULL{marked_only}'
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _reencrypt_row(
|
||||
database: SupportsRawQueries,
|
||||
secret_column: _SecretColumn,
|
||||
row: Mapping[str, object],
|
||||
reencrypted: ReplaceCiphertext,
|
||||
) -> int:
|
||||
stored: Final = _STORED_VALUE.validate_python(row[secret_column.column])
|
||||
migrated, count = replace_ciphertexts(stored, reencrypted)
|
||||
if count == 0:
|
||||
return 0
|
||||
cast_to: Final = "::jsonb" if secret_column.is_json else ""
|
||||
rows_updated: Final = await database.execute_raw(
|
||||
f'UPDATE "{secret_column.table}" SET "{secret_column.column}" = $1{cast_to} '
|
||||
f'WHERE "{secret_column.primary_key}" = $2 AND "{secret_column.column}" = $3{cast_to}',
|
||||
_as_sql_parameter(migrated, secret_column),
|
||||
_PRIMARY_KEY.validate_python(row[secret_column.primary_key]),
|
||||
_as_sql_parameter(stored, secret_column),
|
||||
)
|
||||
return count if rows_updated else 0
|
||||
|
||||
|
||||
def _as_sql_parameter(value: JsonValue, secret_column: _SecretColumn) -> str:
|
||||
return json.dumps(value) if secret_column.is_json else _CIPHERTEXT.validate_python(value)
|
||||
|
||||
|
||||
class NothingToMigrate(Enum):
|
||||
SALT_KEY_ENCRYPTS_STORED_VALUES = "salt_key_encrypts_stored_values"
|
||||
NO_DATABASE = "no_database"
|
||||
NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY = "nothing_encrypted_with_previous_key"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Migrated:
|
||||
migrated: int
|
||||
remaining: int
|
||||
|
||||
|
||||
MigrationOutcome = NothingToMigrate | Migrated
|
||||
|
||||
|
||||
async def migrate_from_previous_master_key(
|
||||
*,
|
||||
previous_master_key: str,
|
||||
master_key: str,
|
||||
salt_key_is_set: bool,
|
||||
database: SupportsRawQueries | None,
|
||||
log: Callable[[str], None],
|
||||
) -> MigrationOutcome:
|
||||
outcome: Final = await _migrate(
|
||||
previous_master_key=previous_master_key,
|
||||
master_key=master_key,
|
||||
salt_key_is_set=salt_key_is_set,
|
||||
database=database,
|
||||
log=log,
|
||||
)
|
||||
log(describe_outcome(outcome))
|
||||
return outcome
|
||||
|
||||
|
||||
async def _migrate(
|
||||
*,
|
||||
previous_master_key: str,
|
||||
master_key: str,
|
||||
salt_key_is_set: bool,
|
||||
database: SupportsRawQueries | None,
|
||||
log: Callable[[str], None],
|
||||
) -> MigrationOutcome:
|
||||
if salt_key_is_set:
|
||||
return NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES
|
||||
if database is None:
|
||||
return NothingToMigrate.NO_DATABASE
|
||||
if previous_master_key == master_key:
|
||||
return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY
|
||||
found: Final = await count_values_encrypted_with(database, previous_master_key)
|
||||
if found == 0:
|
||||
return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY
|
||||
log(f"Re-encrypting {found} stored value(s) from the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} key to the new master key.")
|
||||
migrated: Final = Migrated(
|
||||
migrated=await reencrypt_stored_values(database, from_key=previous_master_key, to_key=master_key),
|
||||
remaining=await count_values_encrypted_with(database, previous_master_key),
|
||||
)
|
||||
another_worker_migrated_everything: Final = migrated == Migrated(migrated=0, remaining=0)
|
||||
return NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY if another_worker_migrated_everything else migrated
|
||||
|
||||
|
||||
def describe_outcome(outcome: MigrationOutcome) -> str:
|
||||
match outcome:
|
||||
case NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES:
|
||||
return (
|
||||
f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is set, but {SALT_KEY_ENV_VAR} is what encrypts your stored "
|
||||
f"values, so there is nothing to migrate. You may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}."
|
||||
)
|
||||
case NothingToMigrate.NO_DATABASE:
|
||||
return (
|
||||
f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is set, but no database is connected, so nothing was migrated. If "
|
||||
f"this proxy has no database, you may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}."
|
||||
)
|
||||
case NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY:
|
||||
return (
|
||||
f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} is still set, but nothing in the database is left to migrate "
|
||||
f"from that key. You may now delete {MIGRATE_FROM_MASTER_KEY_ENV_VAR}."
|
||||
)
|
||||
case Migrated(migrated=migrated, remaining=0):
|
||||
return (
|
||||
f"Done re-encrypting {migrated} stored value(s) with the new master key. You may now delete the "
|
||||
f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR} environment variable."
|
||||
)
|
||||
case Migrated(migrated=migrated, remaining=remaining):
|
||||
return (
|
||||
f"Re-encrypted {migrated} stored value(s), but {remaining} are still encrypted with the previous key "
|
||||
f"because they changed during the migration. Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set and restart "
|
||||
"the proxy to migrate them."
|
||||
)
|
||||
case _:
|
||||
assert_never(outcome)
|
||||
|
|
@ -344,11 +344,13 @@ 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,
|
||||
enforce_master_key_boot_verdict,
|
||||
master_key_boot_verdict,
|
||||
with_stored_secrets_counted,
|
||||
)
|
||||
from litellm.proxy.auth.model_checks import (
|
||||
expand_wildcard_deployments_for_model_info,
|
||||
|
|
@ -472,6 +474,7 @@ from litellm.proxy.config_resolvers.settings_rules import (
|
|||
)
|
||||
from litellm.proxy.container_endpoints.endpoints import router as container_router
|
||||
from litellm.proxy.credential_endpoints.endpoints import router as credential_router
|
||||
from litellm.proxy.db.create_views import SupportsRawQueries
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
|
||||
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
|
||||
|
|
@ -486,6 +489,10 @@ from litellm.proxy.db.gateway_request_tracking import (
|
|||
GatewayRequestRedisBuffer,
|
||||
flush_gateway_requests,
|
||||
)
|
||||
from litellm.proxy.db.master_key_migration import (
|
||||
count_values_encrypted_with_or_none,
|
||||
migrate_from_previous_master_key,
|
||||
)
|
||||
from litellm.proxy.db.proxy_worker_heartbeat import (
|
||||
PROXY_WORKER_HEARTBEAT_INTERVAL_SECONDS,
|
||||
ProxyWorkerHeartbeat,
|
||||
|
|
@ -1131,6 +1138,14 @@ async def _initialize_shared_aiohttp_session():
|
|||
return None
|
||||
|
||||
|
||||
async def _connect_to_count_stored_values() -> SupportsRawQueries:
|
||||
client: Final = prisma_client or PrismaClient(
|
||||
database_url=str(get_secret("DATABASE_URL")), proxy_logging_obj=proxy_logging_obj
|
||||
)
|
||||
await client.connect()
|
||||
return client.db
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
global \
|
||||
|
|
@ -1224,14 +1239,17 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]:
|
|||
await initialize(**worker_config)
|
||||
|
||||
enforce_master_key_boot_verdict(
|
||||
master_key_boot_verdict(
|
||||
master_key=master_key,
|
||||
environment_master_key=os.getenv(MASTER_KEY_ENV_VAR),
|
||||
general_settings=general_settings,
|
||||
config_file_path=user_config_file_path,
|
||||
override_env_is_on=get_secret_bool(UNSAFE_PROXY_OVERRIDE_ENV_VAR) is True,
|
||||
salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None,
|
||||
database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None,
|
||||
await with_stored_secrets_counted(
|
||||
master_key_boot_verdict(
|
||||
master_key=master_key,
|
||||
environment_master_key=os.getenv(MASTER_KEY_ENV_VAR),
|
||||
general_settings=general_settings,
|
||||
config_file_path=user_config_file_path,
|
||||
override_env_is_on=get_secret_bool(UNSAFE_PROXY_OVERRIDE_ENV_VAR) is True,
|
||||
salt_key_is_set=os.getenv(SALT_KEY_ENV_VAR) is not None,
|
||||
database_is_configured=prisma_client is not None or get_secret("DATABASE_URL", None) is not None,
|
||||
),
|
||||
count_values_encrypted_with=partial(count_values_encrypted_with_or_none, _connect_to_count_stored_values),
|
||||
),
|
||||
announce=announce_on_stderr_at_exit,
|
||||
)
|
||||
|
|
@ -1245,6 +1263,16 @@ 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,
|
||||
)
|
||||
|
||||
if prisma_client is not None:
|
||||
|
||||
async def _run_pw_migration():
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
|
|
@ -10,6 +11,7 @@ import pytest
|
|||
from litellm.proxy.auth.master_key_boot_check import (
|
||||
GENERATE_MASTER_KEY_COMMAND,
|
||||
MASTER_KEY_ENV_VAR,
|
||||
MIGRATE_FROM_MASTER_KEY_ENV_VAR,
|
||||
PRINT_NEW_MASTER_KEY_COMMAND,
|
||||
ROTATION_DOCS_URL,
|
||||
UNSAFE_PROXY_OVERRIDE_ENV_VAR,
|
||||
|
|
@ -18,6 +20,7 @@ from litellm.proxy.auth.master_key_boot_check import (
|
|||
EnvironmentSource,
|
||||
MasterKeyBootVerdict,
|
||||
SafeMasterKey,
|
||||
StoredSecretsMigration,
|
||||
UnsafeMasterKeyAllowed,
|
||||
UnsafeMasterKeyError,
|
||||
UnsafeMasterKeyReason,
|
||||
|
|
@ -26,6 +29,7 @@ from litellm.proxy.auth.master_key_boot_check import (
|
|||
enforce_master_key_boot_verdict,
|
||||
master_key_boot_verdict,
|
||||
render_refusal,
|
||||
with_stored_secrets_counted,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -125,38 +129,86 @@ def test_environment_is_the_source_when_yaml_does_not_set_a_master_key():
|
|||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("master_key", "salt_key_is_set", "database_is_configured", "needs_rotation"),
|
||||
("master_key", "salt_key_is_set", "database_is_configured", "migration"),
|
||||
[
|
||||
("sk-1234", False, True, True),
|
||||
("sk-1234", True, True, False),
|
||||
("sk-1234", False, False, False),
|
||||
(None, False, True, False),
|
||||
("", False, True, False),
|
||||
("sk-1234", False, True, StoredSecretsMigration(from_master_key="sk-1234", encrypted_value_count=None)),
|
||||
("", False, True, StoredSecretsMigration(from_master_key="", encrypted_value_count=None)),
|
||||
(" sk-1234\n", False, True, StoredSecretsMigration(from_master_key=" sk-1234\n", encrypted_value_count=None)),
|
||||
("sk-1234", True, True, None),
|
||||
("sk-1234", False, False, None),
|
||||
(None, False, True, None),
|
||||
],
|
||||
)
|
||||
def test_rotation_is_only_needed_when_the_known_key_encrypts_a_database(
|
||||
master_key: str | None, salt_key_is_set: bool, database_is_configured: bool, needs_rotation: bool
|
||||
def test_migration_is_offered_from_the_exact_key_that_may_encrypt_a_database(
|
||||
master_key: str | None,
|
||||
salt_key_is_set: bool,
|
||||
database_is_configured: bool,
|
||||
migration: StoredSecretsMigration | None,
|
||||
):
|
||||
verdict = _verdict(master_key, salt_key_is_set=salt_key_is_set, database_is_configured=database_is_configured)
|
||||
|
||||
assert isinstance(verdict, UnsafeMasterKeyRefused)
|
||||
assert verdict.stored_credentials_need_rotation is needs_rotation
|
||||
assert verdict.migration == migration
|
||||
|
||||
|
||||
def _counted(verdict: MasterKeyBootVerdict, count: int | None) -> tuple[MasterKeyBootVerdict, list[str]]:
|
||||
asked_about: list[str] = []
|
||||
|
||||
async def count_values_encrypted_with(signing_key: str) -> int | None:
|
||||
asked_about.append(signing_key)
|
||||
return count
|
||||
|
||||
return asyncio.run(with_stored_secrets_counted(verdict, count_values_encrypted_with)), asked_about
|
||||
|
||||
|
||||
def test_database_with_nothing_encrypted_needs_no_migration():
|
||||
counted, asked_about = _counted(_verdict("sk-1234", database_is_configured=True), 0)
|
||||
|
||||
assert isinstance(counted, UnsafeMasterKeyRefused)
|
||||
assert counted.migration is None
|
||||
assert asked_about == ["sk-1234"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("count", [4, None])
|
||||
def test_database_with_encrypted_values_or_unreadable_keeps_the_migration(count: int | None):
|
||||
counted, _ = _counted(_verdict("", database_is_configured=True), count)
|
||||
|
||||
assert isinstance(counted, UnsafeMasterKeyRefused)
|
||||
assert counted.migration == StoredSecretsMigration(from_master_key="", encrypted_value_count=count)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"verdict",
|
||||
[
|
||||
SafeMasterKey(),
|
||||
UnsafeMasterKeyAllowed(reason=UnsafeMasterKeyReason.PUBLICLY_KNOWN),
|
||||
_verdict("sk-1234", database_is_configured=False),
|
||||
],
|
||||
)
|
||||
def test_database_is_not_read_when_no_migration_is_on_the_table(verdict: MasterKeyBootVerdict):
|
||||
counted, asked_about = _counted(verdict, 7)
|
||||
|
||||
assert counted == verdict
|
||||
assert asked_about == []
|
||||
|
||||
|
||||
def _refusal(
|
||||
reason: UnsafeMasterKeyReason = UnsafeMasterKeyReason.PUBLICLY_KNOWN,
|
||||
source: ConfigFileSource | EnvironmentSource = EnvironmentSource(),
|
||||
environment_variable_is_set: bool = False,
|
||||
stored_credentials_need_rotation: bool = False,
|
||||
migration: StoredSecretsMigration | None = None,
|
||||
) -> UnsafeMasterKeyRefused:
|
||||
return UnsafeMasterKeyRefused(
|
||||
reason=reason,
|
||||
source=source,
|
||||
environment_variable_is_set=environment_variable_is_set,
|
||||
stored_credentials_need_rotation=stored_credentials_need_rotation,
|
||||
migration=migration,
|
||||
)
|
||||
|
||||
|
||||
_MIGRATION = StoredSecretsMigration(from_master_key="sk-1234", encrypted_value_count=3)
|
||||
|
||||
|
||||
def test_config_refusal_names_the_file_and_tells_it_to_read_the_environment():
|
||||
text = render_refusal(_refusal(source=ConfigFileSource(config_file_path="/app/config.yaml")))
|
||||
|
||||
|
|
@ -208,28 +260,82 @@ def test_unset_key_refusal_says_nothing_supplied_one():
|
|||
assert "Neither general_settings.master_key nor" in text
|
||||
|
||||
|
||||
def test_rotation_guide_appears_only_when_needed():
|
||||
with_rotation = render_refusal(_refusal(stored_credentials_need_rotation=True))
|
||||
without_rotation = render_refusal(_refusal(stored_credentials_need_rotation=False))
|
||||
def test_migration_steps_appear_only_when_the_database_needs_them():
|
||||
with_migration = render_refusal(_refusal(migration=_MIGRATION))
|
||||
without_migration = render_refusal(_refusal(migration=None))
|
||||
|
||||
assert ROTATION_DOCS_URL in with_rotation
|
||||
assert ROTATION_DOCS_URL not in without_rotation
|
||||
assert f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=sk-1234" in with_migration
|
||||
assert "holds 3 value(s) encrypted with this master key" in with_migration
|
||||
assert ROTATION_DOCS_URL in with_migration
|
||||
assert MIGRATE_FROM_MASTER_KEY_ENV_VAR not in without_migration
|
||||
assert ROTATION_DOCS_URL not in without_migration
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", [EnvironmentSource(), ConfigFileSource(config_file_path="/app/config.yaml")])
|
||||
def test_refusal_never_tells_a_user_who_must_rotate_to_save_the_new_key_first(
|
||||
source: ConfigFileSource | EnvironmentSource,
|
||||
):
|
||||
text = render_refusal(_refusal(source=source, stored_credentials_need_rotation=True))
|
||||
def test_unreadable_database_is_reported_as_unchecked_rather_than_counted():
|
||||
text = render_refusal(
|
||||
_refusal(migration=StoredSecretsMigration(from_master_key="sk-1234", encrypted_value_count=None))
|
||||
)
|
||||
|
||||
assert "could not be checked" in text
|
||||
assert "value(s)" not in text
|
||||
assert f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=sk-1234" in text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("from_master_key", "assignment"),
|
||||
[
|
||||
("sk-1234", f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=sk-1234"),
|
||||
("", f"{MIGRATE_FROM_MASTER_KEY_ENV_VAR}="),
|
||||
(" sk-1234", f'{MIGRATE_FROM_MASTER_KEY_ENV_VAR}=" sk-1234"'),
|
||||
],
|
||||
)
|
||||
def test_migrate_from_assignment_carries_the_exact_previous_key(from_master_key: str, assignment: str):
|
||||
text = render_refusal(
|
||||
_refusal(
|
||||
environment_variable_is_set=True,
|
||||
migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1),
|
||||
)
|
||||
)
|
||||
|
||||
assert f" {assignment}\n" in text
|
||||
|
||||
|
||||
def test_migration_with_an_exported_key_replaces_it_in_place_and_numbers_every_step():
|
||||
text = render_refusal(
|
||||
_refusal(
|
||||
source=ConfigFileSource(config_file_path="/app/config.yaml"),
|
||||
environment_variable_is_set=True,
|
||||
migration=_MIGRATION,
|
||||
)
|
||||
)
|
||||
|
||||
assert "tee" not in text
|
||||
assert PRINT_NEW_MASTER_KEY_COMMAND in text
|
||||
assert ".env" not in text
|
||||
assert "os.environ/" not in text
|
||||
assert [line[:2] for line in text.splitlines() if re.match(r"\d\. ", line)] == ["1.", "2.", "3.", "4."]
|
||||
assert text.index("os.environ/") < text.index(MIGRATE_FROM_MASTER_KEY_ENV_VAR + "=") < text.index("Start the proxy")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stored_credentials_need_rotation", [True, False])
|
||||
def test_override_hint_is_the_last_paragraph(stored_credentials_need_rotation: bool):
|
||||
text = render_refusal(_refusal(stored_credentials_need_rotation=stored_credentials_need_rotation))
|
||||
@pytest.mark.skipif(shutil.which("openssl") is None, reason="the printed command shells out to openssl")
|
||||
@pytest.mark.parametrize("from_master_key", ["sk-1234", "", " sk-1234"])
|
||||
def test_printed_migration_commands_save_both_keys_to_the_env_file(tmp_path: Path, from_master_key: str):
|
||||
from dotenv import dotenv_values
|
||||
|
||||
text = render_refusal(
|
||||
_refusal(migration=StoredSecretsMigration(from_master_key=from_master_key, encrypted_value_count=1))
|
||||
)
|
||||
commands = [line.strip() for line in text.splitlines() if line.strip().startswith("echo ")]
|
||||
|
||||
subprocess.run(["bash", "-c", "\n".join(commands)], cwd=tmp_path, capture_output=True, text=True, check=True)
|
||||
|
||||
saved = dotenv_values(tmp_path / ".env")
|
||||
assert saved[MIGRATE_FROM_MASTER_KEY_ENV_VAR] == from_master_key
|
||||
assert _verdict(saved[MASTER_KEY_ENV_VAR]) == SafeMasterKey()
|
||||
assert len(commands) == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("migration", [_MIGRATION, None])
|
||||
def test_override_hint_is_the_last_paragraph(migration: StoredSecretsMigration | None):
|
||||
text = render_refusal(_refusal(migration=migration))
|
||||
last_paragraph = text.split("\n\n")[-1]
|
||||
|
||||
assert UNSAFE_PROXY_OVERRIDE_ENV_VAR in last_paragraph
|
||||
|
|
|
|||
|
|
@ -6,12 +6,16 @@ gate, and the backward-compatibility guarantees that let legacy XSalsa20-Poly130
|
|||
(nacl) ciphertext and new AES values coexist and decrypt correctly.
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
||||
_V2_GCM_PREFIX,
|
||||
decrypt_if_encrypted_with,
|
||||
decrypt_value_helper,
|
||||
encrypt_value,
|
||||
encrypt_value_helper,
|
||||
)
|
||||
|
||||
|
|
@ -185,3 +189,33 @@ def test_decrypt_failure_debug_log_omits_raw_value(monkeypatch):
|
|||
"the failing key should still be named in the breadcrumb"
|
||||
)
|
||||
assert result == secret
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_aes", [False, True])
|
||||
def test_explicit_key_decrypt_reads_only_values_written_under_that_key(monkeypatch, use_aes: bool):
|
||||
if use_aes:
|
||||
_use_aes(monkeypatch)
|
||||
written_with_previous_key = encrypt_value_helper("stored-secret", new_encryption_key="sk-1234")
|
||||
|
||||
assert decrypt_if_encrypted_with(written_with_previous_key, "sk-1234") == "stored-secret"
|
||||
assert decrypt_if_encrypted_with(written_with_previous_key, "sk-another-key") is None
|
||||
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="])
|
||||
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
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_aes", [False, True])
|
||||
def test_explicit_key_decrypt_tells_an_encrypted_empty_string_from_no_ciphertext(monkeypatch, use_aes: bool):
|
||||
if use_aes:
|
||||
_use_aes(monkeypatch)
|
||||
|
||||
assert decrypt_if_encrypted_with(encrypt_value_helper("", new_encryption_key="sk-1234"), "sk-1234") == ""
|
||||
|
||||
|
||||
def test_explicit_key_decrypt_supports_the_empty_master_key():
|
||||
written_with_empty_key = encrypt_value(value="stored-secret", signing_key="")
|
||||
|
||||
assert decrypt_if_encrypted_with(base64.urlsafe_b64encode(written_with_empty_key).decode(), "") == "stored-secret"
|
||||
|
|
|
|||
398
tests/test_litellm/proxy/db/test_master_key_migration.py
Normal file
398
tests/test_litellm/proxy/db/test_master_key_migration.py
Normal file
|
|
@ -0,0 +1,398 @@
|
|||
import json
|
||||
import re
|
||||
from collections.abc import Mapping, Sequence
|
||||
|
||||
import pytest
|
||||
|
||||
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,
|
||||
NothingToMigrate,
|
||||
count_values_encrypted_with,
|
||||
describe_outcome,
|
||||
migrate_from_previous_master_key,
|
||||
reencrypt_stored_values,
|
||||
replace_ciphertexts,
|
||||
)
|
||||
|
||||
PREVIOUS_KEY = "sk-1234"
|
||||
NEW_KEY = "sk-qa-9f2c1e7a44b0d3"
|
||||
UNRELATED_KEY = "sk-some-other-deployment"
|
||||
|
||||
Tables = dict[str, list[dict[str, object]]]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _legacy_algorithm_and_no_salt_key(monkeypatch: pytest.MonkeyPatch):
|
||||
monkeypatch.delenv(SALT_KEY_ENV_VAR, raising=False)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {})
|
||||
|
||||
|
||||
def _encrypted(plaintext: str, key: str = PREVIOUS_KEY) -> str:
|
||||
return str(encrypt_value_helper(plaintext, new_encryption_key=key))
|
||||
|
||||
|
||||
class _FakeDatabase:
|
||||
def __init__(self, tables: Tables, tables_missing_from_the_schema: frozenset[str] = frozenset()) -> None:
|
||||
self.tables = tables
|
||||
self.tables_missing_from_the_schema = tables_missing_from_the_schema
|
||||
self.writes: list[tuple[str, str, str]] = []
|
||||
|
||||
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
|
||||
return [
|
||||
{"table_name": secret_column.table, "column_name": secret_column.column}
|
||||
for secret_column in _SECRET_COLUMNS
|
||||
if secret_column.table not in self.tables_missing_from_the_schema
|
||||
]
|
||||
select = re.fullmatch(r'SELECT "(\w+)", "(\w+)" FROM "(\w+)" WHERE "\2" IS NOT NULL(.*)', query)
|
||||
assert select is not None, query
|
||||
primary_key, column, table, row_filter = select.groups()
|
||||
assert row_filter in ("", f" AND \"{column}\"::text LIKE '%litellm_enc::%'")
|
||||
assert table not in self.tables_missing_from_the_schema, f'relation "{table}" does not exist'
|
||||
return [
|
||||
{primary_key: row[primary_key], column: json.loads(json.dumps(row[column]))}
|
||||
for row in self.tables.get(table, [])
|
||||
if row.get(column) is not None and (not row_filter or "litellm_enc::" in json.dumps(row[column]))
|
||||
]
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
update = re.fullmatch(r'UPDATE "(\w+)" SET "(\w+)" = \$1(::jsonb|) WHERE "(\w+)" = \$2 AND "\2" = \$3\3', query)
|
||||
assert update is not None, query
|
||||
table, column, json_cast, primary_key = update.groups()
|
||||
new_value, row_id, expected = (
|
||||
json.loads(str(arg)) if json_cast and index != 1 else arg for index, arg in enumerate(args)
|
||||
)
|
||||
matching = [row for row in self.tables[table] if row[primary_key] == row_id and row[column] == expected]
|
||||
for row in matching:
|
||||
row[column] = new_value
|
||||
self.writes.append((table, column, str(row_id)))
|
||||
return len(matching)
|
||||
|
||||
|
||||
class _DatabaseThatMustNotBeTouched:
|
||||
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
|
||||
raise AssertionError(f"unexpected read: {query}")
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
raise AssertionError(f"unexpected write: {query}")
|
||||
|
||||
|
||||
def _seeded_tables() -> Tables:
|
||||
return {
|
||||
"LiteLLM_ProxyModelTable": [
|
||||
{
|
||||
"model_id": "model-1",
|
||||
"litellm_params": {
|
||||
"api_key": _encrypted("provider-key"),
|
||||
"model": _encrypted("openai/gpt-5.4-mini"),
|
||||
"rpm": 10,
|
||||
"use_in_pass_through": False,
|
||||
"api_base": None,
|
||||
},
|
||||
},
|
||||
{"model_id": "model-2", "litellm_params": {"api_key": _encrypted("other-deployment", UNRELATED_KEY)}},
|
||||
],
|
||||
"LiteLLM_Config": [
|
||||
{"param_name": "environment_variables", "param_value": {"LANGFUSE_SECRET_KEY": _encrypted("env-secret")}},
|
||||
{"param_name": "general_settings", "param_value": {"proxy_batch_write_at": 10, "ui_name": "plain text"}},
|
||||
{"param_name": "cleared", "param_value": None},
|
||||
],
|
||||
"LiteLLM_MCPServerTable": [
|
||||
{
|
||||
"server_id": "mcp-1",
|
||||
"credentials": {"auth_value": _encrypted("mcp-token"), "aws_region_name": "us-east-1"},
|
||||
"static_headers": _encrypted('{"X-Api-Key": "header-secret"}'),
|
||||
"env_vars": [
|
||||
{"name": "GLOBAL", "scope": "global", "value": _encrypted("global-env")},
|
||||
{"name": "PER_USER", "scope": "user", "value": ""},
|
||||
],
|
||||
"env": {},
|
||||
}
|
||||
],
|
||||
"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("byok-secret")}],
|
||||
"LiteLLM_TeamTable": [
|
||||
{
|
||||
"team_id": "team-1",
|
||||
"metadata": {
|
||||
"logging": [
|
||||
{
|
||||
"callback_name": "langfuse",
|
||||
"callback_vars": {
|
||||
"langfuse_host": "https://example.invalid",
|
||||
"langfuse_secret_key": "litellm_enc::" + _encrypted("team-callback-secret"),
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
},
|
||||
{"team_id": "team-without-callbacks", "metadata": {"note": _encrypted("unmarked, so never selected")}},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
_VALUES_UNDER_THE_PREVIOUS_KEY = 8
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reencryption_moves_every_stored_shape_to_the_new_key_and_nothing_else():
|
||||
tables = _seeded_tables()
|
||||
untouched_before = json.dumps(
|
||||
[tables["LiteLLM_ProxyModelTable"][1], tables["LiteLLM_Config"][1:], tables["LiteLLM_TeamTable"][1]]
|
||||
)
|
||||
|
||||
migrated = await reencrypt_stored_values(_FakeDatabase(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY)
|
||||
|
||||
assert migrated == _VALUES_UNDER_THE_PREVIOUS_KEY
|
||||
model_params = tables["LiteLLM_ProxyModelTable"][0]["litellm_params"]
|
||||
assert isinstance(model_params, dict)
|
||||
assert decrypt_if_encrypted_with(model_params["api_key"], NEW_KEY) == "provider-key"
|
||||
assert decrypt_if_encrypted_with(model_params["model"], NEW_KEY) == "openai/gpt-5.4-mini"
|
||||
assert decrypt_if_encrypted_with(model_params["api_key"], PREVIOUS_KEY) is None
|
||||
assert (model_params["rpm"], model_params["use_in_pass_through"], model_params["api_base"]) == (10, False, None)
|
||||
mcp_server = tables["LiteLLM_MCPServerTable"][0]
|
||||
assert decrypt_if_encrypted_with(mcp_server["static_headers"], NEW_KEY) == '{"X-Api-Key": "header-secret"}'
|
||||
assert mcp_server["credentials"]["aws_region_name"] == "us-east-1"
|
||||
assert decrypt_if_encrypted_with(mcp_server["env_vars"][0]["value"], NEW_KEY) == "global-env"
|
||||
assert mcp_server["env_vars"][1] == {"name": "PER_USER", "scope": "user", "value": ""}
|
||||
assert (
|
||||
decrypt_if_encrypted_with(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"], NEW_KEY) == "byok-secret"
|
||||
)
|
||||
callback_secret = tables["LiteLLM_TeamTable"][0]["metadata"]["logging"][0]["callback_vars"]["langfuse_secret_key"]
|
||||
assert callback_secret.startswith("litellm_enc::")
|
||||
assert decrypt_if_encrypted_with(callback_secret.removeprefix("litellm_enc::"), NEW_KEY) == "team-callback-secret"
|
||||
assert untouched_before == json.dumps(
|
||||
[tables["LiteLLM_ProxyModelTable"][1], tables["LiteLLM_Config"][1:], tables["LiteLLM_TeamTable"][1]]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_count_follows_the_values_from_the_previous_key_to_the_new_one():
|
||||
database = _FakeDatabase(_seeded_tables())
|
||||
|
||||
assert await count_values_encrypted_with(database, PREVIOUS_KEY) == _VALUES_UNDER_THE_PREVIOUS_KEY
|
||||
assert await count_values_encrypted_with(database, NEW_KEY) == 0
|
||||
|
||||
await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY)
|
||||
|
||||
assert await count_values_encrypted_with(database, PREVIOUS_KEY) == 0
|
||||
assert await count_values_encrypted_with(database, NEW_KEY) == _VALUES_UNDER_THE_PREVIOUS_KEY
|
||||
assert await count_values_encrypted_with(database, UNRELATED_KEY) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_rows_holding_values_under_the_previous_key_are_written():
|
||||
database = _FakeDatabase(_seeded_tables())
|
||||
|
||||
await reencrypt_stored_values(database, from_key=PREVIOUS_KEY, to_key=NEW_KEY)
|
||||
|
||||
assert sorted(database.writes) == [
|
||||
("LiteLLM_Config", "param_value", "environment_variables"),
|
||||
("LiteLLM_MCPServerTable", "credentials", "mcp-1"),
|
||||
("LiteLLM_MCPServerTable", "env_vars", "mcp-1"),
|
||||
("LiteLLM_MCPServerTable", "static_headers", "mcp-1"),
|
||||
("LiteLLM_MCPUserCredentials", "credential_b64", "cred-row-1"),
|
||||
("LiteLLM_ProxyModelTable", "litellm_params", "model-1"),
|
||||
("LiteLLM_TeamTable", "metadata", "team-1"),
|
||||
]
|
||||
|
||||
|
||||
@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"})
|
||||
tables = _seeded_tables()
|
||||
database = _FakeDatabase(tables, tables_missing_from_the_schema=missing)
|
||||
|
||||
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 == _VALUES_UNDER_THE_PREVIOUS_KEY - 1
|
||||
assert decrypt_if_encrypted_with(str(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"]), PREVIOUS_KEY)
|
||||
|
||||
|
||||
class _SomeoneEditsEachRowAfterItIsRead(_FakeDatabase):
|
||||
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
|
||||
rows = await super().query_raw(query, *args)
|
||||
for row in self.tables.get("LiteLLM_MCPUserCredentials", []):
|
||||
row["credential_b64"] = "edited-by-an-admin"
|
||||
return rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_value_edited_while_the_migration_runs_is_not_overwritten():
|
||||
tables: Tables = {"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("byok-secret")}]}
|
||||
|
||||
migrated = await reencrypt_stored_values(
|
||||
_SomeoneEditsEachRowAfterItIsRead(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY
|
||||
)
|
||||
|
||||
assert migrated == 0
|
||||
assert tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"] == "edited-by-an-admin"
|
||||
|
||||
|
||||
def test_replacing_ciphertexts_keeps_structure_markers_and_non_strings():
|
||||
value = {"keep": [1, True, None, "plain"], "swap": ["old", {"nested": "litellm_enc::old"}]}
|
||||
|
||||
replaced, count = replace_ciphertexts(value, lambda text: "new" if text == "old" else None)
|
||||
|
||||
assert replaced == {"keep": [1, True, None, "plain"], "swap": ["new", {"nested": "litellm_enc::new"}]}
|
||||
assert count == 2
|
||||
assert value["swap"] == ["old", {"nested": "litellm_enc::old"}]
|
||||
|
||||
|
||||
async def _run(
|
||||
database: _FakeDatabase | _DatabaseThatMustNotBeTouched | None,
|
||||
*,
|
||||
previous_master_key: str = PREVIOUS_KEY,
|
||||
master_key: str = NEW_KEY,
|
||||
salt_key_is_set: bool = False,
|
||||
) -> tuple[object, list[str]]:
|
||||
logged: list[str] = []
|
||||
outcome = await migrate_from_previous_master_key(
|
||||
previous_master_key=previous_master_key,
|
||||
master_key=master_key,
|
||||
salt_key_is_set=salt_key_is_set,
|
||||
database=database,
|
||||
log=logged.append,
|
||||
)
|
||||
return outcome, logged
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_migration_announces_itself_then_says_the_variable_can_go():
|
||||
database = _FakeDatabase(_seeded_tables())
|
||||
|
||||
outcome, logged = await _run(database)
|
||||
|
||||
assert outcome == Migrated(migrated=_VALUES_UNDER_THE_PREVIOUS_KEY, remaining=0)
|
||||
assert len(logged) == 2
|
||||
assert logged[0].startswith(f"Re-encrypting {_VALUES_UNDER_THE_PREVIOUS_KEY} stored value(s)")
|
||||
assert logged[1].startswith(f"Done re-encrypting {_VALUES_UNDER_THE_PREVIOUS_KEY} stored value(s)")
|
||||
assert f"You may now delete the {MIGRATE_FROM_MASTER_KEY_ENV_VAR} environment variable" in logged[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_variable_left_set_after_the_migration_is_a_no_op_with_one_notice():
|
||||
database = _FakeDatabase(_seeded_tables())
|
||||
await _run(database)
|
||||
writes_after_the_migration = list(database.writes)
|
||||
stored_after_the_migration = json.dumps(database.tables)
|
||||
|
||||
outcome, logged = await _run(database)
|
||||
|
||||
assert outcome is NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY
|
||||
assert logged == [describe_outcome(NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY)]
|
||||
assert database.writes == writes_after_the_migration
|
||||
assert json.dumps(database.tables) == stored_after_the_migration
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_previous_master_key_migrates_like_any_other():
|
||||
tables: Tables = {"LiteLLM_CredentialsTable": [{"credential_id": "c1", "credential_values": {"api_key": "x"}}]}
|
||||
tables["LiteLLM_CredentialsTable"][0]["credential_values"] = {"api_key": _encrypted_with_empty_key("cred-secret")}
|
||||
|
||||
outcome, _ = await _run(_FakeDatabase(tables), previous_master_key="")
|
||||
|
||||
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) == "cred-secret"
|
||||
|
||||
|
||||
def _encrypted_with_empty_key(plaintext: str) -> str:
|
||||
import base64
|
||||
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value
|
||||
|
||||
return base64.urlsafe_b64encode(encrypt_value(value=plaintext, signing_key="")).decode()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("salt_key_is_set", "has_database", "master_key", "expected"),
|
||||
[
|
||||
(True, True, NEW_KEY, NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES),
|
||||
(False, False, NEW_KEY, NothingToMigrate.NO_DATABASE),
|
||||
(False, True, PREVIOUS_KEY, NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY),
|
||||
],
|
||||
)
|
||||
async def test_database_is_left_alone_when_the_master_key_cannot_have_encrypted_it(
|
||||
salt_key_is_set: bool, has_database: bool, master_key: str, expected: NothingToMigrate
|
||||
):
|
||||
outcome, logged = await _run(
|
||||
_DatabaseThatMustNotBeTouched() if has_database else None,
|
||||
master_key=master_key,
|
||||
salt_key_is_set=salt_key_is_set,
|
||||
)
|
||||
|
||||
assert outcome is expected
|
||||
assert logged == [describe_outcome(expected)]
|
||||
|
||||
|
||||
class _AnotherWorkerMigratesRightAfterTheFirstCount(_FakeDatabase):
|
||||
def __init__(self, tables: Tables) -> None:
|
||||
super().__init__(tables)
|
||||
self.reads = 0
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]:
|
||||
rows = await super().query_raw(query, *args)
|
||||
self.reads += 1
|
||||
if self.reads == len(_SECRET_COLUMNS):
|
||||
await reencrypt_stored_values(_FakeDatabase(self.tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY)
|
||||
return rows
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_worker_that_loses_the_race_reports_nothing_left_instead_of_zero_values_done():
|
||||
database = _AnotherWorkerMigratesRightAfterTheFirstCount(_seeded_tables())
|
||||
|
||||
outcome, logged = await _run(database)
|
||||
|
||||
assert outcome is NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY
|
||||
assert database.writes == []
|
||||
assert logged[-1] == describe_outcome(NothingToMigrate.NOTHING_ENCRYPTED_WITH_PREVIOUS_KEY)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_values_that_could_not_be_written_keep_the_variable_in_place():
|
||||
tables: Tables = {"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("byok-secret")}]}
|
||||
|
||||
class _EveryWriteLosesToACurrentEdit(_FakeDatabase):
|
||||
async def execute_raw(self, query: str, *args: object) -> int:
|
||||
return 0
|
||||
|
||||
outcome, logged = await _run(_EveryWriteLosesToACurrentEdit(tables))
|
||||
|
||||
assert outcome == Migrated(migrated=0, remaining=1)
|
||||
assert "1 are still encrypted with the previous key" in logged[-1]
|
||||
assert f"Keep {MIGRATE_FROM_MASTER_KEY_ENV_VAR} set" in logged[-1]
|
||||
assert "You may now delete" not in logged[-1]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"outcome", [*NothingToMigrate, Migrated(migrated=5, remaining=0)], ids=lambda outcome: str(outcome)
|
||||
)
|
||||
def test_every_finished_outcome_tells_the_user_the_variable_can_be_deleted(outcome: NothingToMigrate | Migrated):
|
||||
message = describe_outcome(outcome)
|
||||
|
||||
assert "ou may now delete" in message
|
||||
assert MIGRATE_FROM_MASTER_KEY_ENV_VAR in message
|
||||
|
||||
|
||||
def test_salt_key_outcome_names_the_salt_key_as_the_reason():
|
||||
assert SALT_KEY_ENV_VAR in describe_outcome(NothingToMigrate.SALT_KEY_ENCRYPTS_STORED_VALUES)
|
||||
assert SALT_KEY_ENV_VAR not in describe_outcome(NothingToMigrate.NO_DATABASE)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_encrypted_empty_string_is_migrated_like_any_other_value():
|
||||
tables: Tables = {"LiteLLM_MCPUserCredentials": [{"id": "cred-row-1", "credential_b64": _encrypted("")}]}
|
||||
|
||||
migrated = await reencrypt_stored_values(_FakeDatabase(tables), from_key=PREVIOUS_KEY, to_key=NEW_KEY)
|
||||
|
||||
assert migrated == 1
|
||||
assert decrypt_if_encrypted_with(str(tables["LiteLLM_MCPUserCredentials"][0]["credential_b64"]), NEW_KEY) == ""
|
||||
|
|
@ -1633,7 +1633,14 @@ def _boot_with_general_settings(monkeypatch, tmp_path, general_settings):
|
|||
|
||||
config_path = tmp_path / "config.yaml"
|
||||
config_path.write_text(yaml.dump({"general_settings": general_settings}))
|
||||
for name in ("LITELLM_MASTER_KEY", "LITELLM_DANGEROUSLY_ALLOW_UNSAFE_PROXY", "WORKER_CONFIG", "DATABASE_URL"):
|
||||
for name in (
|
||||
"LITELLM_MASTER_KEY",
|
||||
"LITELLM_DANGEROUSLY_ALLOW_UNSAFE_PROXY",
|
||||
"LITELLM_MIGRATE_FROM_MASTER_KEY",
|
||||
"LITELLM_SALT_KEY",
|
||||
"WORKER_CONFIG",
|
||||
"DATABASE_URL",
|
||||
):
|
||||
monkeypatch.delenv(name, raising=False)
|
||||
monkeypatch.setenv("CONFIG_FILE_PATH", str(config_path))
|
||||
scheduler_left_on_a_closed_event_loop_by_an_earlier_test = "litellm.proxy.proxy_server.scheduler"
|
||||
|
|
@ -1649,7 +1656,7 @@ def _boot_with_general_settings(monkeypatch, tmp_path, general_settings):
|
|||
[{"master_key": "sk-1234"}, {"master_key": ""}, {"master_key": None}, {}],
|
||||
ids=["publicly-known", "empty", "yaml-null", "no-general-settings"],
|
||||
)
|
||||
async def test_proxy_startup_refuses_an_unsafe_master_key_before_connecting_to_the_database(
|
||||
async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database_is_unreachable(
|
||||
monkeypatch, tmp_path, general_settings
|
||||
):
|
||||
from fastapi import FastAPI
|
||||
|
|
@ -1657,8 +1664,12 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_before_connecting_to_t
|
|||
from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError
|
||||
from litellm.proxy.proxy_server import proxy_startup_event
|
||||
|
||||
async def unreachable():
|
||||
raise ConnectionError("database is down")
|
||||
|
||||
_, announced = _boot_with_general_settings(monkeypatch, tmp_path, general_settings)
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server._connect_to_count_stored_values", unreachable)
|
||||
|
||||
with pytest.raises(UnsafeMasterKeyError):
|
||||
async with proxy_startup_event(FastAPI()):
|
||||
|
|
@ -1666,6 +1677,68 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_before_connecting_to_t
|
|||
|
||||
assert len(announced) == 1
|
||||
assert "sk-$(openssl rand -hex 32)" in announced[0]
|
||||
key_can_have_encrypted_the_database = general_settings.get("master_key") is not None
|
||||
assert ("could not be checked" in announced[0]) == key_can_have_encrypted_the_database
|
||||
|
||||
|
||||
class _DatabaseWithOneStoredCredential:
|
||||
def __init__(self, ciphertext):
|
||||
self._ciphertext = ciphertext
|
||||
|
||||
async def query_raw(self, query, *args):
|
||||
if "information_schema.columns" in query:
|
||||
return [{"table_name": "LiteLLM_CredentialsTable", "column_name": "credential_values"}]
|
||||
return [{"credential_id": "cred-1", "credential_values": {"api_key": self._ciphertext}}]
|
||||
|
||||
async def execute_raw(self, query, *args):
|
||||
raise AssertionError("a refused boot must not write to the database")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("encrypted_with, asks_to_migrate", [("sk-1234", True), ("sk-some-other-key", False)])
|
||||
async def test_proxy_startup_asks_to_migrate_only_when_the_database_holds_values_under_the_unsafe_key(
|
||||
monkeypatch, tmp_path, encrypted_with, asks_to_migrate
|
||||
):
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy.auth.master_key_boot_check import UnsafeMasterKeyError
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.proxy_server import proxy_startup_event
|
||||
|
||||
database = _DatabaseWithOneStoredCredential(encrypt_value_helper("sk-provider", new_encryption_key=encrypted_with))
|
||||
|
||||
async def connected():
|
||||
return database
|
||||
|
||||
_, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-1234"})
|
||||
monkeypatch.setenv("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable")
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server._connect_to_count_stored_values", connected)
|
||||
|
||||
with pytest.raises(UnsafeMasterKeyError):
|
||||
async with proxy_startup_event(FastAPI()):
|
||||
pass
|
||||
|
||||
assert ("LITELLM_MIGRATE_FROM_MASTER_KEY=sk-1234" in announced[0]) == asks_to_migrate
|
||||
assert ("holds 1 value(s) encrypted with this master key" in announced[0]) == asks_to_migrate
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_startup_says_a_lingering_migrate_from_variable_can_be_deleted(monkeypatch, tmp_path, caplog):
|
||||
from fastapi import FastAPI
|
||||
|
||||
from litellm.proxy.proxy_server import proxy_startup_event
|
||||
|
||||
_boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"})
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
|
||||
monkeypatch.setenv("LITELLM_MIGRATE_FROM_MASTER_KEY", "")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
async with proxy_startup_event(FastAPI()):
|
||||
pass
|
||||
|
||||
notices = [record.getMessage() for record in caplog.records if "LITELLM_MIGRATE_FROM_MASTER_KEY" in record.message]
|
||||
assert len(notices) == 1
|
||||
assert "you may now delete LITELLM_MIGRATE_FROM_MASTER_KEY" in notices[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue