From 1669213eb552fa69ad542c1673c7b7588e5f8ae0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sat, 19 Sep 2026 08:32:46 -0700 Subject: [PATCH] fix(rust): read OCR secrets from the process environment and decline when a secret manager is readable The OCR route called back into Python's get_secret_str for every env fallback. With no secret manager configured that is os.environ behind a GIL hop, and with one configured it blocked a tokio worker on vault I/O and also sent the Azure and GCP identity variables, which Python reads with os.getenv, to the vault. The other Rust routes already read the process environment. Read the process environment here too. When litellm would read secrets from a secret manager, decline the Rust route so the Python route serves the call with the vault-backed keys --- .../crates/python-bridge/python_settings.json | 3 + .../python-bridge/src/python_settings.rs | 74 +++---------------- .../python-bridge/src/routes/ocr/mod.rs | 73 ++++++++++++++++-- litellm/rust_bridge/settings.py | 11 ++- .../test_litellm/rust_bridge/test_settings.py | 25 ++++--- 5 files changed, 100 insertions(+), 86 deletions(-) diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 4ad3edf682d..0af55083bef 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -19,5 +19,8 @@ "vertex_project", "vertex_location", "enable_azure_ad_token_refresh" + ], + "secret_manager": [ + "readable" ] } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 83db4f02500..7ac23a05542 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -1,4 +1,3 @@ -use litellm_core_utils::settings::Lookup; use pyo3::prelude::*; const MODULE: &str = "litellm.rust_bridge.settings"; @@ -8,17 +7,24 @@ pub(crate) enum PythonSettings { Http, UrlPolicy, ProviderDefaults, + SecretManager, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 3] = [Self::Http, Self::UrlPolicy, Self::ProviderDefaults]; + pub(crate) const ALL: [Self; 4] = [ + Self::Http, + Self::UrlPolicy, + Self::ProviderDefaults, + Self::SecretManager, + ]; pub(crate) fn name(self) -> &'static str { match self { Self::Http => "http_settings", Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", + Self::SecretManager => "secret_manager", } } @@ -32,23 +38,6 @@ impl PythonSettings { } } -pub(crate) struct PythonSecrets; - -impl Lookup for PythonSecrets { - fn get(&self, name: &str) -> Option { - Python::attach(|py| { - py.import(MODULE) - .and_then(|module| module.getattr("secret")?.call1((name,))) - .and_then(|value| value.extract::>()) - .unwrap_or_else(|error| { - let _ = - PythonSettings::warn(py, &format!("reading secret {name} failed: {error}")); - None - }) - }) - } -} - #[cfg(test)] pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); @@ -56,10 +45,9 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); mod tests { use std::{collections::BTreeSet, ffi::CString}; - use litellm_core_utils::settings::Lookup; use pyo3::{prelude::*, types::PyDict}; - use super::{CONTRACT, PythonSecrets, PythonSettings}; + use super::{CONTRACT, PythonSettings}; #[test] fn every_settings_group_is_in_the_python_contract() { @@ -83,48 +71,4 @@ mod tests { assert_eq!(read, declared); }); } - - #[test] - fn secrets_come_from_the_python_secret_reader_and_a_failed_read_is_unset() { - Python::initialize(); - Python::attach(|py| { - py.run( - c" -import sys -import types -settings = types.ModuleType('litellm.rust_bridge.settings') -settings.warnings = [] -def secret(name): - if name == 'BROKEN': - raise RuntimeError('vault down') - return {'MISTRAL_API_KEY': 'from-vault'}.get(name) -settings.secret = secret -settings.warn = settings.warnings.append -sys.modules.setdefault('litellm', types.ModuleType('litellm')) -sys.modules.setdefault('litellm.rust_bridge', types.ModuleType('litellm.rust_bridge')) -sys.modules['litellm.rust_bridge.settings'] = settings -", - None, - None, - ) - .unwrap(); - }); - assert_eq!( - PythonSecrets.get("MISTRAL_API_KEY").as_deref(), - Some("from-vault") - ); - assert_eq!(PythonSecrets.get("ABSENT"), None); - assert_eq!(PythonSecrets.get("BROKEN"), None); - Python::attach(|py| { - let warnings: Vec = py - .import("litellm.rust_bridge.settings") - .unwrap() - .getattr("warnings") - .unwrap() - .extract() - .unwrap(); - assert_eq!(warnings.len(), 1); - assert!(warnings[0].contains("BROKEN") && warnings[0].contains("vault down")); - }); - } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index 785f6e48e13..9be3171f70b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -10,17 +10,16 @@ use litellm_auth_gcp::VertexAuth; use litellm_callbacks_legacy::{LegacySurface, PublicCall, run_legacy_call}; use litellm_core::ocr::route::ocr_machine; use litellm_core_utils::settings::ProcessEnvironment; -use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; +use litellm_llms::base_llm::ocr::{ + handler::OcrClient, + settings::{OcrSettings, Secrets}, +}; use pyo3::{ prelude::*, types::{PyDict, PyTuple}, }; -use crate::{ - errors::RustBridgeDeclined, - http, - python_settings::{PythonSecrets, PythonSettings}, -}; +use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings}; const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -42,6 +41,7 @@ fn run_ocr( kwargs: Bound<'_, PyDict>, asynchronous: bool, ) -> PyResult> { + let secrets = process_environment_secrets(&PythonSettings::SecretManager.read(py)?)?; let config = http::call_config(py, &kwargs, asynchronous)?; let client = OcrClient::new( http::pool(), @@ -49,7 +49,7 @@ fn run_ocr( http::url_policy(py)?, VERTEX_AUTH.clone(), ocr_settings(py)?, - Arc::new(PythonSecrets), + secrets, ) .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; run_legacy_call( @@ -62,6 +62,21 @@ fn run_ocr( ) } +#[derive(FromPyObject)] +struct PythonSecretManager { + readable: bool, +} + +fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult { + let manager: PythonSecretManager = secret_manager.extract()?; + if manager.readable { + return Err(RustBridgeDeclined::new_err( + "a readable secret manager is configured and the Rust route only reads the process environment", + )); + } + Ok(Arc::new(ProcessEnvironment)) +} + #[derive(FromPyObject)] struct PythonProviderDefaults { vertex_project: Option, @@ -105,3 +120,47 @@ pub(crate) fn aocr( ) -> PyResult> { run_ocr(py, request, args, kwargs, true) } + +#[cfg(test)] +mod tests { + use pyo3::{prelude::*, types::PyDict}; + + use super::process_environment_secrets; + use crate::errors::RustBridgeDeclined; + + fn secret_manager<'py>(py: Python<'py>, readable: bool) -> Bound<'py, PyAny> { + let locals = PyDict::new(py); + locals.set_item("readable", readable).unwrap(); + py.run( + c"import types\nmanager = types.SimpleNamespace(readable=readable)", + Some(&locals), + Some(&locals), + ) + .unwrap(); + locals.get_item("manager").unwrap().unwrap() + } + + #[test] + fn a_readable_secret_manager_sends_the_call_back_to_python() { + Python::initialize(); + Python::attach(|py| { + let declined = process_environment_secrets(&secret_manager(py, true)) + .err() + .expect("the Rust route declines"); + assert!(declined.is_instance_of::(py)); + }); + } + + #[test] + fn without_a_readable_secret_manager_secrets_are_the_process_environment() { + Python::initialize(); + Python::attach(|py| { + let secrets = process_environment_secrets(&secret_manager(py, false)).unwrap(); + assert_eq!( + secrets.get("LITELLM_RUST_BRIDGE_UNSET_VARIABLE_FOR_TEST"), + None + ); + assert_eq!(secrets.get("PATH"), std::env::var("PATH").ok()); + }); + } +} diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 037d6d9bd27..86450ffbb6b 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -31,16 +31,21 @@ class ProviderDefaults: enable_azure_ad_token_refresh: bool | None +@dataclass(frozen=True, slots=True) +class SecretManager: + readable: bool + + def warn(message: str) -> None: from litellm._logging import verbose_logger verbose_logger.warning("%s", message) -def secret(name: str) -> str | None: - from litellm.secret_managers.main import get_secret_str +def secret_manager() -> SecretManager: + from litellm.secret_managers.main import _should_read_secret_from_secret_manager - return get_secret_str(name) + return SecretManager(readable=_should_read_secret_from_secret_manager()) def provider_defaults() -> ProviderDefaults: diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 44c5ec42b36..6b78ddad44b 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -11,6 +11,7 @@ import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager from litellm.llms.custom_httpx.http_handler import default_user_agent from litellm.rust_bridge import settings +from litellm.secret_managers.main import get_secret_str from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" @@ -23,6 +24,7 @@ def test_the_rust_contract_matches_the_returned_fields() -> None: "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], "provider_defaults": [field.name for field in dataclasses.fields(settings.provider_defaults())], + "secret_manager": [field.name for field in dataclasses.fields(settings.secret_manager())], } @@ -101,25 +103,26 @@ class _VaultSecrets(CustomSecretManager): return self.secrets.get(secret_name) -def test_secret_prefers_the_secret_manager_and_falls_back_to_the_environment_on_a_miss( - monkeypatch: pytest.MonkeyPatch, +@pytest.mark.parametrize( + ("access_mode", "readable"), + [("read_only", True), ("read_and_write", True), ("write_only", False)], +) +def test_secret_manager_is_readable_only_when_litellm_would_read_secrets_from_it( + monkeypatch: pytest.MonkeyPatch, access_mode: str, readable: bool ) -> None: - monkeypatch.setenv("MISTRAL_API_KEY", "stale-env-key") - monkeypatch.setenv("REDUCTO_API_KEY", "env-only-key") + monkeypatch.setenv("MISTRAL_API_KEY", "env-key") monkeypatch.setattr(litellm, "secret_manager_client", _VaultSecrets({"MISTRAL_API_KEY": "vault-key"})) monkeypatch.setattr(litellm, "_key_management_system", KeyManagementSystem.CUSTOM) - monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode="read_only")) + monkeypatch.setattr(litellm, "_key_management_settings", KeyManagementSettings(access_mode=access_mode)) - assert settings.secret("MISTRAL_API_KEY") == "vault-key" - assert settings.secret("REDUCTO_API_KEY") == "env-only-key" - assert settings.secret("ABSENT_KEY") is None + assert settings.secret_manager() == settings.SecretManager(readable=readable) + assert (get_secret_str("MISTRAL_API_KEY") == "vault-key") is readable -def test_secret_reads_the_environment_without_a_secret_manager(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("MISTRAL_API_KEY", "env-key") +def test_secret_manager_is_not_readable_without_a_client(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "secret_manager_client", None) - assert settings.secret("MISTRAL_API_KEY") == "env-key" + assert settings.secret_manager() == settings.SecretManager(readable=False) def test_provider_defaults_read_the_litellm_globals(monkeypatch: pytest.MonkeyPatch) -> None: