From 67895955f6e9f40e78a96c89a642adc03b5d6a3c Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:20:08 +0000 Subject: [PATCH 1/9] feat(proxy): LITELLM_FIPS_MODE startup gate with provider assertion, TLS verify guard and loud password migration Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/fips.py | 151 ++++++++++++++++++ litellm/proxy/proxy_server.py | 31 +++- tests/integration/_support/process.py | 76 +++++++-- .../configuration/test_fips_mode_boot.py | 66 ++++++++ tests/integration/contracts.json | 15 ++ .../proxy/common_utils/test_fips.py | 118 ++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 104 ++++++++++++ 7 files changed, 546 insertions(+), 15 deletions(-) create mode 100644 litellm/proxy/common_utils/fips.py create mode 100644 tests/integration/configuration/test_fips_mode_boot.py create mode 100644 tests/test_litellm/proxy/common_utils/test_fips.py diff --git a/litellm/proxy/common_utils/fips.py b/litellm/proxy/common_utils/fips.py new file mode 100644 index 00000000000..33f35388567 --- /dev/null +++ b/litellm/proxy/common_utils/fips.py @@ -0,0 +1,151 @@ +import hashlib +import os +from collections.abc import Callable +from dataclasses import dataclass +from typing import Final + +from typing_extensions import assert_never + +FIPS_MODE_ENV_VAR: Final = "LITELLM_FIPS_MODE" +SSL_VERIFY_ENV_VAR: Final = "SSL_VERIFY" +SSL_VERIFY_SETTING: Final = "litellm_settings.ssl_verify" +REFUSAL_PREFIX: Final = "LiteLLM proxy refused to start" + +_TRUE_VALUES: Final = frozenset({"true", "1", "yes", "on"}) +_FALSE_VALUES: Final = frozenset({"false", "0", "no", "off", ""}) + + +@dataclass(frozen=True, slots=True) +class FipsModeOff: + pass + + +@dataclass(frozen=True, slots=True) +class FipsModeOn: + pass + + +@dataclass(frozen=True, slots=True) +class MalformedFipsMode: + value: str + + +FipsModeSetting = FipsModeOff | FipsModeOn | MalformedFipsMode + + +@dataclass(frozen=True, slots=True) +class ProviderDoesNotEnforceFips: + pass + + +@dataclass(frozen=True, slots=True) +class TlsVerificationDisabled: + sources: tuple[str, ...] + + +FipsBootRefusal = MalformedFipsMode | ProviderDoesNotEnforceFips | TlsVerificationDisabled +FipsBootVerdict = FipsModeOff | FipsModeOn | FipsBootRefusal + + +class FipsModeError(Exception): + pass + + +def parse_fips_mode(raw: str | None) -> FipsModeSetting: + if raw is None: + return FipsModeOff() + normalized: Final = raw.strip().lower() + if normalized in _TRUE_VALUES: + return FipsModeOn() + if normalized in _FALSE_VALUES: + return FipsModeOff() + return MalformedFipsMode(value=raw) + + +def is_fips_mode(environ: Callable[[str], str | None] = os.environ.get) -> bool: + return isinstance(parse_fips_mode(environ(FIPS_MODE_ENV_VAR)), FipsModeOn) + + +def openssl_enforces_fips() -> bool: + """MD5 is not an approved digest, so an enforcing FIPS provider refuses it even when asked for security use.""" + try: + hashlib.md5(b"", usedforsecurity=True) + except ValueError: + return True + return False + + +def fips_boot_verdict( + *, + raw_fips_mode: str | None, + provider_enforces_fips: Callable[[], bool], + ssl_verify_environment: str | None, + ssl_verify_setting: object, +) -> FipsBootVerdict: + setting: Final = parse_fips_mode(raw_fips_mode) + match setting: + case FipsModeOff() | MalformedFipsMode(): + return setting + case FipsModeOn(): + pass + case _: + assert_never(setting) + disabled: Final = tuple( + source + for source, off in ( + (SSL_VERIFY_ENV_VAR, _is_off(ssl_verify_environment)), + (SSL_VERIFY_SETTING, _is_off(ssl_verify_setting)), + ) + if off + ) + if disabled: + return TlsVerificationDisabled(sources=disabled) + if not provider_enforces_fips(): + return ProviderDoesNotEnforceFips() + return setting + + +def enforce_fips_boot_verdict(verdict: FipsBootVerdict, announce: Callable[[str], object]) -> None: + match verdict: + case FipsModeOff() | FipsModeOn(): + return + case MalformedFipsMode() | ProviderDoesNotEnforceFips() | TlsVerificationDisabled(): + message: Final = render_refusal(verdict) + announce(f"\n{message}\n\n") + raise FipsModeError(message) + case _: + assert_never(verdict) + + +def render_refusal(refusal: FipsBootRefusal) -> str: + match refusal: + case MalformedFipsMode(value=value): + return ( + f"{REFUSAL_PREFIX}: {FIPS_MODE_ENV_VAR}={value} is not a boolean.\n" + f"Set {FIPS_MODE_ENV_VAR} to true or false, or unset it." + ) + case ProviderDoesNotEnforceFips(): + return ( + f"{REFUSAL_PREFIX}: {FIPS_MODE_ENV_VAR} is on but this Python does not enforce FIPS.\n" + "Its OpenSSL still allows non-approved algorithms (MD5 succeeded), so passwords and keys would be\n" + "protected with algorithms the FIPS 140-3 policy forbids. Run the proxy from a FIPS image whose\n" + f"OpenSSL FIPS provider is enabled, or unset {FIPS_MODE_ENV_VAR} on a non-FIPS runtime." + ) + case TlsVerificationDisabled(sources=sources): + return ( + f"{REFUSAL_PREFIX}: {FIPS_MODE_ENV_VAR} is on but TLS certificate verification is disabled by " + f"{' and '.join(sources)}.\nFIPS deployments must verify upstream certificates, so remove the " + "override or point ssl_verify at a CA bundle instead." + ) + case _: + assert_never(refusal) + + +def _is_off(value: object) -> bool: + match value: + case bool(): + return value is False + case str(): + return value.strip().lower() in _FALSE_VALUES - {""} + case _: + return False diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dab7decd4dc..7e2e063ae88 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -404,6 +404,14 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id +from litellm.proxy.common_utils.fips import ( + FIPS_MODE_ENV_VAR, + SSL_VERIFY_ENV_VAR, + enforce_fips_boot_verdict, + fips_boot_verdict, + is_fips_mode, + openssl_enforces_fips, +) from litellm.proxy.common_utils.healthy_model_filter import ( get_hidden_unhealthy_model_names, is_healthy_only_listing_default, @@ -1277,6 +1285,16 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: if isinstance(worker_config, dict): await initialize_from_worker_config(worker_config) + enforce_fips_boot_verdict( + fips_boot_verdict( + raw_fips_mode=os.getenv(FIPS_MODE_ENV_VAR), + provider_enforces_fips=openssl_enforces_fips, + ssl_verify_environment=os.getenv(SSL_VERIFY_ENV_VAR), + ssl_verify_setting=litellm.ssl_verify, + ), + announce=announce_on_stderr_at_exit, + ) + enforce_master_key_boot_verdict( await with_stored_secrets_counted( master_key_boot_verdict( @@ -1316,10 +1334,21 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: try: result: Final = await migrate_passwords_to_scrypt_async(prisma_client) verbose_proxy_logger.info("Password migration: %s", result) + except ValueError as e: + verbose_proxy_logger.error( + "Password migration failed, so plaintext passwords stay unhashed in the database: %s. " + "This is what an OpenSSL FIPS provider reports when the hashing algorithm is not approved.", + e, + ) + if is_fips_mode(): + raise except Exception as e: verbose_proxy_logger.warning("Password migration skipped: %s", e) - asyncio.create_task(_run_pw_migration()) + if is_fips_mode(): + await _run_pw_migration() + else: + asyncio.create_task(_run_pw_migration()) async def _run_agent_grant_id_migration() -> None: from litellm.proxy.agent_endpoints.agent_registry import ( diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 5c44beaa570..6dec960097b 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -68,8 +68,15 @@ def owned_proxy( yield owned.gateway +@dataclass(frozen=True, slots=True) +class LaunchedProxy: + port: int + process: subprocess.Popen[bytes] + log: Path + + @contextmanager -def owned_proxy_process( +def launched_proxy( gateway: Gateway, directory: Path, overrides: Mapping[str, str], @@ -77,7 +84,7 @@ def owned_proxy_process( config: Path | None = None, remove_environment: tuple[str, ...] = (), workers: int = 1, -) -> Iterator[OwnedProxy]: +) -> Iterator[LaunchedProxy]: with socket.socket() as reserve: reserve.bind(("127.0.0.1", 0)) port: Final = reserve.getsockname()[1] @@ -116,18 +123,7 @@ def owned_proxy_process( start_new_session=True, ) try: - with httpx.Client(base_url=f"http://127.0.0.1:{port}", timeout=15, trust_env=False) as client: - deadline: Final = time.monotonic() + 70 - while True: - assert process.poll() is None, "Owned proxy exited before readiness" - try: - if client.get("/health/readiness", timeout=2).status_code == 200: - break - except httpx.TransportError: - pass - assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded" - time.sleep(0.1) - yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), process, log_path) + yield LaunchedProxy(port, process, log_path) finally: root_stopped: Final = stop_root_process(process) residual: Final = group_members(process.pid) @@ -142,3 +138,55 @@ def owned_proxy_process( survivors: Final = group_members(process.pid) assert not survivors, "Owned proxy child survived cleanup" assert root_stopped and not remaining, "Owned proxy required forced cleanup" + + +def refused_boot_log( + gateway: Gateway, + directory: Path, + overrides: Mapping[str, str], + *, + config: Path | None = None, +) -> str: + """Start the proxy and return its log once it exits non-zero instead of becoming ready.""" + with launched_proxy(gateway, directory, overrides, config=config) as launched: + with httpx.Client(base_url=f"http://127.0.0.1:{launched.port}", timeout=15, trust_env=False) as client: + deadline: Final = time.monotonic() + 70 + while launched.process.poll() is None: + try: + ready: Final = client.get("/health/readiness", timeout=2).status_code == 200 + except httpx.TransportError: + ready = False + assert not ready, f"Proxy became ready instead of refusing to boot:\n{launched.log.read_text()}" + assert time.monotonic() < deadline, "Proxy neither exited nor became ready within the deadline" + time.sleep(0.1) + assert launched.process.returncode != 0, ( + f"Proxy exited 0 instead of refusing to boot:\n{launched.log.read_text()}" + ) + return launched.log.read_text() + + +@contextmanager +def owned_proxy_process( + gateway: Gateway, + directory: Path, + overrides: Mapping[str, str], + *, + config: Path | None = None, + remove_environment: tuple[str, ...] = (), + workers: int = 1, +) -> Iterator[OwnedProxy]: + with launched_proxy( + gateway, directory, overrides, config=config, remove_environment=remove_environment, workers=workers + ) as launched: + with httpx.Client(base_url=f"http://127.0.0.1:{launched.port}", timeout=15, trust_env=False) as client: + deadline: Final = time.monotonic() + 70 + while True: + assert launched.process.poll() is None, "Owned proxy exited before readiness" + try: + if client.get("/health/readiness", timeout=2).status_code == 200: + break + except httpx.TransportError: + pass + assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded" + time.sleep(0.1) + yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), launched.process, launched.log) diff --git a/tests/integration/configuration/test_fips_mode_boot.py b/tests/integration/configuration/test_fips_mode_boot.py new file mode 100644 index 00000000000..5cf69f21ab3 --- /dev/null +++ b/tests/integration/configuration/test_fips_mode_boot.py @@ -0,0 +1,66 @@ +"""LITELLM_FIPS_MODE is a boot gate: the proxy refuses to serve unless the process really enforces FIPS. + +Every leg launches the real proxy binary against the suite's Postgres and asserts on what an operator sees: +exit status and the refusal text in the log. Nothing is patched inside the proxy. +""" + +import hashlib +from pathlib import Path +from typing import Final + +import pytest +import yaml + +from tests.integration._support.client import Gateway +from tests.integration._support.process import owned_proxy, refused_boot_log + +REFUSAL: Final = "LiteLLM proxy refused to start" + + +def _this_python_enforces_fips() -> bool: + try: + hashlib.md5(b"probe", usedforsecurity=True) + except ValueError: + return True + return False + + +@pytest.mark.covers("other.configuration.fips_mode.refuses_boot_when_provider_does_not_enforce_fips") +def test_fips_mode_refuses_to_serve_when_this_python_does_not_enforce_fips(gateway: Gateway, tmp_path: Path) -> None: + if _this_python_enforces_fips(): + pytest.skip("Runner OpenSSL enforces FIPS, so this leg cannot observe the non-enforcing refusal") + log: Final = refused_boot_log(gateway, tmp_path, {"LITELLM_FIPS_MODE": "true"}) + assert REFUSAL in log, log + assert "LITELLM_FIPS_MODE" in log and "does not enforce FIPS" in log, log + + +@pytest.mark.covers("other.configuration.fips_mode.refuses_boot_with_tls_verification_disabled") +@pytest.mark.parametrize("source", ("environment", "config")) +def test_fips_mode_refuses_to_serve_with_tls_verification_disabled( + gateway: Gateway, tmp_path: Path, source: str +) -> None: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = tmp_path / "ssl_verify_off.yaml" + settings: Final = {**config.get("litellm_settings", {}), "ssl_verify": False} + path.write_text(yaml.safe_dump({**config, "litellm_settings": settings})) + log: Final = ( + refused_boot_log(gateway, tmp_path, {"LITELLM_FIPS_MODE": "true", "SSL_VERIFY": "false"}) + if source == "environment" + else refused_boot_log(gateway, tmp_path, {"LITELLM_FIPS_MODE": "true"}, config=path) + ) + assert REFUSAL in log, log + assert "TLS certificate verification is disabled" in log, log + assert ("SSL_VERIFY" if source == "environment" else "litellm_settings.ssl_verify") in log, log + + +@pytest.mark.covers("other.configuration.fips_mode.refuses_boot_on_malformed_value") +def test_fips_mode_refuses_to_serve_on_a_value_that_is_not_a_boolean(gateway: Gateway, tmp_path: Path) -> None: + log: Final = refused_boot_log(gateway, tmp_path, {"LITELLM_FIPS_MODE": "enforced"}) + assert REFUSAL in log, log + assert "LITELLM_FIPS_MODE=enforced" in log and "true or false" in log, log + + +@pytest.mark.covers("other.configuration.fips_mode.off_leaves_boot_unchanged") +def test_fips_mode_off_serves_even_with_tls_verification_disabled(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, {"LITELLM_FIPS_MODE": "false", "SSL_VERIFY": "false"}) as candidate: + assert candidate.client.get("/health/readiness").status_code == 200 diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index a8d9cf1df8b..673f62fd731 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -76,6 +76,21 @@ "tests/integration/configuration/test_effective_settings.py::test_credential_value_update_and_model_reload_reach_provider": [ "mgmt.credential.update.saved_value_reaches_wire" ], + "tests/integration/configuration/test_fips_mode_boot.py::test_fips_mode_refuses_to_serve_when_this_python_does_not_enforce_fips": [ + "other.configuration.fips_mode.refuses_boot_when_provider_does_not_enforce_fips" + ], + "tests/integration/configuration/test_fips_mode_boot.py::test_fips_mode_refuses_to_serve_with_tls_verification_disabled[environment]": [ + "other.configuration.fips_mode.refuses_boot_with_tls_verification_disabled" + ], + "tests/integration/configuration/test_fips_mode_boot.py::test_fips_mode_refuses_to_serve_with_tls_verification_disabled[config]": [ + "other.configuration.fips_mode.refuses_boot_with_tls_verification_disabled" + ], + "tests/integration/configuration/test_fips_mode_boot.py::test_fips_mode_refuses_to_serve_on_a_value_that_is_not_a_boolean": [ + "other.configuration.fips_mode.refuses_boot_on_malformed_value" + ], + "tests/integration/configuration/test_fips_mode_boot.py::test_fips_mode_off_serves_even_with_tls_verification_disabled": [ + "other.configuration.fips_mode.off_leaves_boot_unchanged" + ], "tests/integration/management/test_partial_update_sequences.py::test_denied_key_update_preserves_saved_grants_and_serving": [ "mgmt.key.update.denied_request_preserves_effective_state" ], diff --git a/tests/test_litellm/proxy/common_utils/test_fips.py b/tests/test_litellm/proxy/common_utils/test_fips.py new file mode 100644 index 00000000000..e934cafbdda --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_fips.py @@ -0,0 +1,118 @@ +import hashlib + +import pytest + +from litellm.proxy.common_utils.fips import ( + FipsModeError, + FipsModeOff, + FipsModeOn, + MalformedFipsMode, + ProviderDoesNotEnforceFips, + TlsVerificationDisabled, + enforce_fips_boot_verdict, + fips_boot_verdict, + is_fips_mode, + openssl_enforces_fips, + parse_fips_mode, +) + + +def _verdict( + raw: str | None, + *, + enforcing: bool = True, + ssl_env: str | None = None, + ssl_setting: object = True, +): + return fips_boot_verdict( + raw_fips_mode=raw, + provider_enforces_fips=lambda: enforcing, + ssl_verify_environment=ssl_env, + ssl_verify_setting=ssl_setting, + ) + + +@pytest.mark.parametrize("raw", [None, "false", "0", "no", "off", "", " False "]) +def test_unset_and_false_spellings_leave_fips_mode_off(raw): + assert parse_fips_mode(raw) == FipsModeOff() + assert is_fips_mode({"LITELLM_FIPS_MODE": raw}.get) is False + + +@pytest.mark.parametrize("raw", ["true", "1", "yes", "on", " TRUE "]) +def test_true_spellings_turn_fips_mode_on(raw): + assert parse_fips_mode(raw) == FipsModeOn() + assert is_fips_mode({"LITELLM_FIPS_MODE": raw}.get) is True + + +@pytest.mark.parametrize("raw", ["enforced", "2", "strict", "yes please"]) +def test_anything_else_is_malformed_and_refused_with_the_offending_value(raw): + assert parse_fips_mode(raw) == MalformedFipsMode(value=raw) + with pytest.raises(FipsModeError) as refused: + enforce_fips_boot_verdict(_verdict(raw, enforcing=False), announce=lambda _: None) + assert f"LITELLM_FIPS_MODE={raw} is not a boolean" in str(refused.value) + assert "true or false" in str(refused.value) + + +def test_off_never_consults_the_provider_or_tls_settings(): + def explode() -> bool: + raise AssertionError("provider probe must not run while FIPS mode is off") + + verdict = fips_boot_verdict( + raw_fips_mode=None, provider_enforces_fips=explode, ssl_verify_environment="false", ssl_verify_setting=False + ) + assert verdict == FipsModeOff() + enforce_fips_boot_verdict(verdict, announce=lambda _: pytest.fail("nothing to announce when off")) + + +def test_on_with_an_enforcing_provider_and_verified_tls_boots(): + verdict = _verdict("true", enforcing=True, ssl_env="true", ssl_setting="/etc/ssl/certs/ca.pem") + assert verdict == FipsModeOn() + enforce_fips_boot_verdict(verdict, announce=lambda _: pytest.fail("nothing to announce when on")) + + +def test_on_with_a_non_enforcing_provider_is_refused_and_names_the_fix(): + announced = [] + with pytest.raises(FipsModeError) as refused: + enforce_fips_boot_verdict(_verdict("true", enforcing=False), announce=announced.append) + message = str(refused.value) + assert message.startswith("LiteLLM proxy refused to start") + assert "LITELLM_FIPS_MODE is on but this Python does not enforce FIPS" in message + assert "FIPS image" in message + assert announced == [f"\n{message}\n\n"] + + +@pytest.mark.parametrize( + "ssl_env, ssl_setting, sources", + [ + ("false", True, ("SSL_VERIFY",)), + ("0", True, ("SSL_VERIFY",)), + (None, False, ("litellm_settings.ssl_verify",)), + (None, "False", ("litellm_settings.ssl_verify",)), + ("no", False, ("SSL_VERIFY", "litellm_settings.ssl_verify")), + ], +) +def test_disabled_tls_verification_is_refused_naming_every_source(ssl_env, ssl_setting, sources): + verdict = _verdict("true", enforcing=True, ssl_env=ssl_env, ssl_setting=ssl_setting) + assert verdict == TlsVerificationDisabled(sources=sources) + with pytest.raises(FipsModeError) as refused: + enforce_fips_boot_verdict(verdict, announce=lambda _: None) + assert "TLS certificate verification is disabled by " + " and ".join(sources) in str(refused.value) + + +@pytest.mark.parametrize("ssl_setting", [True, "true", "/etc/ssl/certs/ca.pem", None, ""]) +def test_verified_or_custom_bundle_tls_settings_are_not_treated_as_disabled(ssl_setting): + assert _verdict("true", enforcing=True, ssl_setting=ssl_setting) == FipsModeOn() + + +def test_disabled_tls_is_reported_before_the_provider_so_operators_see_config_mistakes_first(): + assert _verdict("true", enforcing=False, ssl_env="false") == TlsVerificationDisabled(sources=("SSL_VERIFY",)) + assert _verdict("true", enforcing=False) == ProviderDoesNotEnforceFips() + + +def test_provider_probe_agrees_with_whether_md5_is_usable_for_security_here(): + try: + hashlib.md5(b"", usedforsecurity=True) + except ValueError: + assert openssl_enforces_fips() is True + else: + assert openssl_enforces_fips() is False diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 89fd9c5c9d4..42ad00d6b13 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1681,6 +1681,110 @@ async def test_proxy_startup_refuses_an_unsafe_master_key_even_when_the_database assert ("could not be checked" in announced[0]) == key_can_have_encrypted_the_database +@pytest.mark.asyncio +async def test_proxy_startup_refuses_fips_mode_when_this_python_does_not_enforce_fips(monkeypatch, tmp_path): + from fastapi import FastAPI + + from litellm.proxy.common_utils.fips import FipsModeError + from litellm.proxy.proxy_server import proxy_startup_event + + _, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.openssl_enforces_fips", lambda: False) + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + + with pytest.raises(FipsModeError): + async with proxy_startup_event(FastAPI()): + pass + + assert len(announced) == 1 + assert "does not enforce FIPS" in announced[0] + + +@pytest.mark.asyncio +async def test_proxy_startup_refuses_fips_mode_when_the_config_disables_tls_verification(monkeypatch, tmp_path): + import yaml + from fastapi import FastAPI + + from litellm.proxy.common_utils.fips import FipsModeError + from litellm.proxy.proxy_server import proxy_startup_event + + config_path, announced = _boot_with_general_settings(monkeypatch, tmp_path, {"master_key": "sk-a-safe-master-key"}) + config_path.write_text( + yaml.dump({"general_settings": {"master_key": "sk-a-safe-master-key"}, "litellm_settings": {"ssl_verify": False}}) + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.openssl_enforces_fips", lambda: True) + monkeypatch.setattr(litellm, "ssl_verify", True) + monkeypatch.setenv("LITELLM_FIPS_MODE", "true") + + with pytest.raises(FipsModeError): + async with proxy_startup_event(FastAPI()): + pass + + assert "TLS certificate verification is disabled by litellm_settings.ssl_verify" in announced[0] + + +class _PrismaClientWhoseUserTableCannotHash: + class _Table: + async def find_many(self, where): + raise ValueError("[digital envelope routines] unsupported") + + class _Db: + litellm_usertable = None + + def __init__(self): + self.db = self._Db() + self.db.litellm_usertable = self._Table() + self.writer_db = self.db + + async def connect(self): + pass + + async def disconnect(self): + pass + + async def check_view_exists(self): + pass + + async def health_check(self): + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fips_mode", ["true", "false"]) +async def test_proxy_startup_surfaces_a_password_migration_crypto_failure(monkeypatch, tmp_path, caplog, fips_mode): + 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("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable") + cannot_hash = _PrismaClientWhoseUserTableCannotHash() + + async def connected(**kwargs): + return cannot_hash + + monkeypatch.setattr("litellm.proxy.proxy_server.ProxyStartupEvent._setup_prisma_client", connected) + monkeypatch.setattr("litellm.proxy.proxy_server.openssl_enforces_fips", lambda: True) + monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + if fips_mode == "true": + with pytest.raises(ValueError, match="digital envelope routines"): + async with proxy_startup_event(FastAPI()): + pass + else: + async with proxy_startup_event(FastAPI()): + await asyncio.sleep(0) + + failures = [r.getMessage() for r in caplog.records if "Password migration failed" in r.getMessage()] + assert len(failures) == 1 + assert "plaintext passwords stay unhashed" in failures[0] + assert "digital envelope routines" in failures[0] + + class _DatabaseWithOneStoredCredential: def __init__(self, ciphertext): self._ciphertext = ciphertext From 5c237840f4055a5a0058d09598ec51a66fae8e27 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:31:01 +0000 Subject: [PATCH 2/9] refactor(proxy): match ssl_verify off detection to runtime str_to_bool semantics Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/fips.py | 4 +++- tests/integration/_support/process.py | 15 ++++++++++----- .../test_litellm/proxy/common_utils/test_fips.py | 6 +++--- 3 files changed, 16 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/common_utils/fips.py b/litellm/proxy/common_utils/fips.py index 33f35388567..8a3c26b3220 100644 --- a/litellm/proxy/common_utils/fips.py +++ b/litellm/proxy/common_utils/fips.py @@ -6,6 +6,8 @@ from typing import Final from typing_extensions import assert_never +from litellm.secret_managers.main import str_to_bool + FIPS_MODE_ENV_VAR: Final = "LITELLM_FIPS_MODE" SSL_VERIFY_ENV_VAR: Final = "SSL_VERIFY" SSL_VERIFY_SETTING: Final = "litellm_settings.ssl_verify" @@ -146,6 +148,6 @@ def _is_off(value: object) -> bool: case bool(): return value is False case str(): - return value.strip().lower() in _FALSE_VALUES - {""} + return str_to_bool(value) is False case _: return False diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index 6dec960097b..f65ab0e8093 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -140,6 +140,13 @@ def launched_proxy( assert root_stopped and not remaining, "Owned proxy required forced cleanup" +def _is_ready(client: httpx.Client) -> bool: + try: + return client.get("/health/readiness", timeout=2).status_code == 200 + except httpx.TransportError: + return False + + def refused_boot_log( gateway: Gateway, directory: Path, @@ -152,11 +159,9 @@ def refused_boot_log( with httpx.Client(base_url=f"http://127.0.0.1:{launched.port}", timeout=15, trust_env=False) as client: deadline: Final = time.monotonic() + 70 while launched.process.poll() is None: - try: - ready: Final = client.get("/health/readiness", timeout=2).status_code == 200 - except httpx.TransportError: - ready = False - assert not ready, f"Proxy became ready instead of refusing to boot:\n{launched.log.read_text()}" + assert not _is_ready(client), ( + f"Proxy became ready instead of refusing to boot:\n{launched.log.read_text()}" + ) assert time.monotonic() < deadline, "Proxy neither exited nor became ready within the deadline" time.sleep(0.1) assert launched.process.returncode != 0, ( diff --git a/tests/test_litellm/proxy/common_utils/test_fips.py b/tests/test_litellm/proxy/common_utils/test_fips.py index e934cafbdda..4b50ae4ad20 100644 --- a/tests/test_litellm/proxy/common_utils/test_fips.py +++ b/tests/test_litellm/proxy/common_utils/test_fips.py @@ -85,10 +85,10 @@ def test_on_with_a_non_enforcing_provider_is_refused_and_names_the_fix(): "ssl_env, ssl_setting, sources", [ ("false", True, ("SSL_VERIFY",)), - ("0", True, ("SSL_VERIFY",)), + (" FALSE ", True, ("SSL_VERIFY",)), (None, False, ("litellm_settings.ssl_verify",)), (None, "False", ("litellm_settings.ssl_verify",)), - ("no", False, ("SSL_VERIFY", "litellm_settings.ssl_verify")), + ("false", False, ("SSL_VERIFY", "litellm_settings.ssl_verify")), ], ) def test_disabled_tls_verification_is_refused_naming_every_source(ssl_env, ssl_setting, sources): @@ -99,7 +99,7 @@ def test_disabled_tls_verification_is_refused_naming_every_source(ssl_env, ssl_s assert "TLS certificate verification is disabled by " + " and ".join(sources) in str(refused.value) -@pytest.mark.parametrize("ssl_setting", [True, "true", "/etc/ssl/certs/ca.pem", None, ""]) +@pytest.mark.parametrize("ssl_setting", [True, "true", "/etc/ssl/certs/ca.pem", None, "", "0", "no"]) def test_verified_or_custom_bundle_tls_settings_are_not_treated_as_disabled(ssl_setting): assert _verdict("true", enforcing=True, ssl_setting=ssl_setting) == FipsModeOn() From 6ae5df66c1d24ba5ac460c2338ea01585753b04a Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:45:39 +0000 Subject: [PATCH 3/9] refactor(proxy): drop tautological fips probe test and satisfy CodeQL return checks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/fips.py | 23 ++++++++----------- .../proxy/common_utils/test_fips.py | 12 ---------- 2 files changed, 10 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/common_utils/fips.py b/litellm/proxy/common_utils/fips.py index 8a3c26b3220..07bd4d2462c 100644 --- a/litellm/proxy/common_utils/fips.py +++ b/litellm/proxy/common_utils/fips.py @@ -121,9 +121,9 @@ def enforce_fips_boot_verdict(verdict: FipsBootVerdict, announce: Callable[[str] def render_refusal(refusal: FipsBootRefusal) -> str: match refusal: - case MalformedFipsMode(value=value): + case MalformedFipsMode(): return ( - f"{REFUSAL_PREFIX}: {FIPS_MODE_ENV_VAR}={value} is not a boolean.\n" + f"{REFUSAL_PREFIX}: {FIPS_MODE_ENV_VAR}={refusal.value} is not a boolean.\n" f"Set {FIPS_MODE_ENV_VAR} to true or false, or unset it." ) case ProviderDoesNotEnforceFips(): @@ -133,21 +133,18 @@ def render_refusal(refusal: FipsBootRefusal) -> str: "protected with algorithms the FIPS 140-3 policy forbids. Run the proxy from a FIPS image whose\n" f"OpenSSL FIPS provider is enabled, or unset {FIPS_MODE_ENV_VAR} on a non-FIPS runtime." ) - case TlsVerificationDisabled(sources=sources): + case TlsVerificationDisabled(): return ( f"{REFUSAL_PREFIX}: {FIPS_MODE_ENV_VAR} is on but TLS certificate verification is disabled by " - f"{' and '.join(sources)}.\nFIPS deployments must verify upstream certificates, so remove the " + f"{' and '.join(refusal.sources)}.\nFIPS deployments must verify upstream certificates, so remove the " "override or point ssl_verify at a CA bundle instead." ) - case _: - assert_never(refusal) + return assert_never(refusal) def _is_off(value: object) -> bool: - match value: - case bool(): - return value is False - case str(): - return str_to_bool(value) is False - case _: - return False + if isinstance(value, bool): + return value is False + if isinstance(value, str): + return str_to_bool(value) is False + return False diff --git a/tests/test_litellm/proxy/common_utils/test_fips.py b/tests/test_litellm/proxy/common_utils/test_fips.py index 4b50ae4ad20..0e8138856c0 100644 --- a/tests/test_litellm/proxy/common_utils/test_fips.py +++ b/tests/test_litellm/proxy/common_utils/test_fips.py @@ -1,5 +1,3 @@ -import hashlib - import pytest from litellm.proxy.common_utils.fips import ( @@ -12,7 +10,6 @@ from litellm.proxy.common_utils.fips import ( enforce_fips_boot_verdict, fips_boot_verdict, is_fips_mode, - openssl_enforces_fips, parse_fips_mode, ) @@ -107,12 +104,3 @@ def test_verified_or_custom_bundle_tls_settings_are_not_treated_as_disabled(ssl_ def test_disabled_tls_is_reported_before_the_provider_so_operators_see_config_mistakes_first(): assert _verdict("true", enforcing=False, ssl_env="false") == TlsVerificationDisabled(sources=("SSL_VERIFY",)) assert _verdict("true", enforcing=False) == ProviderDoesNotEnforceFips() - - -def test_provider_probe_agrees_with_whether_md5_is_usable_for_security_here(): - try: - hashlib.md5(b"", usedforsecurity=True) - except ValueError: - assert openssl_enforces_fips() is True - else: - assert openssl_enforces_fips() is False From e6a5d32cd9925caedbd1eb067b86772031dcf556 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 06:07:49 +0000 Subject: [PATCH 4/9] test(proxy): inject the fake Prisma client through the module boundary instead of patching _setup_prisma_client Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/test_proxy_server.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 42ad00d6b13..39e4a417d19 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -1733,7 +1733,7 @@ class _PrismaClientWhoseUserTableCannotHash: class _Db: litellm_usertable = None - def __init__(self): + def __init__(self, database_url, proxy_logging_obj): self.db = self._Db() self.db.litellm_usertable = self._Table() self.writer_db = self.db @@ -1744,6 +1744,9 @@ class _PrismaClientWhoseUserTableCannotHash: async def disconnect(self): pass + def start_view_setup_task(self): + pass + async def check_view_exists(self): pass @@ -1761,12 +1764,7 @@ async def test_proxy_startup_surfaces_a_password_migration_crypto_failure(monkey _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("DATABASE_URL", "postgresql://nobody:nothing@127.0.0.1:1/unreachable") - cannot_hash = _PrismaClientWhoseUserTableCannotHash() - - async def connected(**kwargs): - return cannot_hash - - monkeypatch.setattr("litellm.proxy.proxy_server.ProxyStartupEvent._setup_prisma_client", connected) + monkeypatch.setattr("litellm.proxy.proxy_server.PrismaClient", _PrismaClientWhoseUserTableCannotHash) monkeypatch.setattr("litellm.proxy.proxy_server.openssl_enforces_fips", lambda: True) monkeypatch.setenv("LITELLM_FIPS_MODE", fips_mode) From 48574cc53cb0c164f6ea7a04e6cf3c02bf018673 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:10:17 +0000 Subject: [PATCH 5/9] feat(mcp): derive bridge and session keys with HKDF-SHA256 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/auth/user_api_key_auth_mcp.py | 17 ++- .../mcp_server/bridge_token_flow.py | 9 +- .../mcp_server/gateway_dcr_flow.py | 36 ++++- .../bridge_credentials.py | 86 ++++++------ .../outbound_credentials/key_derivation.py | 58 ++++++++ .../session_credentials.py | 86 ++++++++---- .../test_bridge_credentials.py | 126 +++++++++++++++++- .../test_session_credentials.py | 119 +++++++++++++++++ 8 files changed, 457 insertions(+), 80 deletions(-) create mode 100644 litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 43d62ba8b22..4ff913aa89d 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -25,6 +25,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti NotBridgeEnvelope, envelope_keys_from_master_key, is_bridge_envelope_shaped, + legacy_envelope_keys_from_master_key, resolve_bridge_envelope, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( @@ -869,7 +870,13 @@ class MCPRequestHandler: await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) keys: Final = envelope_keys_from_master_key(master_key) - result: Final = resolve_bridge_envelope(authorization_value, keys, datetime.now(timezone.utc), server.server_id) + result: Final = resolve_bridge_envelope( + authorization_value, + keys, + datetime.now(timezone.utc), + server.server_id, + legacy_keys=legacy_envelope_keys_from_master_key(master_key), + ) match result: case BridgeEnvelopeAdmitted(): return result @@ -1005,6 +1012,7 @@ class MCPRequestHandler: SessionBearerInvalid, SessionSigningConfigError, active_session_signing_keys, + legacy_session_keys_from_master_key, resolve_session_bearer, ) from litellm.proxy.proxy_server import master_key @@ -1018,7 +1026,12 @@ class MCPRequestHandler: if isinstance(keys, SessionSigningConfigError): verbose_logger.error("mcp gateway session admission rejected: %s", keys.detail) raise HTTPException(status_code=500, detail="Server misconfigured: mcp_session_token_signing is invalid") - result: Final = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc)) + result: Final = resolve_session_bearer( + authorization_value, + keys, + datetime.now(timezone.utc), + legacy_keys=legacy_session_keys_from_master_key(master_key, keys), + ) match result: case SessionBearerAdmitted(): try: diff --git a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py index e5d11271c67..ed41ee5b286 100644 --- a/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py +++ b/litellm/proxy/_experimental/mcp_server/bridge_token_flow.py @@ -867,6 +867,7 @@ async def _prepare_bridge_refresh( from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( # noqa: PLC0415 # inline import avoids a module-load circular import BridgeRefreshOpened, envelope_keys_from_master_key, + legacy_envelope_keys_from_master_key, open_bridge_refresh_envelope, ) from litellm.proxy.proxy_server import ( # noqa: PLC0415 # inline import avoids a module-load circular import @@ -878,7 +879,13 @@ async def _prepare_bridge_refresh( if not refresh_value: return "invalid_refresh" keys: Final = envelope_keys_from_master_key(master_key) - opened: Final = open_bridge_refresh_envelope(refresh_value, keys, datetime.now(timezone.utc), mcp_server.server_id) + opened: Final = open_bridge_refresh_envelope( + refresh_value, + keys, + datetime.now(timezone.utc), + mcp_server.server_id, + legacy_keys=legacy_envelope_keys_from_master_key(master_key), + ) if not isinstance(opened, BridgeRefreshOpened): return "invalid_refresh" failure: Final = await _revalidate_active_subject(opened.identity) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index e66504af47a..ab922a67b2d 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -67,6 +67,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionRefreshOpened, SessionSigningConfigError, active_session_signing_keys, + legacy_session_keys_from_master_key, open_session_refresh_bearer, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -75,8 +76,11 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MintedSessionToken, OpenedSessionToken, SessionAudience, + SessionExpired, + SessionKeys, SessionPrincipal, SessionSigningKeys, + SessionTokenOpenError, is_session_refresh_token, is_session_token, mint_session_refresh_token, @@ -1236,6 +1240,7 @@ async def aggregate_token( keys=keys, now=now, issue=issue, + legacy_keys=legacy_session_keys_from_master_key(master_key, keys), ) if grant_type == TOKEN_EXCHANGE_GRANT_TYPE: return await _token_exchange_grant( @@ -1395,10 +1400,13 @@ async def _refresh_token_grant( keys: SessionSigningKeys, now: datetime, issue: _GrantIssuer, + legacy_keys: SessionKeys | None = None, ) -> Response: if not refresh_token: return _oauth_error(400, "invalid_request", "refresh_token is required") - opened: Final = open_session_refresh_bearer(refresh_token, keys, now, expected_client_id=client_id) + opened: Final = open_session_refresh_bearer( + refresh_token, keys, now, expected_client_id=client_id, legacy_keys=legacy_keys + ) if not isinstance(opened, SessionRefreshOpened): return _oauth_error(400, "invalid_grant", "the refresh token is invalid for this client") if _resource_conflicts_with_scope(request, resource, opened.principal.resource_server_id): @@ -1459,7 +1467,13 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non verbose_logger.error("mcp_gateway_dcr revoke rejected: %s", keys.detail) return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid") now: Final = datetime.now(timezone.utc) - opened: Final = open_session_refresh_bearer(token, keys, now, expected_client_id=client_id) + opened: Final = open_session_refresh_bearer( + token, + keys, + now, + expected_client_id=client_id, + legacy_keys=legacy_session_keys_from_master_key(master_key, keys), + ) if isinstance(opened, SessionRefreshOpened): burned: Final = await _SingleUseGuard(cache).claim( f"{_USED_REFRESH_CACHE_PREFIX}{opened.jti}", SESSION_REFRESH_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS @@ -1469,6 +1483,19 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS) +def _open_with_legacy( + open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], + token: str, + keys: SessionSigningKeys, + legacy_keys: SessionKeys | None, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + opened: Final = open_token(token, keys, now) + if isinstance(opened, OpenedSessionToken) or isinstance(opened, SessionExpired) or legacy_keys is None: + return opened + return open_token(token, legacy_keys, now) + + def _inactive_introspection_response() -> Response: """RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason (wrong family, bad signature, expired, revoked, or a deactivated user), answers 200 @@ -1527,10 +1554,11 @@ async def introspect_gateway_token( verbose_logger.error("mcp_gateway_dcr introspect rejected: %s", keys.detail) return _oauth_error(500, "server_error", keys.detail) now: Final = datetime.now(timezone.utc) + legacy_keys: Final = legacy_session_keys_from_master_key(master_key, keys) if is_session_token(token): - opened = open_session_token(token, keys, now) + opened = _open_with_legacy(open_session_token, token, keys, legacy_keys, now) elif is_session_refresh_token(token): - opened = open_session_refresh_token(token, keys, now) + opened = _open_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now) else: return _inactive_introspection_response() if not isinstance(opened, OpenedSessionToken): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index ab5fa65480e..749813df4d0 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -9,7 +9,8 @@ to the upstream MCP server (consumer). This module is the pure surface for both token-endpoint and admission wiring live in their respective call sites. """ -import hashlib +import os +from collections.abc import Callable from datetime import datetime from functools import lru_cache from typing import Final, Literal, TypeAlias @@ -32,56 +33,55 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import open_envelope, open_refresh_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + hkdf_sha256, + legacy_kdf_grace_enabled, + legacy_scrypt, +) _SIGNING_KEY_DOMAIN: Final = b"litellm-mcp-bridge:envelope-signing:" _ENCRYPTION_KEY_DOMAIN: Final = b"litellm-mcp-bridge:envelope-encryption:" -# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation, which -# makes offline guessing of a candidate master key memory-hard rather than a bare hash comparison. -_SCRYPT_N: Final = 2**15 -_SCRYPT_R: Final = 8 -_SCRYPT_P: Final = 1 -# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales -# with every work factor and a future p or r bump does not trip "memory limit exceeded". -_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 -_DERIVED_KEY_BYTES: Final = 32 - @lru_cache(maxsize=8) def envelope_keys_from_master_key(master_key: str) -> EnvelopeKeys: """Derive the envelope signing and encryption keys from the proxy master key. - A memory-hard scrypt KDF (RFC 7914) over two distinct domain-label salts yields two - independent 256-bit subkeys from the one secret, so the producer (mint) and consumer - (open) agree on keys without persisting any. scrypt is used rather than a bare hash or - HMAC so that a captured envelope is not a cheap offline oracle for the master key: each - candidate guess costs a full memory-hard derivation, which is what protects a deployment - whose master key is weaker than it should be. The result is cached (the master key is - fixed for a process), so the KDF runs once per key and adds nothing to the per-request - admission path. The derivation is deterministic; rotating ``master_key`` invalidates - every outstanding envelope, which is the intended behavior for a signing-key change. + HKDF-SHA256 (RFC 5869) over the master key, once per domain label as ``info``, yields + two independent 256-bit subkeys from the one secret, so the producer (mint) and consumer + (open) agree on keys without persisting any. No memory-hard guessing cost is needed: + the master key is a high-entropy secret, so a captured envelope is not a cheap offline + oracle for it either way. The result is cached (the master key is fixed for a process), + so the KDF runs once per key and adds nothing to the per-request admission path. The + derivation is deterministic; rotating ``master_key`` invalidates every outstanding + envelope, which is the intended behavior for a signing-key change. """ - signing: Final = hashlib.scrypt( - master_key.encode(), - salt=_SIGNING_KEY_DOMAIN, - n=_SCRYPT_N, - r=_SCRYPT_R, - p=_SCRYPT_P, - maxmem=_SCRYPT_MAXMEM, - dklen=_DERIVED_KEY_BYTES, - ).hex() - encryption: Final = hashlib.scrypt( - master_key.encode(), - salt=_ENCRYPTION_KEY_DOMAIN, - n=_SCRYPT_N, - r=_SCRYPT_R, - p=_SCRYPT_P, - maxmem=_SCRYPT_MAXMEM, - dklen=_DERIVED_KEY_BYTES, - ).hex() + return EnvelopeKeys( + signing_key=SecretStr(hkdf_sha256(master_key, _SIGNING_KEY_DOMAIN)), + encryption_key=SecretStr(hkdf_sha256(master_key, _ENCRYPTION_KEY_DOMAIN)), + ) + + +@lru_cache(maxsize=8) +def _legacy_envelope_keys(master_key: str) -> EnvelopeKeys | None: + """The pre-rotation scrypt subkeys, kept so envelopes minted before the KDF rotation + can still be opened during the opt-in grace window. ``None`` when scrypt is unavailable.""" + signing: Final = legacy_scrypt(master_key, _SIGNING_KEY_DOMAIN) + encryption: Final = legacy_scrypt(master_key, _ENCRYPTION_KEY_DOMAIN) + if signing is None or encryption is None: + return None return EnvelopeKeys(signing_key=SecretStr(signing), encryption_key=SecretStr(encryption)) +def legacy_envelope_keys_from_master_key( + master_key: str, environ: Callable[[str], str | None] = os.environ.get +) -> EnvelopeKeys | None: + """The legacy envelope keys when the operator enabled the KDF grace window, else ``None``.""" + if not legacy_kdf_grace_enabled(environ): + return None + return _legacy_envelope_keys(master_key) + + def build_bridge_token_response( identity: EnvelopeIdentity, grant: UpstreamTokenGrant, @@ -136,6 +136,7 @@ def open_bridge_refresh_envelope( keys: EnvelopeKeys, now: datetime, expected_server_id: str, + legacy_keys: EnvelopeKeys | None = None, ) -> BridgeRefreshResult: """Open a refresh envelope a bridge ``oauth_delegate`` client presented on a refresh_token grant. @@ -150,7 +151,9 @@ def open_bridge_refresh_envelope( candidate: Final = _strip_bearer(refresh_value) if not is_refresh_envelope(candidate): return BridgeRefreshInvalid() - opened: Final = open_refresh_envelope(candidate, keys, now) + opened = open_refresh_envelope(candidate, keys, now) + if not isinstance(opened, OpenedRefreshEnvelope) and legacy_keys is not None: + opened = open_refresh_envelope(candidate, legacy_keys, now) if not isinstance(opened, OpenedRefreshEnvelope): return BridgeRefreshInvalid() if opened.identity.server_id != expected_server_id: @@ -207,6 +210,7 @@ def resolve_bridge_envelope( keys: EnvelopeKeys, now: datetime, expected_server_id: str, + legacy_keys: EnvelopeKeys | None = None, ) -> BridgeEnvelopeResult: """Classify an ``Authorization`` value presented to a bridge ``oauth_delegate`` server. @@ -233,7 +237,9 @@ def resolve_bridge_envelope( return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() - opened: Final = open_envelope(candidate, keys, now) + opened = open_envelope(candidate, keys, now) + if not isinstance(opened, OpenedEnvelope) and legacy_keys is not None: + opened = open_envelope(candidate, legacy_keys, now) if not isinstance(opened, OpenedEnvelope): return BridgeEnvelopeInvalid() if opened.identity.server_id != expected_server_id: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py new file mode 100644 index 00000000000..98c2218f2eb --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py @@ -0,0 +1,58 @@ +"""Key derivation for the MCP outbound-credential token families. + +Bridge envelopes and gateway session tokens are keyed by HKDF-SHA256 (RFC 5869) over the +high-entropy proxy ``master_key`` with a per-family domain label as ``info``. The legacy +scrypt derivation is kept for one release as an opt-in read fallback so tokens minted +before the rotation still open during the grace window; it is never used under FIPS. +""" + +import hashlib +import os +from collections.abc import Callable +from typing import Final + +from cryptography.hazmat.primitives.hashes import SHA256 +from cryptography.hazmat.primitives.kdf.hkdf import HKDF + +from litellm.proxy.common_utils.fips import is_fips_mode + +LEGACY_KDF_GRACE_ENV_VAR: Final = "LITELLM_MCP_LEGACY_KDF_GRACE" +DERIVED_KEY_BYTES: Final = 32 + +# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation. +_SCRYPT_N: Final = 2**15 +_SCRYPT_R: Final = 8 +_SCRYPT_P: Final = 1 +# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales +# with every work factor and a future p or r bump does not trip "memory limit exceeded". +_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 + + +def hkdf_sha256(master_key: str, info: bytes) -> str: + """Derive a 256-bit subkey from ``master_key`` bound to the domain label ``info``.""" + return HKDF(algorithm=SHA256(), length=DERIVED_KEY_BYTES, salt=None, info=info).derive(master_key.encode()).hex() + + +def legacy_scrypt(master_key: str, salt: bytes) -> str | None: + """The pre-rotation derivation. ``None`` when hashlib.scrypt is unavailable (some + OpenSSL builds, including FIPS providers, do not expose it).""" + try: + return hashlib.scrypt( + master_key.encode(), + salt=salt, + n=_SCRYPT_N, + r=_SCRYPT_R, + p=_SCRYPT_P, + maxmem=_SCRYPT_MAXMEM, + dklen=DERIVED_KEY_BYTES, + ).hex() + except (AttributeError, ValueError): + return None + + +def legacy_kdf_grace_enabled(environ: Callable[[str], str | None] = os.environ.get) -> bool: + """True only when the operator opted the legacy read fallback in and the deployment is + not in FIPS mode (scrypt is not a FIPS-approved KDF).""" + if is_fips_mode(environ): + return False + return (environ(LEGACY_KDF_GRACE_ENV_VAR) or "").strip().lower() == "true" diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index df2bbdba345..02b231fed27 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -8,20 +8,26 @@ recovered litellm user (consumer), reloading the live user record and policy bef anything runs. This module is the pure surface for both sides; the token-endpoint and admission wiring live in their respective call sites. -The signing key is derived with the same memory-hard scrypt construction as +The signing key is derived with the same HKDF-SHA256 construction as :func:`~.bridge_credentials.envelope_keys_from_master_key` but under a distinct domain label, so session tokens and bridge envelopes never share key material: a token of one family is unverifiable in the other by key separation, on top of the distinct issuers, prefixes, and claim shapes. """ -import hashlib +import os +from collections.abc import Callable from datetime import datetime from functools import lru_cache from typing import Final, Literal, TypeAlias from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + hkdf_sha256, + legacy_kdf_grace_enabled, + legacy_scrypt, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( AsymmetricSessionKeys, OpenedSessionToken, @@ -38,39 +44,49 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i _SESSION_SIGNING_KEY_DOMAIN: Final = b"litellm-mcp-gateway:session-signing:" -# scrypt work factors (RFC 7914), identical to the envelope KDF: memory-hard so a captured -# session token is not a cheap offline oracle for the master key. -_SCRYPT_N: Final = 2**15 -_SCRYPT_R: Final = 8 -_SCRYPT_P: Final = 1 -_SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 -_DERIVED_KEY_BYTES: Final = 32 - @lru_cache(maxsize=8) def session_keys_from_master_key(master_key: str) -> SessionKeys: """Derive the session signing key from the proxy master key. - A memory-hard scrypt KDF (RFC 7914) over a session-specific domain-label salt yields a - 256-bit subkey from the one secret, so the producer (mint) and consumer (open) agree on - the key without persisting any. The domain label differs from both envelope labels in - :mod:`.bridge_credentials`, so compromise or misuse of one token family never crosses - into the other. The result is cached (the master key is fixed for a process); rotating - ``master_key`` invalidates every outstanding session, which is the intended behavior - for a signing-key change. + HKDF-SHA256 (RFC 5869) over the master key with a session-specific domain label as + ``info`` yields a 256-bit subkey from the one secret, so the producer (mint) and + consumer (open) agree on the key without persisting any. The label differs from both + envelope labels in :mod:`.bridge_credentials`, so compromise or misuse of one token + family never crosses into the other. The result is cached (the master key is fixed for + a process); rotating ``master_key`` invalidates every outstanding session, which is + the intended behavior for a signing-key change. """ - signing: Final = hashlib.scrypt( - master_key.encode(), - salt=_SESSION_SIGNING_KEY_DOMAIN, - n=_SCRYPT_N, - r=_SCRYPT_R, - p=_SCRYPT_P, - maxmem=_SCRYPT_MAXMEM, - dklen=_DERIVED_KEY_BYTES, - ).hex() + return SessionKeys(signing_key=SecretStr(hkdf_sha256(master_key, _SESSION_SIGNING_KEY_DOMAIN))) + + +@lru_cache(maxsize=8) +def _legacy_session_keys(master_key: str) -> SessionKeys | None: + """The pre-rotation scrypt signing key, kept so session tokens minted before the KDF + rotation can still be opened during the opt-in grace window. ``None`` when scrypt is + unavailable.""" + signing: Final = legacy_scrypt(master_key, _SESSION_SIGNING_KEY_DOMAIN) + if signing is None: + return None return SessionKeys(signing_key=SecretStr(signing)) +def legacy_session_keys_from_master_key( + master_key: str, + active: SessionSigningKeys, + environ: Callable[[str], str | None] = os.environ.get, +) -> SessionKeys | None: + """The legacy session key when the operator enabled the KDF grace window, else ``None``. + + The fallback only applies to the default master-key HS256 path: a configured + ``mcp_session_token_signing`` RS256 key was never scrypt-derived, so an + :class:`AsymmetricSessionKeys` ``active`` gets no legacy key even with grace on. + """ + if not isinstance(active, SessionKeys) or not legacy_kdf_grace_enabled(environ): + return None + return _legacy_session_keys(master_key) + + class SessionSigningPreviousKey(BaseModel): """One retired key in ``mcp_session_token_signing.previous_public_keys``: its ``kid`` and the PEM public half (inline or an ``os.environ/`` reference).""" @@ -214,6 +230,7 @@ def resolve_session_bearer( authorization_value: str, keys: SessionSigningKeys, now: datetime, + legacy_keys: SessionKeys | None = None, ) -> SessionBearerResult: """Classify an ``Authorization`` value presented at the aggregate MCP edge. @@ -232,7 +249,13 @@ def resolve_session_bearer( return SessionBearerInvalid() if not is_session_token(candidate): return NotSessionBearer() - opened: Final = open_session_token(candidate, keys, now) + opened = open_session_token(candidate, keys, now) + if ( + not isinstance(opened, OpenedSessionToken) + and not isinstance(opened, SessionExpired) + and legacy_keys is not None + ): + opened = open_session_token(candidate, legacy_keys, now) if isinstance(opened, OpenedSessionToken): return SessionBearerAdmitted(principal=opened.principal) return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) @@ -265,6 +288,7 @@ def open_session_refresh_bearer( keys: SessionSigningKeys, now: datetime, expected_client_id: str, + legacy_keys: SessionKeys | None = None, ) -> SessionRefreshResult: """Open a session refresh token presented on a ``refresh_token`` grant. @@ -279,7 +303,13 @@ def open_session_refresh_bearer( candidate: Final = _strip_bearer(refresh_value) if not is_session_refresh_token(candidate): return SessionRefreshInvalid() - opened: Final = open_session_refresh_token(candidate, keys, now) + opened = open_session_refresh_token(candidate, keys, now) + if ( + not isinstance(opened, OpenedSessionToken) + and not isinstance(opened, SessionExpired) + and legacy_keys is not None + ): + opened = open_session_refresh_token(candidate, legacy_keys, now) if not isinstance(opened, OpenedSessionToken): return SessionRefreshInvalid() if opened.principal.client_id != expected_client_id: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index f8fb22469f1..6718285ef0a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -8,6 +8,8 @@ unopenable, and envelope minted for a different server); the producer helper rou through the consumer; and no path leaks the upstream token in a repr. """ +import hashlib +import hmac from datetime import datetime, timedelta, timezone import pytest @@ -23,9 +25,17 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti build_bridge_token_response, envelope_keys_from_master_key, is_bridge_envelope_shaped, + legacy_envelope_keys_from_master_key, open_bridge_refresh_envelope, resolve_bridge_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, +) +from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( ENVELOPE_PREFIX, EnvelopeIdentity, @@ -146,11 +156,10 @@ def test_key_derivation_signing_key_meets_hs256_floor_for_short_master_key(): assert len(keys.signing_key.get_secret_value()) >= 32 -def test_key_derivation_is_cached_so_the_memory_hard_kdf_runs_once_per_key(): - """The scrypt KDF is intentionally expensive to resist offline guessing, so it must be cached: - repeated calls for the same master key return the identical object rather than re-deriving, - keeping the per-request admission path free. Returning a distinct object each call would mean - the cache was dropped and every open would pay the memory-hard cost.""" +def test_key_derivation_is_cached_so_the_hkdf_runs_once_per_key(): + """Repeated calls for the same master key return the identical object rather than + re-deriving, keeping the per-request admission path free. Returning a distinct object + each call would mean the cache was dropped and every open would pay the KDF cost.""" first = envelope_keys_from_master_key("sk-cache-probe-key-9988776655") assert envelope_keys_from_master_key("sk-cache-probe-key-9988776655") is first @@ -318,3 +327,110 @@ def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer(): assert is_bridge_envelope_shaped("Bearer sk-some-litellm-key") is False assert is_bridge_envelope_shaped("plain-upstream-token") is False assert is_bridge_envelope_shaped("") is False + + +def _rfc5869(ikm: bytes, info: bytes) -> bytes: + """Independent RFC 5869 HKDF-SHA256, hand-written on hmac/hashlib so a wrong KDF + construction in the product cannot agree with it.""" + prk = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() + return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:32] + + +def _environ(values: dict[str, str]): + return values.get + + +def _legacy_keys(master_key: str) -> EnvelopeKeys: + def derive(salt: bytes) -> str: + return hashlib.scrypt( + master_key.encode(), salt=salt, n=2**15, r=8, p=1, maxmem=128 * 2**15 * 8 * 2, dklen=32 + ).hex() + + return EnvelopeKeys( + signing_key=SecretStr(derive(b"litellm-mcp-bridge:envelope-signing:")), + encryption_key=SecretStr(derive(b"litellm-mcp-bridge:envelope-encryption:")), + ) + + +def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): + keys = envelope_keys_from_master_key("sk-1234") + assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-signing:").hex() + assert ( + keys.encryption_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-encryption:").hex() + ) + # literal pin computed from the same hand implementation (RFC 5869, salt = 32 zero bytes) + assert keys.signing_key.get_secret_value() == "331032bfe2b8d86bd00586ea69c0d854140019ca0e0adbe4dd0cc9e06cb36605" + + +def test_session_key_derivation_matches_an_independent_rfc5869_vector(): + keys = session_keys_from_master_key("sk-1234") + assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-gateway:session-signing:").hex() + assert keys.signing_key.get_secret_value() == "511f07b16ceeb5fba1660033785d4b9a4a6e80c749a51d639a688c287671a0e6" + + +def test_the_three_derived_keys_are_pairwise_distinct(): + envelope = envelope_keys_from_master_key(_MASTER_KEY) + session = session_keys_from_master_key(_MASTER_KEY) + derived = { + envelope.signing_key.get_secret_value(), + envelope.encryption_key.get_secret_value(), + session.signing_key.get_secret_value(), + } + assert len(derived) == 3 + + +def test_new_envelopes_open_under_the_hkdf_keys_with_no_legacy_fallback(): + keys = envelope_keys_from_master_key(_MASTER_KEY) + result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeAdmitted) + + +def test_legacy_envelope_opens_during_the_grace_window(): + legacy = _legacy_keys(_MASTER_KEY) + token = _sealed_token(legacy) + grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + assert grace == legacy + result = resolve_bridge_envelope( + token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace + ) + assert isinstance(result, BridgeEnvelopeAdmitted) + assert result.identity == _IDENTITY + + +def test_legacy_envelope_is_invalid_without_grace(): + token = _sealed_token(_legacy_keys(_MASTER_KEY)) + assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({})) is None + result = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + assert isinstance(result, BridgeEnvelopeInvalid) + + +def test_grace_is_disabled_under_fips(): + grace = legacy_envelope_keys_from_master_key( + _MASTER_KEY, + environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}), + ) + assert grace is None + + +@pytest.mark.parametrize("value", ["yes", "1", "on", "TRUE ", "True"]) +def test_grace_is_enabled_only_by_case_insensitive_true(value: str): + result = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: value})) + expected = value.strip().lower() == "true" + assert (result is not None) is expected + + +def test_legacy_refresh_envelope_opens_during_grace(): + legacy = _legacy_keys(_MASTER_KEY) + refresh = _sealed_refresh(legacy) + grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + result = open_bridge_refresh_envelope( + refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace + ) + assert isinstance(result, BridgeRefreshOpened) + assert result.refresh.refresh_token.get_secret_value() == _UPSTREAM_REFRESH + + +def test_legacy_refresh_envelope_is_invalid_without_grace(): + refresh = _sealed_refresh(_legacy_keys(_MASTER_KEY)) + result = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + assert isinstance(result, BridgeRefreshInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 992b1e8632b..396e7b965e1 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -1,5 +1,6 @@ """Tests for the session-token KDF and the edge/token-endpoint resolvers.""" +import hashlib from datetime import datetime, timedelta, timezone import pytest @@ -11,6 +12,7 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials import session_credentials from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( NotSessionBearer, SessionBearerAdmitted, @@ -19,11 +21,15 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionRefreshOpened, SessionSigningConfigError, is_session_bearer_shaped, + legacy_session_keys_from_master_key, open_session_refresh_bearer, resolve_session_bearer, resolve_session_signing_keys, session_keys_from_master_key, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_TTL_SECONDS, AsymmetricSessionKeys, @@ -34,6 +40,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i mint_session_token, session_public_key_pem, ) +from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) MASTER_KEY = "sk-master-key-for-tests" @@ -230,3 +237,115 @@ def test_signing_config_error_detail_never_leaks_key_material(): ) assert isinstance(resolved, SessionSigningConfigError) assert pem.splitlines()[1] not in resolved.detail + + +def _environ(values: dict[str, str]): + return values.get + + +def _legacy_session_keys(master_key: str) -> SessionKeys: + signing = hashlib.scrypt( + master_key.encode(), + salt=b"litellm-mcp-gateway:session-signing:", + n=2**15, + r=8, + p=1, + maxmem=128 * 2**15 * 8 * 2, + dklen=32, + ).hex() + return SessionKeys(signing_key=SecretStr(signing)) + + +def test_new_session_tokens_open_under_the_hkdf_key_with_no_legacy_fallback(): + result = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) + assert isinstance(result, SessionBearerAdmitted) + + +def test_legacy_session_token_opens_during_the_grace_window(): + legacy = _legacy_session_keys(MASTER_KEY) + minted = mint_session_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + assert grace == legacy + result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) + assert isinstance(result, SessionBearerAdmitted) + assert result.principal == PRINCIPAL + + +def test_legacy_session_token_is_invalid_without_grace(): + minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({})) is None + result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is False + + +def test_legacy_session_key_is_never_derived_for_configured_rs256_keys(): + active = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") + grace = legacy_session_keys_from_master_key( + MASTER_KEY, active, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}) + ) + assert grace is None + + +def test_grace_is_disabled_under_fips(): + grace = legacy_session_keys_from_master_key( + MASTER_KEY, + KEYS, + environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}), + ) + assert grace is None + + +def test_legacy_session_refresh_token_opens_during_grace(): + legacy = _legacy_session_keys(MASTER_KEY) + minted = mint_session_refresh_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + result = open_session_refresh_bearer( + minted.token.get_secret_value(), + KEYS, + NOW, + expected_client_id="llm_client_abc", + legacy_keys=grace, + ) + assert isinstance(result, SessionRefreshOpened) + assert result.principal == PRINCIPAL + + +def test_legacy_session_refresh_token_is_invalid_without_grace(): + minted = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + result = open_session_refresh_bearer( + minted.token.get_secret_value(), KEYS, NOW, expected_client_id="llm_client_abc" + ) + assert isinstance(result, SessionRefreshInvalid) + + +def test_expired_token_under_the_primary_key_still_reports_expired_during_grace(): + legacy = _legacy_session_keys(MASTER_KEY) + minted = mint_session_token(PRINCIPAL, legacy, NOW) + assert isinstance(minted, MintedSessionToken) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) + assert isinstance(result, SessionBearerInvalid) + assert result.expired is True + + +def test_fallback_tries_the_active_key_before_legacy(monkeypatch): + calls: list[object] = [] + real = session_credentials.open_session_token + + def spy(candidate, keys, now): + calls.append(keys) + return real(candidate, keys, now) + + monkeypatch.setattr(session_credentials, "open_session_token", spy) + grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + assert isinstance(minted, MintedSessionToken) + result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) + assert isinstance(result, SessionBearerAdmitted) + assert calls[0] is KEYS From 3062035adaf588aff668431181b0d4c251e9ec36 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:10:17 +0000 Subject: [PATCH 6/9] test(integration): cover mcp kdf rotation window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/mcp/test_mcp_kdf_rotation.py | 243 ++++++++++++++++++ 1 file changed, 243 insertions(+) create mode 100644 tests/integration/mcp/test_mcp_kdf_rotation.py diff --git a/tests/integration/mcp/test_mcp_kdf_rotation.py b/tests/integration/mcp/test_mcp_kdf_rotation.py new file mode 100644 index 00000000000..b915759dbda --- /dev/null +++ b/tests/integration/mcp/test_mcp_kdf_rotation.py @@ -0,0 +1,243 @@ +"""Bridge envelopes and gateway session tokens are keyed by HKDF-SHA256 over the master key. + +Every cell mints a bearer in the test under keys the test derives itself (HKDF-SHA256 for the current +construction, scrypt for the legacy one) and presents it to a real proxy. The proxy is the only party +that derives keys from the master key, so admission proves which KDF it runs and whether the legacy +grace window is honoured. +""" + +import hashlib +import uuid +from datetime import datetime, timezone +from pathlib import Path +from typing import Final + +import httpx +import pytest +from cryptography.hazmat.primitives.hashes import SHA256 +from cryptography.hazmat.primitives.kdf.hkdf import HKDF +from integration._support.client import Gateway +from integration._support.mcp import McpPeer, mcp_peer, register_mcp, tool_calls +from integration._support.oauth_server import oauth_server +from integration._support.process import owned_proxy +from pydantic import SecretStr + +from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( + EnvelopeKeys, + SealedEnvelope, + UpstreamTokenGrant, + key_hash_identity, + mint_envelope, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + MintedSessionToken, + SessionKeys, + SessionPrincipal, + mint_session_token, +) +from litellm.proxy.utils import hash_token + +ACCEPT: Final = {"Accept": "application/json, text/event-stream"} +TOOLS_LIST: Final = {"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}} +ENVELOPE_SIGNING: Final = b"litellm-mcp-bridge:envelope-signing:" +ENVELOPE_ENCRYPTION: Final = b"litellm-mcp-bridge:envelope-encryption:" +SESSION_SIGNING: Final = b"litellm-mcp-gateway:session-signing:" +GRACE: Final = {"LITELLM_MCP_LEGACY_KDF_GRACE": "true"} + + +def _hkdf(master_key: str, info: bytes) -> str: + return HKDF(algorithm=SHA256(), length=32, salt=None, info=info).derive(master_key.encode()).hex() + + +def _scrypt(master_key: str, salt: bytes) -> str: + return hashlib.scrypt(master_key.encode(), salt=salt, n=2**15, r=8, p=1, maxmem=2**27, dklen=32).hex() + + +def _hkdf_envelope_keys(master_key: str) -> EnvelopeKeys: + return EnvelopeKeys( + signing_key=SecretStr(_hkdf(master_key, ENVELOPE_SIGNING)), + encryption_key=SecretStr(_hkdf(master_key, ENVELOPE_ENCRYPTION)), + ) + + +def _scrypt_envelope_keys(master_key: str) -> EnvelopeKeys: + return EnvelopeKeys( + signing_key=SecretStr(_scrypt(master_key, ENVELOPE_SIGNING)), + encryption_key=SecretStr(_scrypt(master_key, ENVELOPE_ENCRYPTION)), + ) + + +def _hkdf_session_keys(master_key: str) -> SessionKeys: + return SessionKeys(signing_key=SecretStr(_hkdf(master_key, SESSION_SIGNING))) + + +def _scrypt_session_keys(master_key: str) -> SessionKeys: + return SessionKeys(signing_key=SecretStr(_scrypt(master_key, SESSION_SIGNING))) + + +def _envelope(identity_server: str, key: str, keys: EnvelopeKeys, upstream_token: str) -> str: + sealed: Final = mint_envelope( + key_hash_identity(server_id=identity_server, key_hash=hash_token(key)), + UpstreamTokenGrant(access_token=SecretStr(upstream_token), token_type="Bearer", expires_in=600), + keys, + datetime.now(timezone.utc), + ) + assert isinstance(sealed, SealedEnvelope), sealed + return sealed.token.get_secret_value() + + +def _session(user_id: str, keys: SessionKeys) -> str: + minted: Final = mint_session_token( + SessionPrincipal(user_id=user_id, client_id="dcr-integration"), keys, datetime.now(timezone.utc) + ) + assert isinstance(minted, MintedSessionToken), minted + return minted.token.get_secret_value() + + +def _register_bridge(scenario, peer: McpPeer, alias: str, issuer: str) -> str: + return register_mcp( + scenario, + peer, + alias, + auth_type="oauth_delegate", + dcr_bridge=True, + issuer=issuer, + authorization_url=issuer + "/authorize", + token_url=issuer + "/token", + registration_url=issuer + "/register", + ) + + +def _tools_list(client: httpx.Client, path: str, bearer: str) -> httpx.Response: + return client.post(path, headers={**ACCEPT, "Authorization": f"Bearer {bearer}"}, json=TOOLS_LIST) + + +def _peer_authorizations(peer: McpPeer) -> tuple[bytes | None, ...]: + return tuple( + value if isinstance(value := call["headers"].get(b"authorization"), bytes) else None + for call in peer.drain() + if isinstance(call["headers"], dict) + ) + + +def _assert_admitted(response: httpx.Response, peer: McpPeer, upstream_token: str) -> None: + assert response.status_code == 200, response.text + assert "tools" in response.text, response.text + forwarded: Final = _peer_authorizations(peer) + assert forwarded and set(forwarded) == {f"Bearer {upstream_token}".encode()}, forwarded + + +def _assert_rejected(response: httpx.Response, peer: McpPeer) -> None: + assert response.status_code == 401, response.text + assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers + assert peer.drain() == () + + +def test_envelope_minted_under_hkdf_sha256_keys_is_admitted(gateway: Gateway) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + upstream_token: Final = "up-" + uuid.uuid4().hex + peer.drain() + envelope: Final = _envelope(identity, key, _hkdf_envelope_keys(gateway.key), upstream_token) + _assert_admitted(_tools_list(gateway.client, f"/{alias}/mcp", envelope), peer, upstream_token) + + +def test_envelope_minted_under_legacy_scrypt_keys_is_rejected_without_grace(gateway: Gateway) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + peer.drain() + envelope: Final = _envelope(identity, key, _scrypt_envelope_keys(gateway.key), "up-" + uuid.uuid4().hex) + _assert_rejected(_tools_list(gateway.client, f"/{alias}/mcp", envelope), peer) + + +@pytest.mark.parametrize("kdf", ("hkdf", "scrypt")) +def test_envelope_of_either_kdf_is_admitted_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path, kdf: str +) -> None: + with ( + owned_proxy(gateway, tmp_path, GRACE) as graced, + mcp_peer() as peer, + oauth_server() as auth, + graced.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + upstream_token: Final = "up-" + uuid.uuid4().hex + peer.drain() + keys: Final = _hkdf_envelope_keys(graced.key) if kdf == "hkdf" else _scrypt_envelope_keys(graced.key) + envelope: Final = _envelope(identity, key, keys, upstream_token) + _assert_admitted(_tools_list(graced.client, f"/{alias}/mcp", envelope), peer, upstream_token) + + +def test_envelope_under_a_foreign_master_key_is_rejected_during_grace(gateway: Gateway, tmp_path: Path) -> None: + with ( + owned_proxy(gateway, tmp_path, GRACE) as graced, + mcp_peer() as peer, + oauth_server() as auth, + graced.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + peer.drain() + foreign: Final = "sk-foreign-" + uuid.uuid4().hex + for keys in (_hkdf_envelope_keys(foreign), _scrypt_envelope_keys(foreign)): + _assert_rejected(_tools_list(graced.client, f"/{alias}/mcp", _envelope(identity, key, keys, "up")), peer) + + +def test_grace_variable_set_to_anything_but_true_keeps_legacy_envelopes_rejected( + gateway: Gateway, tmp_path: Path +) -> None: + with ( + owned_proxy(gateway, tmp_path, {"LITELLM_MCP_LEGACY_KDF_GRACE": "yes"}) as proxy, + mcp_peer() as peer, + oauth_server() as auth, + proxy.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + peer.drain() + envelope: Final = _envelope(identity, key, _scrypt_envelope_keys(proxy.key), "up") + _assert_rejected(_tools_list(proxy.client, f"/{alias}/mcp", envelope), peer) + + +def test_session_token_minted_under_hkdf_sha256_key_is_admitted(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + token: Final = _session(scenario.user(), _hkdf_session_keys(gateway.key)) + response: Final = _tools_list(gateway.client, "/mcp", token) + assert response.status_code == 200, response.text + assert "tools" in response.text, response.text + + +def test_session_token_minted_under_legacy_scrypt_key_is_rejected_without_grace(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + token: Final = _session(scenario.user(), _scrypt_session_keys(gateway.key)) + response: Final = _tools_list(gateway.client, "/mcp", token) + assert response.status_code == 401, response.text + assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers + + +@pytest.mark.parametrize("kdf", ("hkdf", "scrypt")) +def test_session_token_of_either_kdf_is_admitted_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path, kdf: str +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + keys: Final = _hkdf_session_keys(graced.key) if kdf == "hkdf" else _scrypt_session_keys(graced.key) + response: Final = _tools_list(graced.client, "/mcp", _session(scenario.user(), keys)) + assert response.status_code == 200, response.text + assert "tools" in response.text, response.text + + +def test_tampered_session_token_is_rejected_during_grace(gateway: Gateway, tmp_path: Path) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + token: Final = _session(scenario.user(), _scrypt_session_keys(graced.key)) + tampered: Final = token[:-2] + ("AA" if token[-2:] != "AA" else "BB") + response: Final = _tools_list(graced.client, "/mcp", tampered) + assert response.status_code == 401, response.text + assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers From 31c456837af7065b3374bed23309855bd8b9f2b4 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:17:45 +0000 Subject: [PATCH 7/9] refactor(mcp): single-assignment kdf fallback helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/gateway_dcr_flow.py | 20 +---- .../bridge_credentials.py | 37 ++++++-- .../outbound_credentials/key_derivation.py | 4 +- .../session_credentials.py | 34 +++++--- .../test_bridge_credentials.py | 87 +++++++++++-------- .../test_session_credentials.py | 85 ++++++++---------- 6 files changed, 142 insertions(+), 125 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index ab922a67b2d..508c6f51593 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -68,6 +68,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionSigningConfigError, active_session_signing_keys, legacy_session_keys_from_master_key, + open_session_credential_with_legacy, open_session_refresh_bearer, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -76,11 +77,9 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i MintedSessionToken, OpenedSessionToken, SessionAudience, - SessionExpired, SessionKeys, SessionPrincipal, SessionSigningKeys, - SessionTokenOpenError, is_session_refresh_token, is_session_token, mint_session_refresh_token, @@ -1483,19 +1482,6 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS) -def _open_with_legacy( - open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], - token: str, - keys: SessionSigningKeys, - legacy_keys: SessionKeys | None, - now: datetime, -) -> OpenedSessionToken | SessionTokenOpenError: - opened: Final = open_token(token, keys, now) - if isinstance(opened, OpenedSessionToken) or isinstance(opened, SessionExpired) or legacy_keys is None: - return opened - return open_token(token, legacy_keys, now) - - def _inactive_introspection_response() -> Response: """RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason (wrong family, bad signature, expired, revoked, or a deactivated user), answers 200 @@ -1556,9 +1542,9 @@ async def introspect_gateway_token( now: Final = datetime.now(timezone.utc) legacy_keys: Final = legacy_session_keys_from_master_key(master_key, keys) if is_session_token(token): - opened = _open_with_legacy(open_session_token, token, keys, legacy_keys, now) + opened = open_session_credential_with_legacy(open_session_token, token, keys, legacy_keys, now) elif is_session_refresh_token(token): - opened = _open_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now) + opened = open_session_credential_with_legacy(open_session_refresh_token, token, keys, legacy_keys, now) else: return _inactive_introspection_response() if not isinstance(opened, OpenedSessionToken): diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py index 749813df4d0..fb5a2e38236 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/bridge_credentials.py @@ -21,6 +21,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import EnvelopeIdentity, EnvelopeKeys, EnvelopeMintError, + EnvelopeOpenError, OpenedEnvelope, OpenedRefreshEnvelope, RefreshCredential, @@ -131,6 +132,34 @@ class BridgeRefreshInvalid(BaseModel): BridgeRefreshResult: TypeAlias = BridgeRefreshOpened | BridgeRefreshInvalid +def _open_envelope_with_fallback( + candidate: str, + keys: EnvelopeKeys, + legacy_keys: EnvelopeKeys | None, + now: datetime, +) -> OpenedEnvelope | EnvelopeOpenError: + """Open an access envelope under the active keys, then under the legacy keys when + grace supplied them.""" + primary: Final = open_envelope(candidate, keys, now) + if isinstance(primary, OpenedEnvelope) or legacy_keys is None: + return primary + return open_envelope(candidate, legacy_keys, now) + + +def _open_refresh_envelope_with_fallback( + candidate: str, + keys: EnvelopeKeys, + legacy_keys: EnvelopeKeys | None, + now: datetime, +) -> OpenedRefreshEnvelope | EnvelopeOpenError: + """Open a refresh envelope under the active keys, then under the legacy keys when + grace supplied them.""" + primary: Final = open_refresh_envelope(candidate, keys, now) + if isinstance(primary, OpenedRefreshEnvelope) or legacy_keys is None: + return primary + return open_refresh_envelope(candidate, legacy_keys, now) + + def open_bridge_refresh_envelope( refresh_value: str, keys: EnvelopeKeys, @@ -151,9 +180,7 @@ def open_bridge_refresh_envelope( candidate: Final = _strip_bearer(refresh_value) if not is_refresh_envelope(candidate): return BridgeRefreshInvalid() - opened = open_refresh_envelope(candidate, keys, now) - if not isinstance(opened, OpenedRefreshEnvelope) and legacy_keys is not None: - opened = open_refresh_envelope(candidate, legacy_keys, now) + opened: Final = _open_refresh_envelope_with_fallback(candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedRefreshEnvelope): return BridgeRefreshInvalid() if opened.identity.server_id != expected_server_id: @@ -237,9 +264,7 @@ def resolve_bridge_envelope( return BridgeEnvelopeInvalid() if not is_envelope(candidate): return NotBridgeEnvelope() - opened = open_envelope(candidate, keys, now) - if not isinstance(opened, OpenedEnvelope) and legacy_keys is not None: - opened = open_envelope(candidate, legacy_keys, now) + opened: Final = _open_envelope_with_fallback(candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedEnvelope): return BridgeEnvelopeInvalid() if opened.identity.server_id != expected_server_id: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py index 98c2218f2eb..9ea90386bc8 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/key_derivation.py @@ -19,12 +19,10 @@ from litellm.proxy.common_utils.fips import is_fips_mode LEGACY_KDF_GRACE_ENV_VAR: Final = "LITELLM_MCP_LEGACY_KDF_GRACE" DERIVED_KEY_BYTES: Final = 32 -# scrypt work factors (RFC 7914). n=2**15 with r=8/p=1 costs ~50ms and ~32MB per derivation. +# RFC 7914 work factors of the pre-rotation derivation; maxmem is twice the working set. _SCRYPT_N: Final = 2**15 _SCRYPT_R: Final = 8 _SCRYPT_P: Final = 1 -# scrypt's working-set is ~128 * N * r * p bytes; cap at twice that so the maxmem ceiling scales -# with every work factor and a future p or r bump does not trip "memory limit exceeded". _SCRYPT_MAXMEM: Final = 128 * _SCRYPT_N * _SCRYPT_R * _SCRYPT_P * 2 diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index 02b231fed27..396a821b0fa 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -36,6 +36,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SessionPrincipal, SessionRotatedPublicKey, SessionSigningKeys, + SessionTokenOpenError, is_session_refresh_token, is_session_token, open_session_refresh_token, @@ -249,13 +250,7 @@ def resolve_session_bearer( return SessionBearerInvalid() if not is_session_token(candidate): return NotSessionBearer() - opened = open_session_token(candidate, keys, now) - if ( - not isinstance(opened, OpenedSessionToken) - and not isinstance(opened, SessionExpired) - and legacy_keys is not None - ): - opened = open_session_token(candidate, legacy_keys, now) + opened: Final = open_session_credential_with_legacy(open_session_token, candidate, keys, legacy_keys, now) if isinstance(opened, OpenedSessionToken): return SessionBearerAdmitted(principal=opened.principal) return SessionBearerInvalid(expired=isinstance(opened, SessionExpired)) @@ -283,6 +278,23 @@ class SessionRefreshInvalid(BaseModel): SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid +def open_session_credential_with_legacy( + open_token: Callable[[str, SessionSigningKeys, datetime], OpenedSessionToken | SessionTokenOpenError], + token: str, + keys: SessionSigningKeys, + legacy_keys: SessionKeys | None, + now: datetime, +) -> OpenedSessionToken | SessionTokenOpenError: + """Open a session credential under the active keys, then under the legacy keys when + grace supplied them. An :class:`OpenedSessionToken` or :class:`SessionExpired` under + the active keys is final: expiry must still report as expiry, so only a signature or + format failure falls through to the legacy derivation.""" + primary: Final = open_token(token, keys, now) + if isinstance(primary, (OpenedSessionToken, SessionExpired)) or legacy_keys is None: + return primary + return open_token(token, legacy_keys, now) + + def open_session_refresh_bearer( refresh_value: str, keys: SessionSigningKeys, @@ -303,13 +315,7 @@ def open_session_refresh_bearer( candidate: Final = _strip_bearer(refresh_value) if not is_session_refresh_token(candidate): return SessionRefreshInvalid() - opened = open_session_refresh_token(candidate, keys, now) - if ( - not isinstance(opened, OpenedSessionToken) - and not isinstance(opened, SessionExpired) - and legacy_keys is not None - ): - opened = open_session_refresh_token(candidate, legacy_keys, now) + opened: Final = open_session_credential_with_legacy(open_session_refresh_token, candidate, keys, legacy_keys, now) if not isinstance(opened, OpenedSessionToken): return SessionRefreshInvalid() if opened.principal.client_id != expected_client_id: diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py index 6718285ef0a..673c418041b 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_bridge_credentials.py @@ -10,7 +10,10 @@ through the consumer; and no path leaks the upstream token in a repr. import hashlib import hmac +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final import pytest from pydantic import SecretStr @@ -29,13 +32,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credenti open_bridge_refresh_envelope, resolve_bridge_envelope, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( - LEGACY_KDF_GRACE_ENV_VAR, -) -from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( - session_keys_from_master_key, -) -from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( ENVELOPE_PREFIX, EnvelopeIdentity, @@ -47,6 +43,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import key_hash_identity, mint_envelope, ) +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, +) +from litellm.proxy.common_utils.fips import FIPS_MODE_ENV_VAR _NOW = datetime(2026, 7, 9, 12, 0, 0, tzinfo=timezone.utc) _MASTER_KEY = "sk-master-key-for-derivation-tests-0123456789" @@ -332,11 +335,11 @@ def test_is_bridge_envelope_shaped_rejects_non_envelope_bearer(): def _rfc5869(ikm: bytes, info: bytes) -> bytes: """Independent RFC 5869 HKDF-SHA256, hand-written on hmac/hashlib so a wrong KDF construction in the product cannot agree with it.""" - prk = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() + prk: Final = hmac.new(b"\x00" * 32, ikm, hashlib.sha256).digest() return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:32] -def _environ(values: dict[str, str]): +def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]: return values.get @@ -353,7 +356,7 @@ def _legacy_keys(master_key: str) -> EnvelopeKeys: def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): - keys = envelope_keys_from_master_key("sk-1234") + keys: Final = envelope_keys_from_master_key("sk-1234") assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-signing:").hex() assert ( keys.encryption_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-bridge:envelope-encryption:").hex() @@ -363,34 +366,38 @@ def test_envelope_key_derivation_matches_an_independent_rfc5869_vector(): def test_session_key_derivation_matches_an_independent_rfc5869_vector(): - keys = session_keys_from_master_key("sk-1234") + keys: Final = session_keys_from_master_key("sk-1234") assert keys.signing_key.get_secret_value() == _rfc5869(b"sk-1234", b"litellm-mcp-gateway:session-signing:").hex() assert keys.signing_key.get_secret_value() == "511f07b16ceeb5fba1660033785d4b9a4a6e80c749a51d639a688c287671a0e6" def test_the_three_derived_keys_are_pairwise_distinct(): - envelope = envelope_keys_from_master_key(_MASTER_KEY) - session = session_keys_from_master_key(_MASTER_KEY) - derived = { - envelope.signing_key.get_secret_value(), - envelope.encryption_key.get_secret_value(), - session.signing_key.get_secret_value(), - } + envelope: Final = envelope_keys_from_master_key(_MASTER_KEY) + session: Final = session_keys_from_master_key(_MASTER_KEY) + derived: Final = frozenset( + ( + envelope.signing_key.get_secret_value(), + envelope.encryption_key.get_secret_value(), + session.signing_key.get_secret_value(), + ) + ) assert len(derived) == 3 def test_new_envelopes_open_under_the_hkdf_keys_with_no_legacy_fallback(): - keys = envelope_keys_from_master_key(_MASTER_KEY) - result = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) + keys: Final = envelope_keys_from_master_key(_MASTER_KEY) + result: Final = resolve_bridge_envelope(_sealed_token(keys), keys, _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeAdmitted) def test_legacy_envelope_opens_during_the_grace_window(): - legacy = _legacy_keys(_MASTER_KEY) - token = _sealed_token(legacy) - grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + legacy: Final = _legacy_keys(_MASTER_KEY) + token: Final = _sealed_token(legacy) + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) assert grace == legacy - result = resolve_bridge_envelope( + result: Final = resolve_bridge_envelope( token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace ) assert isinstance(result, BridgeEnvelopeAdmitted) @@ -398,32 +405,36 @@ def test_legacy_envelope_opens_during_the_grace_window(): def test_legacy_envelope_is_invalid_without_grace(): - token = _sealed_token(_legacy_keys(_MASTER_KEY)) - assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({})) is None - result = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + token: Final = _sealed_token(_legacy_keys(_MASTER_KEY)) + assert legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ(MappingProxyType({}))) is None + result: Final = resolve_bridge_envelope(token, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) assert isinstance(result, BridgeEnvelopeInvalid) def test_grace_is_disabled_under_fips(): - grace = legacy_envelope_keys_from_master_key( + grace: Final = legacy_envelope_keys_from_master_key( _MASTER_KEY, - environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}), + environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})), ) assert grace is None -@pytest.mark.parametrize("value", ["yes", "1", "on", "TRUE ", "True"]) +@pytest.mark.parametrize("value", ("yes", "1", "on", "TRUE ", "True")) def test_grace_is_enabled_only_by_case_insensitive_true(value: str): - result = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: value})) - expected = value.strip().lower() == "true" + result: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: value})) + ) + expected: Final = value.strip().lower() == "true" assert (result is not None) is expected def test_legacy_refresh_envelope_opens_during_grace(): - legacy = _legacy_keys(_MASTER_KEY) - refresh = _sealed_refresh(legacy) - grace = legacy_envelope_keys_from_master_key(_MASTER_KEY, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - result = open_bridge_refresh_envelope( + legacy: Final = _legacy_keys(_MASTER_KEY) + refresh: Final = _sealed_refresh(legacy) + grace: Final = legacy_envelope_keys_from_master_key( + _MASTER_KEY, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + result: Final = open_bridge_refresh_envelope( refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID, legacy_keys=grace ) assert isinstance(result, BridgeRefreshOpened) @@ -431,6 +442,6 @@ def test_legacy_refresh_envelope_opens_during_grace(): def test_legacy_refresh_envelope_is_invalid_without_grace(): - refresh = _sealed_refresh(_legacy_keys(_MASTER_KEY)) - result = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) + refresh: Final = _sealed_refresh(_legacy_keys(_MASTER_KEY)) + result: Final = open_bridge_refresh_envelope(refresh, envelope_keys_from_master_key(_MASTER_KEY), _NOW, _SERVER_ID) assert isinstance(result, BridgeRefreshInvalid) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 396e7b965e1..13886982b37 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -1,7 +1,10 @@ """Tests for the session-token KDF and the edge/token-endpoint resolvers.""" import hashlib +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone +from types import MappingProxyType +from typing import Final import pytest from cryptography.hazmat.primitives import serialization @@ -12,7 +15,9 @@ from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials import session_credentials +from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( + LEGACY_KDF_GRACE_ENV_VAR, +) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( NotSessionBearer, SessionBearerAdmitted, @@ -27,9 +32,6 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent resolve_session_signing_keys, session_keys_from_master_key, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.key_derivation import ( - LEGACY_KDF_GRACE_ENV_VAR, -) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_TTL_SECONDS, AsymmetricSessionKeys, @@ -239,12 +241,12 @@ def test_signing_config_error_detail_never_leaks_key_material(): assert pem.splitlines()[1] not in resolved.detail -def _environ(values: dict[str, str]): +def _environ(values: Mapping[str, str]) -> Callable[[str], str | None]: return values.get def _legacy_session_keys(master_key: str) -> SessionKeys: - signing = hashlib.scrypt( + signing: Final = hashlib.scrypt( master_key.encode(), salt=b"litellm-mcp-gateway:session-signing:", n=2**15, @@ -257,53 +259,57 @@ def _legacy_session_keys(master_key: str) -> SessionKeys: def test_new_session_tokens_open_under_the_hkdf_key_with_no_legacy_fallback(): - result = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) + result: Final = resolve_session_bearer(f"Bearer {_access_token()}", KEYS, NOW) assert isinstance(result, SessionBearerAdmitted) def test_legacy_session_token_opens_during_the_grace_window(): - legacy = _legacy_session_keys(MASTER_KEY) - minted = mint_session_token(PRINCIPAL, legacy, NOW) + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_token(PRINCIPAL, legacy, NOW) assert isinstance(minted, MintedSessionToken) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) assert grace == legacy - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) assert isinstance(result, SessionBearerAdmitted) assert result.principal == PRINCIPAL def test_legacy_session_token_is_invalid_without_grace(): - minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + minted: Final = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) assert isinstance(minted, MintedSessionToken) - assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({})) is None - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW) + assert legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ(MappingProxyType({}))) is None + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW) assert isinstance(result, SessionBearerInvalid) assert result.expired is False def test_legacy_session_key_is_never_derived_for_configured_rs256_keys(): - active = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") - grace = legacy_session_keys_from_master_key( - MASTER_KEY, active, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"}) + active: Final = AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem()), kid="k") + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, active, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) ) assert grace is None def test_grace_is_disabled_under_fips(): - grace = legacy_session_keys_from_master_key( + grace: Final = legacy_session_keys_from_master_key( MASTER_KEY, KEYS, - environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"}), + environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true", FIPS_MODE_ENV_VAR: "true"})), ) assert grace is None def test_legacy_session_refresh_token_opens_during_grace(): - legacy = _legacy_session_keys(MASTER_KEY) - minted = mint_session_refresh_token(PRINCIPAL, legacy, NOW) + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_refresh_token(PRINCIPAL, legacy, NOW) assert isinstance(minted, MintedSessionToken) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - result = open_session_refresh_bearer( + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + result: Final = open_session_refresh_bearer( minted.token.get_secret_value(), KEYS, NOW, @@ -315,37 +321,22 @@ def test_legacy_session_refresh_token_opens_during_grace(): def test_legacy_session_refresh_token_is_invalid_without_grace(): - minted = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) + minted: Final = mint_session_refresh_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) assert isinstance(minted, MintedSessionToken) - result = open_session_refresh_bearer( + result: Final = open_session_refresh_bearer( minted.token.get_secret_value(), KEYS, NOW, expected_client_id="llm_client_abc" ) assert isinstance(result, SessionRefreshInvalid) def test_expired_token_under_the_primary_key_still_reports_expired_during_grace(): - legacy = _legacy_session_keys(MASTER_KEY) - minted = mint_session_token(PRINCIPAL, legacy, NOW) + legacy: Final = _legacy_session_keys(MASTER_KEY) + minted: Final = mint_session_token(PRINCIPAL, legacy, NOW) assert isinstance(minted, MintedSessionToken) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - later = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) + grace: Final = legacy_session_keys_from_master_key( + MASTER_KEY, KEYS, environ=_environ(MappingProxyType({LEGACY_KDF_GRACE_ENV_VAR: "true"})) + ) + later: Final = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + result: Final = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, later, legacy_keys=grace) assert isinstance(result, SessionBearerInvalid) assert result.expired is True - - -def test_fallback_tries_the_active_key_before_legacy(monkeypatch): - calls: list[object] = [] - real = session_credentials.open_session_token - - def spy(candidate, keys, now): - calls.append(keys) - return real(candidate, keys, now) - - monkeypatch.setattr(session_credentials, "open_session_token", spy) - grace = legacy_session_keys_from_master_key(MASTER_KEY, KEYS, environ=_environ({LEGACY_KDF_GRACE_ENV_VAR: "true"})) - minted = mint_session_token(PRINCIPAL, _legacy_session_keys(MASTER_KEY), NOW) - assert isinstance(minted, MintedSessionToken) - result = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", KEYS, NOW, legacy_keys=grace) - assert isinstance(result, SessionBearerAdmitted) - assert calls[0] is KEYS From a466d8e08f751f24f5f362c0f6f74c492f966eec Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:49:54 +0000 Subject: [PATCH 8/9] test(integration): cover legacy kdf refresh, revoke and introspect paths Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../integration/mcp/test_mcp_kdf_rotation.py | 267 +++++++++++++++++- 1 file changed, 259 insertions(+), 8 deletions(-) diff --git a/tests/integration/mcp/test_mcp_kdf_rotation.py b/tests/integration/mcp/test_mcp_kdf_rotation.py index b915759dbda..69d6dc12a23 100644 --- a/tests/integration/mcp/test_mcp_kdf_rotation.py +++ b/tests/integration/mcp/test_mcp_kdf_rotation.py @@ -17,23 +17,33 @@ import pytest from cryptography.hazmat.primitives.hashes import SHA256 from cryptography.hazmat.primitives.kdf.hkdf import HKDF from integration._support.client import Gateway -from integration._support.mcp import McpPeer, mcp_peer, register_mcp, tool_calls +from integration._support.mcp import McpPeer, mcp_peer, register_mcp from integration._support.oauth_server import oauth_server from integration._support.process import owned_proxy from pydantic import SecretStr +from typing_extensions import ReadOnly, TypedDict from litellm.proxy._experimental.mcp_server.outbound_credentials.envelope import ( EnvelopeKeys, + OpenedEnvelope, + OpenedRefreshEnvelope, + RefreshCredential, SealedEnvelope, UpstreamTokenGrant, key_hash_identity, mint_envelope, + mint_refresh_envelope, + open_envelope, + open_refresh_envelope, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( MintedSessionToken, + OpenedSessionToken, SessionKeys, SessionPrincipal, + mint_session_refresh_token, mint_session_token, + open_session_refresh_token, ) from litellm.proxy.utils import hash_token @@ -45,6 +55,30 @@ SESSION_SIGNING: Final = b"litellm-mcp-gateway:session-signing:" GRACE: Final = {"LITELLM_MCP_LEGACY_KDF_GRACE": "true"} +class _RefreshGrantForm(TypedDict): + grant_type: ReadOnly[str] + refresh_token: ReadOnly[str] + client_id: ReadOnly[str] + + +class _RevokeForm(TypedDict): + token: ReadOnly[str] + client_id: ReadOnly[str] + + +class _IntrospectForm(TypedDict): + token: ReadOnly[str] + + +class _ClientRegistration(TypedDict): + redirect_uris: ReadOnly[tuple[str, ...]] + client_name: ReadOnly[str] + + +class _ObjectPermission(TypedDict): + mcp_servers: ReadOnly[tuple[str, ...]] + + def _hkdf(master_key: str, info: bytes) -> str: return HKDF(algorithm=SHA256(), length=32, salt=None, info=info).derive(master_key.encode()).hex() @@ -86,14 +120,51 @@ def _envelope(identity_server: str, key: str, keys: EnvelopeKeys, upstream_token return sealed.token.get_secret_value() -def _session(user_id: str, keys: SessionKeys) -> str: +def _session(user_id: str, keys: SessionKeys, client_id: str = "dcr-integration") -> str: minted: Final = mint_session_token( - SessionPrincipal(user_id=user_id, client_id="dcr-integration"), keys, datetime.now(timezone.utc) + SessionPrincipal(user_id=user_id, client_id=client_id), keys, datetime.now(timezone.utc) ) assert isinstance(minted, MintedSessionToken), minted return minted.token.get_secret_value() +def _refresh_envelope(identity_server: str, key: str, keys: EnvelopeKeys, upstream_refresh: str) -> str: + sealed: Final = mint_refresh_envelope( + key_hash_identity(server_id=identity_server, key_hash=hash_token(key)), + RefreshCredential(refresh_token=SecretStr(upstream_refresh), scope="tools.read tools.call", expires_in=600), + keys, + datetime.now(timezone.utc), + ) + assert isinstance(sealed, SealedEnvelope), sealed + return sealed.token.get_secret_value() + + +def _session_refresh(user_id: str, client_id: str, keys: SessionKeys) -> str: + minted: Final = mint_session_refresh_token( + SessionPrincipal(user_id=user_id, client_id=client_id), keys, datetime.now(timezone.utc) + ) + assert isinstance(minted, MintedSessionToken), minted + return minted.token.get_secret_value() + + +def _upstream_refresh(auth) -> str: + issued: Final = auth.issue("authorization_code", "kdf-client", "integration-user", "tools.read tools.call") + token: Final = issued["refresh_token"] + assert isinstance(token, str) and token, issued + return token + + +def _register_gateway_client(client: httpx.Client) -> str: + registered: Final = client.post( + "/register", + json=_ClientRegistration(redirect_uris=("http://127.0.0.1/callback",), client_name="kdf-rotation"), + ) + assert registered.status_code == 201, registered.text + client_id: Final = registered.json().get("client_id") + assert isinstance(client_id, str) and client_id, registered.text + return client_id + + def _register_bridge(scenario, peer: McpPeer, alias: str, issuer: str) -> str: return register_mcp( scenario, @@ -137,7 +208,7 @@ def test_envelope_minted_under_hkdf_sha256_keys_is_admitted(gateway: Gateway) -> with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: alias: Final = "kdf" + uuid.uuid4().hex[:8] identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) - key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) upstream_token: Final = "up-" + uuid.uuid4().hex peer.drain() envelope: Final = _envelope(identity, key, _hkdf_envelope_keys(gateway.key), upstream_token) @@ -148,7 +219,7 @@ def test_envelope_minted_under_legacy_scrypt_keys_is_rejected_without_grace(gate with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: alias: Final = "kdf" + uuid.uuid4().hex[:8] identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) - key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) peer.drain() envelope: Final = _envelope(identity, key, _scrypt_envelope_keys(gateway.key), "up-" + uuid.uuid4().hex) _assert_rejected(_tools_list(gateway.client, f"/{alias}/mcp", envelope), peer) @@ -166,7 +237,7 @@ def test_envelope_of_either_kdf_is_admitted_during_the_legacy_grace_window( ): alias: Final = "kdf" + uuid.uuid4().hex[:8] identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) - key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) upstream_token: Final = "up-" + uuid.uuid4().hex peer.drain() keys: Final = _hkdf_envelope_keys(graced.key) if kdf == "hkdf" else _scrypt_envelope_keys(graced.key) @@ -183,7 +254,7 @@ def test_envelope_under_a_foreign_master_key_is_rejected_during_grace(gateway: G ): alias: Final = "kdf" + uuid.uuid4().hex[:8] identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) - key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) peer.drain() foreign: Final = "sk-foreign-" + uuid.uuid4().hex for keys in (_hkdf_envelope_keys(foreign), _scrypt_envelope_keys(foreign)): @@ -201,7 +272,7 @@ def test_grace_variable_set_to_anything_but_true_keeps_legacy_envelopes_rejected ): alias: Final = "kdf" + uuid.uuid4().hex[:8] identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) - key: Final = scenario.key(user_id=scenario.user(), object_permission={"mcp_servers": [identity]}) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) peer.drain() envelope: Final = _envelope(identity, key, _scrypt_envelope_keys(proxy.key), "up") _assert_rejected(_tools_list(proxy.client, f"/{alias}/mcp", envelope), peer) @@ -241,3 +312,183 @@ def test_tampered_session_token_is_rejected_during_grace(gateway: Gateway, tmp_p response: Final = _tools_list(graced.client, "/mcp", tampered) assert response.status_code == 401, response.text assert "invalid_token" in response.headers.get("www-authenticate", ""), response.headers + + +def test_bridge_refresh_envelope_minted_under_hkdf_sha256_keys_is_renewed(gateway: Gateway) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + envelope: Final = _refresh_envelope(identity, key, _hkdf_envelope_keys(gateway.key), _upstream_refresh(auth)) + response: Final = gateway.client.post( + f"/{alias}/token", + headers=(("x-litellm-api-key", key),), + data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=envelope, client_id="kdf-client"), + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + assert renewed["access_token"] != envelope, response.text + opened: Final = open_refresh_envelope( + renewed["refresh_token"], _hkdf_envelope_keys(gateway.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedRefreshEnvelope), opened + + +def test_bridge_refresh_envelope_minted_under_legacy_scrypt_keys_is_rejected_without_grace( + gateway: Gateway, +) -> None: + with mcp_peer() as peer, oauth_server() as auth, gateway.scenario() as scenario: + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + envelope: Final = _refresh_envelope(identity, key, _scrypt_envelope_keys(gateway.key), _upstream_refresh(auth)) + response: Final = gateway.client.post( + f"/{alias}/token", + headers=(("x-litellm-api-key", key),), + data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=envelope, client_id="kdf-client"), + ) + assert response.status_code == 400, response.text + assert response.json()["error"] == "invalid_grant", response.text + assert auth.token_requests() == (), "a rejected envelope must never reach the upstream token endpoint" + + +@pytest.mark.parametrize("kdf", ("hkdf", "scrypt")) +def test_bridge_refresh_envelope_of_either_kdf_is_renewed_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path, kdf: str +) -> None: + with ( + owned_proxy(gateway, tmp_path, GRACE) as graced, + mcp_peer() as peer, + oauth_server() as auth, + graced.scenario() as scenario, + ): + alias: Final = "kdf" + uuid.uuid4().hex[:8] + identity: Final = _register_bridge(scenario, peer, alias, auth.issuer) + key: Final = scenario.key(user_id=scenario.user(), object_permission=_ObjectPermission(mcp_servers=(identity,))) + peer.drain() + keys: Final = _hkdf_envelope_keys(graced.key) if kdf == "hkdf" else _scrypt_envelope_keys(graced.key) + envelope: Final = _refresh_envelope(identity, key, keys, _upstream_refresh(auth)) + response: Final = graced.client.post( + f"/{alias}/token", + headers=(("x-litellm-api-key", key),), + data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=envelope, client_id="kdf-client"), + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + opened: Final = open_refresh_envelope( + renewed["refresh_token"], _hkdf_envelope_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedRefreshEnvelope), opened + access: Final = open_envelope( + renewed["access_token"], _hkdf_envelope_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(access, OpenedEnvelope), access + _assert_admitted( + _tools_list(graced.client, f"/{alias}/mcp", renewed["access_token"]), + peer, + access.grant.access_token.get_secret_value(), + ) + + +def test_session_refresh_token_minted_under_hkdf_sha256_key_is_renewed(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + client_id: Final = _register_gateway_client(gateway.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _hkdf_session_keys(gateway.key)) + response: Final = gateway.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + assert renewed["refresh_token"] != refresh, response.text + opened: Final = open_session_refresh_token( + renewed["refresh_token"], _hkdf_session_keys(gateway.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedSessionToken), opened + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_is_rejected_without_grace(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + client_id: Final = _register_gateway_client(gateway.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(gateway.key)) + response: Final = gateway.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 400, response.text + assert response.json()["error"] == "invalid_grant", response.text + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_is_renewed_during_the_legacy_grace_window( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + client_id: Final = _register_gateway_client(graced.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(graced.key)) + response: Final = graced.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 200, response.text + renewed: Final = response.json() + opened: Final = open_session_refresh_token( + renewed["refresh_token"], _hkdf_session_keys(graced.key), datetime.now(timezone.utc) + ) + assert isinstance(opened, OpenedSessionToken), opened + admitted: Final = _tools_list(graced.client, "/mcp", renewed["access_token"]) + assert admitted.status_code == 200, admitted.text + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_is_revoked_during_grace( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + client_id: Final = _register_gateway_client(graced.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(graced.key)) + revoked: Final = graced.client.post("/revoke", data=_RevokeForm(token=refresh, client_id=client_id)) + assert revoked.status_code == 200, revoked.text + response: Final = graced.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 400, response.text + assert "invalid_grant" in response.text, response.text + + +def test_session_refresh_token_minted_under_legacy_scrypt_key_revokes_cleanly_without_grace( + gateway: Gateway, +) -> None: + with gateway.scenario() as scenario: + client_id: Final = _register_gateway_client(gateway.client) + refresh: Final = _session_refresh(scenario.user(), client_id, _scrypt_session_keys(gateway.key)) + revoked: Final = gateway.client.post("/revoke", data=_RevokeForm(token=refresh, client_id=client_id)) + assert revoked.status_code == 200, revoked.text + response: Final = gateway.client.post( + "/token", data=_RefreshGrantForm(grant_type="refresh_token", refresh_token=refresh, client_id=client_id) + ) + assert response.status_code == 400, response.text + assert "invalid_grant" in response.text, response.text + + +@pytest.mark.parametrize(("kdf", "active"), (("hkdf", True), ("scrypt", False))) +def test_session_access_token_introspects_by_its_kdf_without_grace(gateway: Gateway, kdf: str, active: bool) -> None: + with gateway.scenario() as scenario: + keys: Final = _hkdf_session_keys(gateway.key) if kdf == "hkdf" else _scrypt_session_keys(gateway.key) + key: Final = scenario.key(user_id=scenario.user()) + response: Final = gateway.client.post( + "/introspect", + headers=(("x-litellm-api-key", key),), + data=_IntrospectForm(token=_session(scenario.user(), keys)), + ) + assert response.status_code == 200, response.text + assert response.json()["active"] is active, response.text + + +def test_session_access_token_minted_under_legacy_scrypt_key_introspects_active_during_grace( + gateway: Gateway, tmp_path: Path +) -> None: + with owned_proxy(gateway, tmp_path, GRACE) as graced, graced.scenario() as scenario: + key: Final = scenario.key(user_id=scenario.user()) + response: Final = graced.client.post( + "/introspect", + headers=(("x-litellm-api-key", key),), + data=_IntrospectForm(token=_session(scenario.user(), _scrypt_session_keys(graced.key))), + ) + assert response.status_code == 200, response.text + assert response.json()["active"] is True, response.text From b403b04776e3ab72b2ff685a6d17d08de60435c0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 08:50:43 +0000 Subject: [PATCH 9/9] test(integration): type the upstream refresh helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/mcp/test_mcp_kdf_rotation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/mcp/test_mcp_kdf_rotation.py b/tests/integration/mcp/test_mcp_kdf_rotation.py index 69d6dc12a23..9fdce17d7d8 100644 --- a/tests/integration/mcp/test_mcp_kdf_rotation.py +++ b/tests/integration/mcp/test_mcp_kdf_rotation.py @@ -18,7 +18,7 @@ from cryptography.hazmat.primitives.hashes import SHA256 from cryptography.hazmat.primitives.kdf.hkdf import HKDF from integration._support.client import Gateway from integration._support.mcp import McpPeer, mcp_peer, register_mcp -from integration._support.oauth_server import oauth_server +from integration._support.oauth_server import AuthorizationServer, oauth_server from integration._support.process import owned_proxy from pydantic import SecretStr from typing_extensions import ReadOnly, TypedDict @@ -147,7 +147,7 @@ def _session_refresh(user_id: str, client_id: str, keys: SessionKeys) -> str: return minted.token.get_secret_value() -def _upstream_refresh(auth) -> str: +def _upstream_refresh(auth: AuthorizationServer) -> str: issued: Final = auth.issue("authorization_code", "kdf-client", "integration-user", "tools.read tools.call") token: Final = issued["refresh_token"] assert isinstance(token, str) and token, issued