diff --git a/litellm/proxy/auth/master_key_boot_check.py b/litellm/proxy/auth/master_key_boot_check.py index def80f551e0..68cbf2ea6e1 100644 --- a/litellm/proxy/auth/master_key_boot_check.py +++ b/litellm/proxy/auth/master_key_boot_check.py @@ -1,5 +1,3 @@ -"""Decides at boot whether the proxy may start with the master key it resolved.""" - import atexit import sys from collections.abc import Callable, Mapping @@ -57,6 +55,7 @@ class UnsafeMasterKeyAllowed: class UnsafeMasterKeyRefused: reason: UnsafeMasterKeyReason source: MasterKeySource + environment_variable_is_set: bool stored_credentials_need_rotation: bool @@ -90,6 +89,7 @@ def master_key_boot_verdict( if MASTER_KEY_SETTING in general_settings and not config_file_only_relays_the_environment 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 ), @@ -116,11 +116,7 @@ def enforce_master_key_boot_verdict(verdict: MasterKeyBootVerdict, announce: Cal def announce_on_stderr_at_exit(message: str) -> None: - """Keeps the fix as the last thing on screen and away from the log handlers. - - A failed lifespan prints a traceback hundreds of lines long (a frame pair per included router) that buries - anything written before it, and the log redactor strips the key-shaped command from anything sent to a logger. - """ + """A logger would redact the key-shaped command and the lifespan traceback would bury it, so print at exit.""" atexit.register(_flush_stdout_then_write_stderr, message) @@ -133,7 +129,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.source), + _ROTATE_INSTEAD_OF_REPLACING if refusal.stored_credentials_need_rotation else _fix_steps(refusal), _OVERRIDE_HINT, ) ) @@ -164,6 +160,14 @@ _SAVE_KEY_STEP: Final = ( f" {MASTER_KEY_ENV_VAR} environment variable instead." ) +_REPLACE_EXPORTED_KEY_STEP: Final = ( + "Generate a key:\n" + f" {PRINT_NEW_MASTER_KEY_COMMAND}\n" + f" Put it in place of the current {MASTER_KEY_ENV_VAR} value wherever that is set: a shell export, your\n" + " container or deployment environment, or its line in .env. Do not just add it to .env, because a value\n" + " 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" @@ -213,15 +217,16 @@ def _source_line(refusal: UnsafeMasterKeyRefused) -> str: assert_never(refusal.source) -def _fix_steps(source: MasterKeySource) -> str: - match source: - case ConfigFileSource(): +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: 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. {_SAVE_KEY_STEP}" + f"2. {set_key_step}" ) case EnvironmentSource(): - return f"1. {_SAVE_KEY_STEP}" + return f"1. {set_key_step}" case _: - assert_never(source) + assert_never(refusal.source) diff --git a/tests/test_litellm/proxy/auth/test_master_key_boot_check.py b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py index fcd209d3d39..2bb33276a0c 100644 --- a/tests/test_litellm/proxy/auth/test_master_key_boot_check.py +++ b/tests/test_litellm/proxy/auth/test_master_key_boot_check.py @@ -1,6 +1,7 @@ import re import shutil import subprocess +import sys from collections.abc import Mapping from pathlib import Path @@ -21,6 +22,7 @@ from litellm.proxy.auth.master_key_boot_check import ( UnsafeMasterKeyError, UnsafeMasterKeyReason, UnsafeMasterKeyRefused, + announce_on_stderr_at_exit, enforce_master_key_boot_verdict, master_key_boot_verdict, render_refusal, @@ -144,10 +146,14 @@ def test_rotation_is_only_needed_when_the_known_key_encrypts_a_database( def _refusal( reason: UnsafeMasterKeyReason = UnsafeMasterKeyReason.PUBLICLY_KNOWN, source: ConfigFileSource | EnvironmentSource = EnvironmentSource(), + environment_variable_is_set: bool = False, stored_credentials_need_rotation: bool = False, ) -> UnsafeMasterKeyRefused: return UnsafeMasterKeyRefused( - reason=reason, source=source, stored_credentials_need_rotation=stored_credentials_need_rotation + reason=reason, + source=source, + environment_variable_is_set=environment_variable_is_set, + stored_credentials_need_rotation=stored_credentials_need_rotation, ) @@ -160,13 +166,42 @@ def test_config_refusal_names_the_file_and_tells_it_to_read_the_environment(): def test_environment_refusal_gives_the_command_without_a_config_step(): - text = render_refusal(_refusal(source=EnvironmentSource())) + text = render_refusal(_refusal(reason=UnsafeMasterKeyReason.NOT_SET, source=EnvironmentSource())) assert f"the {MASTER_KEY_ENV_VAR} environment variable" in text assert GENERATE_MASTER_KEY_COMMAND in text assert "os.environ/" not in text +@pytest.mark.parametrize( + ("master_key", "general_settings", "environment_master_key", "is_set"), + [ + (None, {}, None, False), + ("sk-1234", {"master_key": "sk-1234"}, None, False), + ("sk-1234", {}, "sk-1234", True), + ("sk-1234", {"master_key": "sk-1234"}, "", True), + ], +) +def test_refusal_records_whether_the_environment_variable_is_already_set( + master_key: str | None, general_settings: Mapping[str, object], environment_master_key: str | None, is_set: bool +): + refusal = _verdict(master_key, general_settings, environment_master_key=environment_master_key) + + assert isinstance(refusal, UnsafeMasterKeyRefused) + assert refusal.environment_variable_is_set is is_set + + +@pytest.mark.parametrize("source", [EnvironmentSource(), ConfigFileSource(config_file_path="/app/config.yaml")]) +def test_refusal_never_tells_a_user_with_an_exported_key_to_append_to_the_env_file( + source: ConfigFileSource | EnvironmentSource, +): + text = render_refusal(_refusal(source=source, environment_variable_is_set=True)) + + assert PRINT_NEW_MASTER_KEY_COMMAND in text + assert "tee" not in text + assert "wins over .env" in text + + def test_unset_key_refusal_says_nothing_supplied_one(): text = render_refusal(_refusal(reason=UnsafeMasterKeyReason.NOT_SET, source=EnvironmentSource())) @@ -242,3 +277,19 @@ def test_safe_and_overridden_keys_boot_without_announcing(verdict: MasterKeyBoot enforce_master_key_boot_verdict(verdict, announce=announced.append) assert announced == [] + + +def test_announced_fix_is_the_last_thing_a_crashing_process_prints(): + crash_after_announcing = ( + "from litellm.proxy.auth.master_key_boot_check import announce_on_stderr_at_exit\n" + "announce_on_stderr_at_exit('THE FIX')\n" + "print('buffered stdout')\n" + "raise RuntimeError('lifespan failed')\n" + ) + + completed = subprocess.run([sys.executable, "-c", crash_after_announcing], capture_output=True, text=True) + + assert completed.returncode != 0 + assert "RuntimeError: lifespan failed" in completed.stderr + assert completed.stderr.endswith("THE FIX") + assert completed.stdout == "buffered stdout\n"