From ef2ab74c7aecab4f970827ec21d740ed3aaa430b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:33:38 +0000 Subject: [PATCH] refactor(rust): back the vault secret manager with vaultrs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 91 +++- litellm-rust/Cargo.toml | 3 + .../crates/secrets-hashicorp/Cargo.toml | 4 +- .../secrets-hashicorp/src/cert_login.rs | 25 ++ .../crates/secrets-hashicorp/src/error.rs | 8 +- .../crates/secrets-hashicorp/src/lib.rs | 3 +- .../secrets-hashicorp/src/secret_manager.rs | 406 +++++++++--------- .../secrets-hashicorp/tests/secret_manager.rs | 293 +++++++++---- litellm-rust/crates/secrets/tests/handler.rs | 39 +- 9 files changed, 586 insertions(+), 286 deletions(-) create mode 100644 litellm-rust/crates/secrets-hashicorp/src/cert_login.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6f05bba9417..7b59533ad06 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2763,13 +2763,15 @@ dependencies = [ "litellm-core-utils", "litellm-secrets-types", "moka", - "reqwest 0.12.28", "rstest", + "rustify", + "rustify_derive", "serde", "serde_json", "tempfile", "thiserror 2.0.19", "tokio", + "vaultrs", "veil", "wiremock", ] @@ -3975,6 +3977,40 @@ dependencies = [ "semver", ] +[[package]] +name = "rustify" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4800ce4c1cc2fec12c559dae2ddbf0e17fcee7569b796e6d75898efef443368b" +dependencies = [ + "anyhow", + "async-trait", + "bytes", + "http 1.4.2", + "reqwest 0.13.5", + "rustify_derive", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 1.0.69", + "tracing", + "url", +] + +[[package]] +name = "rustify_derive" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78ea7fda74240f7410d0198b603a8a2f662acc7d76b6667a49f9b162cd8d9b4f" +dependencies = [ + "proc-macro2", + "quote", + "regex", + "serde_urlencoded", + "syn 1.0.109", + "synstructure 0.12.6", +] + [[package]] name = "rustix" version = "1.1.5" @@ -4502,6 +4538,17 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.119" @@ -4533,6 +4580,18 @@ dependencies = [ "futures-core", ] +[[package]] +name = "synstructure" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f36bdaa60a83aca3921b5259d5400cbf5e90fc51931376a9bd4a0eb79aa7210f" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", + "unicode-xid", +] + [[package]] name = "synstructure" version = "0.13.2" @@ -4947,6 +5006,7 @@ version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" dependencies = [ + "log", "pin-project-lite", "tracing-attributes", "tracing-core", @@ -5130,6 +5190,12 @@ version = "1.13.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + [[package]] name = "unicode_categories" version = "0.1.1" @@ -5189,6 +5255,25 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" +[[package]] +name = "vaultrs" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30ffcc0e81025065dda612ec1e26a3d81bb16ef3062354873d17a35965d68522" +dependencies = [ + "async-trait", + "derive_builder", + "http 1.4.2", + "reqwest 0.13.5", + "rustify", + "rustify_derive", + "serde", + "serde_json", + "thiserror 2.0.19", + "tracing", + "url", +] + [[package]] name = "veil" version = "0.3.0" @@ -5632,7 +5717,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.119", - "synstructure", + "synstructure 0.13.2", ] [[package]] @@ -5673,7 +5758,7 @@ dependencies = [ "proc-macro2", "quote", "syn 2.0.119", - "synstructure", + "synstructure 0.13.2", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 5030a94b140..edc9f5a91f6 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -51,6 +51,9 @@ reqwest = { version = "0.12", default-features = false, features = ["json", "mul rstest = "0.26.1" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustify = "=0.7.0" +rustify_derive = "=0.5.5" +vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] } rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 0656980a033..c646e02ef09 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -9,11 +9,13 @@ repository.workspace = true litellm-core-utils.workspace = true litellm-secrets-types.workspace = true moka.workspace = true -reqwest.workspace = true +rustify.workspace = true +rustify_derive.workspace = true serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true +vaultrs.workspace = true veil.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs new file mode 100644 index 00000000000..1eb09a89715 --- /dev/null +++ b/litellm-rust/crates/secrets-hashicorp/src/cert_login.rs @@ -0,0 +1,25 @@ +#[derive(Debug, rustify_derive::Endpoint)] +#[endpoint(path = "/auth/{self.mount}/login", method = "POST")] +pub struct CertLoginRequest { + #[endpoint(skip)] + pub mount: String, + #[endpoint(skip)] + #[allow(dead_code)] + pub name: Option, + #[endpoint(raw)] + body: Vec, +} + +impl CertLoginRequest { + pub fn new(name: Option) -> Self { + let body: Vec = match name.as_deref() { + Some(name) => serde_json::to_vec(&serde_json::json!({ "name": name })).unwrap(), + None => b"{}".to_vec(), + }; + Self { + mount: "cert".to_owned(), + name, + body, + } + } +} diff --git a/litellm-rust/crates/secrets-hashicorp/src/error.rs b/litellm-rust/crates/secrets-hashicorp/src/error.rs index 3f6a5b66475..e26033af085 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/error.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/error.rs @@ -4,12 +4,14 @@ pub enum Error { EnterpriseRequired, #[error("invalid secret name")] InvalidSecretName(#[from] litellm_secrets_types::Error), - #[error("HashiCorp Vault request failed")] - Request( + #[error("HashiCorp Vault client failed")] + Client( #[from] #[redact] - reqwest::Error, + vaultrs::error::ClientError, ), + #[error("HashiCorp Vault client settings are invalid: {message}")] + ClientSettings { message: String }, #[error("HashiCorp Vault TLS identity could not be configured for {path}: {message}")] TlsIdentity { path: std::path::PathBuf, diff --git a/litellm-rust/crates/secrets-hashicorp/src/lib.rs b/litellm-rust/crates/secrets-hashicorp/src/lib.rs index a561a58b59e..0c2b05647f8 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/lib.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/lib.rs @@ -1,9 +1,10 @@ #![forbid(unsafe_code)] +mod cert_login; mod config; mod error; pub mod secret_manager; pub use config::{AppRoleAuth, HashicorpVaultConfig, TlsCertAuth}; pub use error::Error; -pub use secret_manager::HashicorpVault; +pub use secret_manager::{HashicorpVault, SecretLocation}; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 52cc143c1e3..7937fa93f8f 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashMap, fmt, sync::Arc, time::{Duration, Instant}, @@ -9,26 +10,39 @@ use litellm_secrets_types::{ BaseSecretManager, SecretValue, async_rotate_secret, validate_secret_name, }; use moka::future::Cache; -use reqwest::{Client, Identity, RequestBuilder}; -use serde_json::{Value, json}; +use rustify::errors::ClientError as RustifyClientError; +use serde_json::Value; use tokio::sync::Mutex; +use vaultrs::{ + api, + auth::approle, + client::{Identity, VaultClient, VaultClientSettingsBuilder}, + error::ClientError, + kv2, +}; -use crate::{Error, HashicorpVaultConfig, TlsCertAuth}; +use crate::{Error, HashicorpVaultConfig, TlsCertAuth, cert_login::CertLoginRequest}; const CACHE_CAPACITY: u64 = 200; #[derive(Clone)] -struct CachedToken { - token: SecretValue, +struct CachedClient { + client: Arc, expires_at: Option, } +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SecretLocation { + pub namespace: Option, + pub mount: String, + pub path: String, +} + #[derive(Clone)] pub struct HashicorpVault { - client: Arc, config: HashicorpVaultConfig, cache: Cache, - auth_token: Arc>>, + auth_client: Arc>>, } impl fmt::Debug for HashicorpVault { @@ -45,17 +59,12 @@ impl HashicorpVault { environment: Arc, enterprise_enabled: bool, ) -> Result { - if !enterprise_enabled { - return Err(Error::EnterpriseRequired); - } let config: HashicorpVaultConfig = HashicorpVaultConfig::from_environment(environment.as_ref())?; - let client: Client = client_for_config(&config)?; - Self::with_client(client, config, enterprise_enabled) + Self::from_config(config, enterprise_enabled) } - pub fn with_client( - client: Client, + pub fn from_config( config: HashicorpVaultConfig, enterprise_enabled: bool, ) -> Result { @@ -67,47 +76,27 @@ impl HashicorpVault { .time_to_live(config.refresh_interval) .build(); Ok(Self { - client: Arc::new(client), config, cache, - auth_token: Arc::new(Mutex::new(None)), + auth_client: Arc::new(Mutex::new(None)), }) } - pub fn secret_url(&self, secret_name: &str) -> Result { + pub fn secret_location(&self, secret_name: &str) -> Result { validate_secret_name(secret_name).map_err(Error::InvalidSecretName)?; - let namespace: String = self - .config - .secret_namespace() - .map(|value| format!("{value}/")) - .unwrap_or_default(); - let path_prefix: String = self - .config - .path_prefix - .as_deref() - .map(|value| format!("{value}/")) - .unwrap_or_default(); - Ok(format!( - "{}/v1/{}{}/data/{}{}", - self.config.address, namespace, self.config.mount, path_prefix, secret_name - )) - } - - pub fn login_url(&self) -> Option { - self.config.approle.as_ref().map_or_else( - || { - self.config - .tls_cert - .as_ref() - .map(|_| format!("{}/v1/auth/cert/login", self.config.address)) - }, - |approle| { - Some(format!( - "{}/v1/auth/{}/login", - self.config.address, approle.mount_path - )) - }, - ) + let path: String = [ + self.config.path_prefix.clone(), + Some(secret_name.to_owned()), + ] + .into_iter() + .flatten() + .collect::>() + .join("/"); + Ok(SecretLocation { + namespace: self.config.secret_namespace().map(str::to_owned), + mount: self.config.mount.clone(), + path, + }) } pub fn config(&self) -> &HashicorpVaultConfig { @@ -115,40 +104,24 @@ impl HashicorpVault { } pub async fn async_read_secret(&self, secret_name: &str) -> Result, Error> { - let url: String = self.secret_url(secret_name)?; - if let Some(value) = self.cache.get(&url).await { + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + if let Some(value) = self.cache.get(&cache_key).await { return Ok(Some(value)); } - let token: SecretValue = self.vault_token().await?; - let response: reqwest::Response = self - .client - .get(&url) - .header("X-Vault-Token", token.expose()) - .send() - .await?; - if response.status() == reqwest::StatusCode::NOT_FOUND { - return Ok(None); - } - if !response.status().is_success() { - return Err(Error::Status { - status: response.status().as_u16(), - }); - } - let body = response.bytes().await?; - let body: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload)?; - let data: &Value = body - .get("data") - .and_then(Value::as_object) - .and_then(|value| value.get("data")) - .ok_or(Error::MalformedPayload)?; - let data: &serde_json::Map = - data.as_object().ok_or(Error::MalformedPayload)?; + let client: Arc = self.vault_client().await?; + let data: HashMap = + match kv2::read(client.as_ref(), &location.mount, &location.path).await { + Ok(data) => data, + Err(error) if api_status(&error) == Some(404) => return Ok(None), + Err(error) => return Err(map_api_error(error, ErrorContext::Read)), + }; let Some(value) = data.get("key") else { return Ok(None); }; let value: &str = value.as_str().ok_or(Error::NonStringValue)?; let value: SecretValue = SecretValue::new(value); - self.cache.insert(url, value.clone()).await; + self.cache.insert(cache_key, value.clone()).await; Ok(Some(value)) } @@ -158,47 +131,39 @@ impl HashicorpVault { value: SecretValue, description: Option<&str>, ) -> Result { - let url: String = self.secret_url(secret_name)?; - let data: Value = match description { - Some(description) => json!({"key": value.expose(), "description": description}), - None => json!({"key": value.expose()}), + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + let data: HashMap = match description { + Some(description) => [ + ("key".to_owned(), Value::String(value.expose().to_owned())), + ( + "description".to_owned(), + Value::String(description.to_owned()), + ), + ] + .into_iter() + .collect(), + None => [("key".to_owned(), Value::String(value.expose().to_owned()))] + .into_iter() + .collect(), }; - let token: SecretValue = self.vault_token().await?; - let response: reqwest::Response = self - .client - .post(&url) - .header("X-Vault-Token", token.expose()) - .json(&json!({"data": data})) - .send() - .await?; - if !response.status().is_success() { - return Err(Error::Status { - status: response.status().as_u16(), - }); - } - self.cache.invalidate(&url).await; - let body = response.bytes().await?; - if body.is_empty() { - return Ok(Value::Null); - } - serde_json::from_slice(&body).map_err(|_| Error::MalformedPayload) + let client: Arc = self.vault_client().await?; + let metadata = kv2::set(client.as_ref(), &location.mount, &location.path, &data) + .await + .map_err(|error| map_api_error(error, ErrorContext::Secret))?; + self.cache.invalidate(&cache_key).await; + serde_json::to_value(metadata) + .map_err(|source| Error::Client(ClientError::JsonParseError { source })) } pub async fn async_delete_secret(&self, secret_name: &str) -> Result<(), Error> { - let url: String = self.secret_url(secret_name)?; - let token: SecretValue = self.vault_token().await?; - let response: reqwest::Response = self - .client - .delete(&url) - .header("X-Vault-Token", token.expose()) - .send() - .await?; - if !response.status().is_success() { - return Err(Error::Status { - status: response.status().as_u16(), - }); - } - self.cache.invalidate(&url).await; + let location: SecretLocation = self.secret_location(secret_name)?; + let cache_key: String = cache_key(&location); + let client: Arc = self.vault_client().await?; + kv2::delete_latest(client.as_ref(), &location.mount, &location.path) + .await + .map_err(|error| map_api_error(error, ErrorContext::Secret))?; + self.cache.invalidate(&cache_key).await; Ok(()) } @@ -211,69 +176,76 @@ impl HashicorpVault { async_rotate_secret(self, current_name, new_name, value).await } - async fn vault_token(&self) -> Result { - let mut cached: tokio::sync::MutexGuard<'_, Option> = - self.auth_token.lock().await; + async fn vault_client(&self) -> Result, Error> { + let mut cached = self.auth_client.lock().await; if let Some(entry) = cached.as_ref() && entry .expires_at .is_none_or(|expires_at| expires_at > Instant::now()) { - return Ok(entry.token.clone()); + return Ok(entry.client.clone()); } - let Some(login_url) = self.login_url() else { - let Some(token) = self.config.token.clone() else { - return Err(Error::NoAuthConfigured); + + let (client, expires_at): (VaultClient, Option) = + match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) { + (Some(approle), _) => { + let login_client: VaultClient = + self.build_client(self.config.login_namespace(), "")?; + let auth = approle::login( + &login_client, + &approle.mount_path, + &approle.role_id, + approle.secret_id.expose(), + ) + .await + .map_err(|error| map_api_error(error, ErrorContext::Login))?; + ( + self.build_client(self.config.secret_namespace(), &auth.client_token)?, + token_expiry(auth.lease_duration), + ) + } + (None, Some(tls)) => { + let login_client: VaultClient = + self.build_client(self.config.login_namespace(), "")?; + let endpoint: CertLoginRequest = CertLoginRequest::new(tls.role.clone()); + let auth = api::auth(&login_client, endpoint) + .await + .map_err(|error| map_api_error(error, ErrorContext::Login))?; + ( + self.build_client(self.config.secret_namespace(), &auth.client_token)?, + token_expiry(auth.lease_duration), + ) + } + (None, None) => { + let token: SecretValue = + self.config.token.clone().ok_or(Error::NoAuthConfigured)?; + ( + self.build_client(self.config.secret_namespace(), token.expose())?, + None, + ) + } }; - return Ok(token); - }; - let body: Value = match (self.config.approle.as_ref(), self.config.tls_cert.as_ref()) { - (Some(approle), _) => { - json!({"role_id": approle.role_id, "secret_id": approle.secret_id.expose()}) - } - (None, Some(tls)) => tls - .role - .as_deref() - .map_or_else(|| json!({}), |role| json!({"name": role})), - (None, None) => { - let Some(token) = self.config.token.clone() else { - return Err(Error::NoAuthConfigured); - }; - return Ok(token); - } - }; - let request: RequestBuilder = with_namespace( - self.client.post(login_url).json(&body), - self.config.login_namespace(), - ); - let response: reqwest::Response = request.send().await?; - if !response.status().is_success() { - return Err(Error::LoginStatus { - status: response.status().as_u16(), - }); - } - let body = response.bytes().await?; - let payload: Value = serde_json::from_slice(&body).map_err(|_| Error::MalformedLogin)?; - let auth: &serde_json::Map = payload - .get("auth") - .and_then(Value::as_object) - .ok_or(Error::MalformedLogin)?; - let token: SecretValue = SecretValue::new( - auth.get("client_token") - .and_then(Value::as_str) - .ok_or(Error::MalformedLogin)?, - ); - let lease_duration: u64 = auth - .get("lease_duration") - .and_then(Value::as_u64) - .ok_or(Error::MalformedLogin)?; - let expires_at: Option = - (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)); - *cached = Some(CachedToken { - token: token.clone(), + let client: Arc = Arc::new(client); + *cached = Some(CachedClient { + client: client.clone(), expires_at, }); - Ok(token) + Ok(client) + } + + fn build_client(&self, namespace: Option<&str>, token: &str) -> Result { + let settings = VaultClientSettingsBuilder::default() + .address(&self.config.address) + .token(token.to_owned()) + .namespace(namespace.map(str::to_owned)) + .identity(identity_for(self.config.tls_cert.as_ref())?) + .ca_certs(Vec::new()) + .verify(true) + .build() + .map_err(|message| Error::ClientSettings { + message: message.to_string(), + })?; + VaultClient::new(settings).map_err(Error::Client) } } @@ -304,34 +276,84 @@ impl BaseSecretManager for HashicorpVault { } } -fn client_for_config(config: &HashicorpVaultConfig) -> Result { - let builder: reqwest::ClientBuilder = match config.tls_cert.as_ref() { - Some(tls) => Client::builder().identity(identity_for(tls)?), - None => Client::builder(), - }; - builder.build().map_err(Error::Request) +#[derive(Clone, Copy)] +enum ErrorContext { + Login, + Read, + Secret, } -fn with_namespace(request: RequestBuilder, namespace: Option<&str>) -> RequestBuilder { - match namespace { - Some(namespace) => request.header("X-Vault-Namespace", namespace), - None => request, +fn cache_key(location: &SecretLocation) -> String { + format!( + "{:?}/{}/{}", + location.namespace, location.mount, location.path + ) +} + +fn identity_for(tls: Option<&TlsCertAuth>) -> Result, Error> { + tls.map(|tls| { + let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + })?; + let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { + path: tls.key_path.clone(), + message: source.to_string(), + })?; + Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { + Error::TlsIdentity { + path: tls.cert_path.clone(), + message: source.to_string(), + } + }) + }) + .transpose() +} + +fn map_api_error(error: ClientError, context: ErrorContext) -> Error { + match error { + ClientError::APIError { code, .. } => match context { + ErrorContext::Login => Error::LoginStatus { status: code }, + ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, + }, + ClientError::JsonParseError { source } => match context { + ErrorContext::Login => Error::MalformedLogin, + ErrorContext::Read => Error::MalformedPayload, + ErrorContext::Secret => Error::Client(ClientError::JsonParseError { source }), + }, + ClientError::ResponseEmptyError | ClientError::ResponseDataEmptyError => { + malformed_response(context) + } + ClientError::RestClientError { source } => match source { + RustifyClientError::ServerResponseError { code, .. } => match context { + ErrorContext::Login => Error::LoginStatus { status: code }, + ErrorContext::Read | ErrorContext::Secret => Error::Status { status: code }, + }, + RustifyClientError::ResponseParseError { .. } => malformed_response(context), + source => Error::Client(ClientError::RestClientError { source }), + }, + error => Error::Client(error), } } -fn identity_for(tls: &TlsCertAuth) -> Result { - let cert: Vec = std::fs::read(&tls.cert_path).map_err(|source| Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - })?; - let key: Vec = std::fs::read(&tls.key_path).map_err(|source| Error::TlsIdentity { - path: tls.key_path.clone(), - message: source.to_string(), - })?; - Identity::from_pem(&[cert.as_slice(), key.as_slice()].concat()).map_err(|source| { - Error::TlsIdentity { - path: tls.cert_path.clone(), - message: source.to_string(), - } - }) +fn api_status(error: &ClientError) -> Option { + match error { + ClientError::APIError { code, .. } => Some(*code), + ClientError::RestClientError { + source: RustifyClientError::ServerResponseError { code, .. }, + } => Some(*code), + _ => None, + } +} + +fn malformed_response(context: ErrorContext) -> Error { + match context { + ErrorContext::Login => Error::MalformedLogin, + ErrorContext::Read => Error::MalformedPayload, + ErrorContext::Secret => Error::Client(ClientError::ResponseDataEmptyError), + } +} + +fn token_expiry(lease_duration: u64) -> Option { + (lease_duration > 0).then(|| Instant::now() + Duration::from_secs(lease_duration)) } diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index b7a81f49ad4..c52db46e41e 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -22,7 +22,51 @@ fn config(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVaultConfig } fn manager(server: &MockServer, values: &[(&str, &str)]) -> HashicorpVault { - HashicorpVault::with_client(reqwest::Client::new(), config(server, values), true).unwrap() + HashicorpVault::from_config(config(server, values), true).unwrap() +} + +fn auth_response(token: &str, lease_duration: u64) -> serde_json::Value { + json!({ + "auth": { + "client_token": token, + "accessor": "", + "policies": [], + "token_policies": [], + "metadata": null, + "lease_duration": lease_duration, + "renewable": false, + "entity_id": "", + "token_type": "service", + "orphan": false + }, + "lease_id": "", + "lease_duration": lease_duration, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }) +} + +fn read_response(data: serde_json::Value) -> serde_json::Value { + json!({ + "data": { + "data": data, + "metadata": { + "created_time": "", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 1 + } + }, + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }) } #[tokio::test] @@ -32,7 +76,7 @@ async fn token_reads_use_vault_headers_and_cache_values() { .and(path("/v1/secret/data/name")) .and(header("X-Vault-Token", "token")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(1) .mount(&server) @@ -48,6 +92,12 @@ async fn token_reads_use_vault_headers_and_cache_values() { .expose(), "value" ); + let requests = server.received_requests().await.unwrap(); + assert!( + requests + .iter() + .all(|request| !request.headers.contains_key("X-Vault-Namespace")) + ); assert_eq!( manager .async_read_secret("name") @@ -63,9 +113,10 @@ async fn token_reads_use_vault_headers_and_cache_values() { async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { let server: MockServer = MockServer::start().await; Mock::given(method("GET")) - .and(path("/v1/team-a/kv-prod/data/virtual-keys/name")) + .and(path("/v1/kv-prod/data/virtual-keys/name")) + .and(header("X-Vault-Namespace", "team-a")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(1) .mount(&server) @@ -80,10 +131,10 @@ async fn namespace_mount_and_prefix_are_sanitized_in_the_url() { ], ); - assert_eq!( - manager.secret_url("name").unwrap(), - format!("{}/v1/team-a/kv-prod/data/virtual-keys/name", server.uri()) - ); + let location = manager.secret_location("name").unwrap(); + assert_eq!(location.namespace.as_deref(), Some("team-a")); + assert_eq!(location.mount, "kv-prod"); + assert_eq!(location.path, "virtual-keys/name"); assert!(manager.async_read_secret("name").await.unwrap().is_some()); } @@ -96,12 +147,15 @@ fn trailing_address_slashes_are_removed() { }); let config: HashicorpVaultConfig = HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); - let manager: HashicorpVault = - HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + let manager: HashicorpVault = HashicorpVault::from_config(config, true).unwrap(); assert_eq!( - manager.secret_url("name").unwrap(), - "http://vault.test:8200/v1/secret/data/name" + manager.secret_location("name").unwrap(), + litellm_secrets_hashicorp::SecretLocation { + namespace: None, + mount: "secret".to_owned(), + path: "name".to_owned(), + } ); } @@ -127,21 +181,26 @@ async fn approle_login_uses_namespace_and_reuses_the_token() { .and(path("/v1/auth/custom-approle/login")) .and(header("X-Vault-Namespace", "login-root")) .and(body_json(json!({"role_id": "role", "secret_id": "secret"}))) - .respond_with(ResponseTemplate::new(200).set_body_json( - json!({"auth": {"client_token": "login-token", "lease_duration": 3600}}), - )) + .respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 3600))) .expect(1) .mount(&server) .await; Mock::given(method("GET")) - .and(path("/v1/secret-root/secret/data/name")) + .and(path("/v1/secret/data/name")) .and(header("X-Vault-Token", "login-token")) + .and(header("X-Vault-Namespace", "secret-root")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(1) .mount(&server) .await; + Mock::given(method("GET")) + .and(path("/v1/secret/data/name-2")) + .respond_with(ResponseTemplate::new(404).set_body_json(json!({"errors": ["missing"]}))) + .expect(1) + .mount(&server) + .await; let manager: HashicorpVault = manager( &server, &[ @@ -162,17 +221,13 @@ async fn approle_tokens_expire_after_the_vault_lease() { let server: MockServer = MockServer::start().await; Mock::given(method("POST")) .and(path("/v1/auth/approle/login")) - .respond_with( - ResponseTemplate::new(200).set_body_json( - json!({"auth": {"client_token": "login-token", "lease_duration": 1}}), - ), - ) + .respond_with(ResponseTemplate::new(200).set_body_json(auth_response("login-token", 1))) .expect(2) .mount(&server) .await; Mock::given(method("GET")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(2) .mount(&server) @@ -201,54 +256,98 @@ async fn tls_login_posts_the_role_and_uses_the_client_identity() { std::fs::write(&key_path, TEST_PRIVATE_KEY).unwrap(); Mock::given(method("POST")) .and(path("/v1/auth/cert/login")) - .and(body_json(json!({"name": "vault-role"}))) - .respond_with( - ResponseTemplate::new(200).set_body_json( - json!({"auth": {"client_token": "cert-token", "lease_duration": 0}}), - ), - ) - .expect(1) + .and(header("X-Vault-Namespace", "login-ns")) + .respond_with(ResponseTemplate::new(200).set_body_json(auth_response("cert-token", 0))) + .expect(2) .mount(&server) .await; Mock::given(method("GET")) .and(path("/v1/secret/data/name")) + .and(header("X-Vault-Token", "cert-token")) + .and(header("X-Vault-Namespace", "secret-ns")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) + .expect(2) .mount(&server) .await; - let manager: HashicorpVault = HashicorpVault::new( - { - let environment_values: HashMap = HashMap::from([ - ("HCP_VAULT_ADDR".to_owned(), server.uri()), - ( - "HCP_VAULT_CLIENT_CERT".to_owned(), - cert_path.to_str().unwrap().to_owned(), - ), - ( - "HCP_VAULT_CLIENT_KEY".to_owned(), - key_path.to_str().unwrap().to_owned(), - ), - ("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()), - ]); - Arc::new(move |name: &str| environment_values.get(name).cloned()) - }, - true, - ) - .unwrap(); - - assert!(manager.async_read_secret("name").await.unwrap().is_some()); - assert_eq!( - manager.login_url().as_deref(), - Some(format!("{}/v1/auth/cert/login", server.uri()).as_str()) + let role_values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + key_path.to_str().unwrap().to_owned(), + ), + ("HCP_VAULT_CERT_ROLE".to_owned(), "vault-role".to_owned()), + ( + "HCP_VAULT_LOGIN_NAMESPACE".to_owned(), + "login-ns".to_owned(), + ), + ( + "HCP_VAULT_SECRET_NAMESPACE".to_owned(), + "secret-ns".to_owned(), + ), + ]); + let role_environment: Arc = + Arc::new(move |name: &str| role_values.get(name).cloned()); + let role_manager: HashicorpVault = HashicorpVault::new(role_environment, true).unwrap(); + assert!( + role_manager + .async_read_secret("name") + .await + .unwrap() + .is_some() ); + + let no_role_values: HashMap = HashMap::from([ + ("HCP_VAULT_ADDR".to_owned(), server.uri()), + ( + "HCP_VAULT_CLIENT_CERT".to_owned(), + cert_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_CLIENT_KEY".to_owned(), + key_path.to_str().unwrap().to_owned(), + ), + ( + "HCP_VAULT_LOGIN_NAMESPACE".to_owned(), + "login-ns".to_owned(), + ), + ( + "HCP_VAULT_SECRET_NAMESPACE".to_owned(), + "secret-ns".to_owned(), + ), + ]); + let no_role_environment: Arc = + Arc::new(move |name: &str| no_role_values.get(name).cloned()); + let no_role_manager: HashicorpVault = HashicorpVault::new(no_role_environment, true).unwrap(); + assert!( + no_role_manager + .async_read_secret("name") + .await + .unwrap() + .is_some() + ); + let login_bodies: Vec = server + .received_requests() + .await + .unwrap() + .iter() + .filter(|request| request.method.as_str() == "POST") + .map(|request| serde_json::from_slice(&request.body).unwrap()) + .collect(); + assert!(login_bodies.contains(&json!({"name": "vault-role"}))); + assert!(login_bodies.contains(&json!({}))); } #[rstest::rstest] -#[case::missing(404, json!({}), 0)] -#[case::malformed(200, json!({}), 1)] -#[case::missing_key(200, json!({"data": {"data": {}}}), 0)] -#[case::non_string(200, json!({"data": {"data": {"key": 1}}}), 2)] +#[case::missing(404, json!({"errors": ["missing"]}), 0)] +#[case::malformed(200, json!({"data": "invalid"}), 1)] +#[case::missing_key(200, json!({}), 0)] +#[case::non_string(200, json!({"key": 1}), 2)] #[tokio::test] async fn read_responses_distinguish_absence_and_malformed_payloads( #[case] status: u16, @@ -257,7 +356,13 @@ async fn read_responses_distinguish_absence_and_malformed_payloads( ) { let server: MockServer = MockServer::start().await; Mock::given(method("GET")) - .respond_with(ResponseTemplate::new(status).set_body_json(body)) + .respond_with(ResponseTemplate::new(status).set_body_json( + if status == 200 && expected != 1 { + read_response(body) + } else { + body + }, + )) .expect(1) .mount(&server) .await; @@ -279,7 +384,7 @@ async fn write_and_delete_invalidate_the_read_cache() { Mock::given(method("GET")) .and(path("/v1/secret/data/name")) .respond_with( - ResponseTemplate::new(200).set_body_json(json!({"data": {"data": {"key": "value"}}})), + ResponseTemplate::new(200).set_body_json(read_response(json!({"key": "value"}))), ) .expect(2) .mount(&server) @@ -289,7 +394,21 @@ async fn write_and_delete_invalidate_the_read_cache() { .and(body_json( json!({"data": {"key": "updated", "description": "description"}}), )) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({"data": {"version": 2}}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "data": { + "created_time": "", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 2 + }, + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }))) .expect(1) .mount(&server) .await; @@ -331,12 +450,9 @@ async fn no_auth_and_invalid_names_fail_without_requests() { #[tokio::test] async fn debug_output_redacts_authentication_values() { let server: MockServer = MockServer::start().await; - let manager: HashicorpVault = HashicorpVault::with_client( - reqwest::Client::new(), - config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), - true, - ) - .unwrap(); + let manager: HashicorpVault = + HashicorpVault::from_config(config(&server, &[("HCP_VAULT_TOKEN", "token-value")]), true) + .unwrap(); let debug: String = format!("{manager:?}"); assert!(!debug.contains("token-value")); assert!(!debug.contains("secret-id")); @@ -365,13 +481,35 @@ fn configuration_matches_python_parity_fixture() { Arc::new(move |name: &str| values.get(name).cloned()); let config: HashicorpVaultConfig = HashicorpVaultConfig::from_environment(environment.as_ref()).unwrap(); - let manager: HashicorpVault = - HashicorpVault::with_client(reqwest::Client::new(), config, true).unwrap(); + let manager: HashicorpVault = HashicorpVault::from_config(config.clone(), true).unwrap(); + let location = manager.secret_location(&case.secret_name).unwrap(); + let namespace = location + .namespace + .as_deref() + .map(|namespace| format!("{namespace}/")) + .unwrap_or_default(); assert_eq!( - manager.secret_url(&case.secret_name).unwrap(), + format!( + "{}/v1/{}{}/data/{}", + config.address, namespace, location.mount, location.path + ), case.expected_secret_url ); - assert_eq!(manager.login_url(), case.expected_login_url); + let login_url = config.approle.as_ref().map_or_else( + || { + config + .tls_cert + .as_ref() + .map(|_| format!("{}/v1/auth/cert/login", config.address)) + }, + |approle| { + Some(format!( + "{}/v1/auth/{}/login", + config.address, approle.mount_path + )) + }, + ); + assert_eq!(login_url, case.expected_login_url); assert_eq!( manager.config().login_namespace(), case.expected_login_namespace.as_deref() @@ -391,8 +529,15 @@ async fn live_vault_round_trip() { let manager: HashicorpVault = HashicorpVault::new(environment, true).unwrap(); let name: String = std::env::var("LITELLM_VAULT_LIVE_SECRET_NAME").unwrap(); let value: SecretValue = SecretValue::new("native-live-value"); - let url: String = manager.secret_url(&name).unwrap(); - println!("native provenance: {} {}", module_path!(), url); + let location = manager.secret_location(&name).unwrap(); + println!( + "native provenance: {} vaultrs {} {:?} {} {}", + module_path!(), + manager.config().address, + location.namespace, + location.mount, + location.path + ); manager .async_write_secret(&name, value.clone(), None) .await diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index 58941171fd9..c169d411f8f 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -124,10 +124,24 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let found_server = MockServer::start().await; Mock::given(method("GET")) .and(path("/v1/secret/data/KEY")) - .respond_with( - ResponseTemplate::new(200) - .set_body_json(serde_json::json!({"data": {"data": {"key": "remote"}}})), - ) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "data": { + "data": {"key": "remote"}, + "metadata": { + "created_time": "", + "deletion_time": "", + "custom_metadata": null, + "destroyed": false, + "version": 1 + } + }, + "lease_id": "", + "lease_duration": 0, + "renewable": false, + "request_id": "", + "warnings": null, + "wrap_info": null + }))) .mount(&found_server) .await; let found_environment: Arc = Arc::new({ @@ -139,8 +153,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { } }); let found_config = HashicorpVaultConfig::from_environment(found_environment.as_ref()).unwrap(); - let found_manager = - HashicorpVault::with_client(reqwest::Client::new(), found_config, true).unwrap(); + let found_manager = HashicorpVault::from_config(found_config, true).unwrap(); let found_resolver = SecretResolver::new( Arc::new(SecretManagerState::new( SecretManager::HashicorpVault(found_manager), @@ -164,7 +177,9 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let missing_server = MockServer::start().await; Mock::given(method("GET")) - .respond_with(ResponseTemplate::new(404)) + .respond_with( + ResponseTemplate::new(404).set_body_json(serde_json::json!({"errors": ["missing"]})), + ) .mount(&missing_server) .await; let missing_environment: Arc = Arc::new({ @@ -177,8 +192,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }); let missing_config = HashicorpVaultConfig::from_environment(missing_environment.as_ref()).unwrap(); - let missing_manager = - HashicorpVault::with_client(reqwest::Client::new(), missing_config, true).unwrap(); + let missing_manager = HashicorpVault::from_config(missing_config, true).unwrap(); let missing_state = SecretManagerState::new( SecretManager::HashicorpVault(missing_manager), KeyManagementSettings { @@ -198,7 +212,9 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { let failed_server = MockServer::start().await; Mock::given(method("GET")) - .respond_with(ResponseTemplate::new(500)) + .respond_with( + ResponseTemplate::new(500).set_body_json(serde_json::json!({"errors": ["failed"]})), + ) .mount(&failed_server) .await; let failed_environment: Arc = Arc::new({ @@ -211,8 +227,7 @@ async fn hashicorp_handler_resolves_found_missing_and_failed_values() { }); let failed_config = HashicorpVaultConfig::from_environment(failed_environment.as_ref()).unwrap(); - let failed_manager = - HashicorpVault::with_client(reqwest::Client::new(), failed_config, true).unwrap(); + let failed_manager = HashicorpVault::from_config(failed_config, true).unwrap(); let failed_state = SecretManagerState::new( SecretManager::HashicorpVault(failed_manager), KeyManagementSettings {