mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
refactor(rust): define consistent secret lookup contracts
This commit is contained in:
Yujong Lee
parent
82bc67b122
commit
e3f69fc4d8
16 changed files with 522 additions and 775 deletions
254
litellm-rust/Cargo.lock
generated
254
litellm-rust/Cargo.lock
generated
|
|
@ -1515,15 +1515,6 @@ dependencies = [
|
|||
"version_check",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getopts"
|
||||
version = "0.2.24"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cfe4fbac503b8d1f88e6676011885f34b7174f46e59956bba534ba83abded4df"
|
||||
dependencies = [
|
||||
"unicode-width",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "getrandom"
|
||||
version = "0.2.17"
|
||||
|
|
@ -2227,27 +2218,6 @@ version = "2.12.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2"
|
||||
|
||||
[[package]]
|
||||
name = "is-macro"
|
||||
version = "0.3.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8267aa6001e25494f3015f9663bbd88a18240c74483afa5f0934a1b3e4c388e9"
|
||||
dependencies = [
|
||||
"heck",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.11.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b1c173a5686ce8bfa551b3563d0c2170bf24ca44da99c7ca4bfdab5418c3fe57"
|
||||
dependencies = [
|
||||
"either",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "itertools"
|
||||
version = "0.13.0"
|
||||
|
|
@ -2394,12 +2364,6 @@ dependencies = [
|
|||
"zeroize",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "lalrpop-util"
|
||||
version = "0.20.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "507460a910eb7b32ee961886ff48539633b788a36b65692b95f225b844c82553"
|
||||
|
||||
[[package]]
|
||||
name = "lazy_static"
|
||||
version = "1.5.0"
|
||||
|
|
@ -2716,14 +2680,12 @@ dependencies = [
|
|||
"moka",
|
||||
"reqwest 0.12.28",
|
||||
"rstest",
|
||||
"rustpython-parser",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"strum",
|
||||
"tempfile",
|
||||
"thiserror 2.0.19",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"wiremock",
|
||||
]
|
||||
|
||||
|
|
@ -2809,7 +2771,7 @@ dependencies = [
|
|||
"base64 0.22.1",
|
||||
"rand 0.8.7",
|
||||
"rstest",
|
||||
"rustc-hash 2.1.3",
|
||||
"rustc-hash",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.19",
|
||||
|
|
@ -2985,16 +2947,6 @@ dependencies = [
|
|||
"minimal-lexical",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-bigint"
|
||||
version = "0.4.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367"
|
||||
dependencies = [
|
||||
"num-integer",
|
||||
"num-traits",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num-bigint"
|
||||
version = "0.5.1"
|
||||
|
|
@ -3178,44 +3130,6 @@ version = "2.3.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220"
|
||||
|
||||
[[package]]
|
||||
name = "phf"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1fd6780a80ae0c52cc120a26a1a42c1ae51b247a253e4e06113d23d2c2edd078"
|
||||
dependencies = [
|
||||
"phf_shared",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "phf_codegen"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "aef8048c789fa5e851558d709946d6d79a8ff88c0440c587967f8e94bfb1216a"
|
||||
dependencies = [
|
||||
"phf_generator",
|
||||
"phf_shared",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "phf_generator"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3c80231409c20246a13fddb31776fb942c38553c51e871f8cbd687a4cfb5843d"
|
||||
dependencies = [
|
||||
"phf_shared",
|
||||
"rand 0.8.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "phf_shared"
|
||||
version = "0.11.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "67eabc2ef2a60eb7faa00097bd1ffdb5bd28e62bf39990626a582201b7a754e5"
|
||||
dependencies = [
|
||||
"siphasher",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pin-project"
|
||||
version = "1.1.13"
|
||||
|
|
@ -3488,7 +3402,7 @@ dependencies = [
|
|||
"pin-project-lite",
|
||||
"quinn-proto",
|
||||
"quinn-udp",
|
||||
"rustc-hash 2.1.3",
|
||||
"rustc-hash",
|
||||
"rustls 0.23.42",
|
||||
"socket2 0.6.5",
|
||||
"thiserror 2.0.19",
|
||||
|
|
@ -3510,7 +3424,7 @@ dependencies = [
|
|||
"rand 0.10.2",
|
||||
"rand_pcg",
|
||||
"ring",
|
||||
"rustc-hash 2.1.3",
|
||||
"rustc-hash",
|
||||
"rustls 0.23.42",
|
||||
"rustls-pki-types",
|
||||
"slab",
|
||||
|
|
@ -3689,7 +3603,7 @@ dependencies = [
|
|||
"arcstr",
|
||||
"combine",
|
||||
"itoa",
|
||||
"num-bigint 0.5.1",
|
||||
"num-bigint",
|
||||
"percent-encoding",
|
||||
"ryu",
|
||||
"sha1_smol",
|
||||
|
|
@ -3918,12 +3832,6 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2"
|
||||
|
||||
[[package]]
|
||||
name = "rustc-hash"
|
||||
version = "2.1.3"
|
||||
|
|
@ -4051,63 +3959,6 @@ dependencies = [
|
|||
"untrusted",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustpython-ast"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4cdaf8ee5c1473b993b398c174641d3aa9da847af36e8d5eb8291930b72f31a5"
|
||||
dependencies = [
|
||||
"is-macro",
|
||||
"num-bigint 0.4.8",
|
||||
"rustpython-parser-core",
|
||||
"static_assertions",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustpython-parser"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "868f724daac0caf9bd36d38caf45819905193a901e8f1c983345a68e18fb2abb"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"is-macro",
|
||||
"itertools 0.11.0",
|
||||
"lalrpop-util",
|
||||
"log",
|
||||
"num-bigint 0.4.8",
|
||||
"num-traits",
|
||||
"phf",
|
||||
"phf_codegen",
|
||||
"rustc-hash 1.1.0",
|
||||
"rustpython-ast",
|
||||
"rustpython-parser-core",
|
||||
"tiny-keccak",
|
||||
"unic-emoji-char",
|
||||
"unic-ucd-ident",
|
||||
"unicode_names2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustpython-parser-core"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4b6c12fa273825edc7bccd9a734f0ad5ba4b8a2f4da5ff7efe946f066d0f4ad"
|
||||
dependencies = [
|
||||
"is-macro",
|
||||
"memchr",
|
||||
"rustpython-parser-vendored",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustpython-parser-vendored"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "04fcea49a4630a3a5d940f4d514dc4f575ed63c14c3e3ed07146634aed7f67a6"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
"once_cell",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rustversion"
|
||||
version = "1.0.23"
|
||||
|
|
@ -4412,12 +4263,6 @@ version = "0.1.5"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e"
|
||||
|
||||
[[package]]
|
||||
name = "siphasher"
|
||||
version = "1.0.3"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649"
|
||||
|
||||
[[package]]
|
||||
name = "slab"
|
||||
version = "0.4.12"
|
||||
|
|
@ -4648,7 +4493,7 @@ dependencies = [
|
|||
"fancy-regex 0.17.0",
|
||||
"lazy_static",
|
||||
"regex",
|
||||
"rustc-hash 2.1.3",
|
||||
"rustc-hash",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4681,15 +4526,6 @@ dependencies = [
|
|||
"time-core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tiny-keccak"
|
||||
version = "2.0.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2c9d3793400a45f954c52e73d068316d76b6f4e36977e3fcebb13a2721e80237"
|
||||
dependencies = [
|
||||
"crunchy",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tinystr"
|
||||
version = "0.8.3"
|
||||
|
|
@ -5124,58 +4960,6 @@ version = "0.1.4"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94"
|
||||
|
||||
[[package]]
|
||||
name = "unic-char-property"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a8c57a407d9b6fa02b4795eb81c5b6652060a15a7903ea981f3d723e6c0be221"
|
||||
dependencies = [
|
||||
"unic-char-range",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unic-char-range"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0398022d5f700414f6b899e10b8348231abf9173fa93144cbc1a43b9793c1fbc"
|
||||
|
||||
[[package]]
|
||||
name = "unic-common"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "80d7ff825a6a654ee85a63e80f92f054f904f21e7d12da4e22f9834a4aaa35bc"
|
||||
|
||||
[[package]]
|
||||
name = "unic-emoji-char"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0b07221e68897210270a38bde4babb655869637af0f69407f96053a34f76494d"
|
||||
dependencies = [
|
||||
"unic-char-property",
|
||||
"unic-char-range",
|
||||
"unic-ucd-version",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unic-ucd-ident"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e230a37c0381caa9219d67cf063aa3a375ffed5bf541a452db16e744bdab6987"
|
||||
dependencies = [
|
||||
"unic-char-property",
|
||||
"unic-char-range",
|
||||
"unic-ucd-version",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unic-ucd-version"
|
||||
version = "0.9.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "96bd2f2237fe450fcd0a1d2f5f4e91711124f7857ba2e964247776ebeeb7b0c4"
|
||||
dependencies = [
|
||||
"unic-common",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicase"
|
||||
version = "2.9.0"
|
||||
|
|
@ -5203,40 +4987,12 @@ version = "1.13.3"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8"
|
||||
|
||||
[[package]]
|
||||
name = "unicode-width"
|
||||
version = "0.2.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254"
|
||||
|
||||
[[package]]
|
||||
name = "unicode_categories"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e"
|
||||
|
||||
[[package]]
|
||||
name = "unicode_names2"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d1673eca9782c84de5f81b82e4109dcfb3611c8ba0d52930ec4a9478f547b2dd"
|
||||
dependencies = [
|
||||
"phf",
|
||||
"unicode_names2_generator",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "unicode_names2_generator"
|
||||
version = "1.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b91e5b84611016120197efd7dc93ef76774f4e084cd73c9fb3ea4a86c570c56e"
|
||||
dependencies = [
|
||||
"getopts",
|
||||
"log",
|
||||
"phf_codegen",
|
||||
"rand 0.8.7",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "untrusted"
|
||||
version = "0.9.0"
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ pub enum Error {
|
|||
Timeout,
|
||||
#[error("AWS KMS decrypt failed")]
|
||||
Decrypt(#[from] #[redact] Box<SdkError<aws_sdk_kms::operation::decrypt::DecryptError>>),
|
||||
#[error("AWS Secrets Manager request preparation failed")]
|
||||
#[error("AWS Secrets Manager read failed")]
|
||||
Read(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError>>),
|
||||
#[error("AWS Secrets Manager create failed")]
|
||||
Create(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::create_secret::CreateSecretError>>),
|
||||
|
|
@ -22,6 +22,8 @@ pub enum Error {
|
|||
Delete(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::delete_secret::DeleteSecretError>>),
|
||||
#[error("AWS Secrets Manager replication failed")]
|
||||
Replicate(#[from] #[redact] Box<SdkError<aws_sdk_secretsmanager::operation::replicate_secret_to_regions::ReplicateSecretToRegionsError>>),
|
||||
#[error("AWS Secrets Manager response has no string payload")]
|
||||
MissingString,
|
||||
#[error("primary secret is not a JSON object")]
|
||||
PrimarySecret,
|
||||
#[error(transparent)]
|
||||
|
|
|
|||
|
|
@ -106,23 +106,24 @@ impl AwsSecretsManagerV2 {
|
|||
} else {
|
||||
self.async_read_secret(primary).await?
|
||||
};
|
||||
let object: Value = serde_json::from_str(
|
||||
value
|
||||
.as_ref()
|
||||
.map(SecretValue::expose)
|
||||
.filter(|v| !v.is_empty())
|
||||
.unwrap_or("{}"),
|
||||
)
|
||||
.map_err(|_| Error::PrimarySecret)?;
|
||||
let Some(value) = value else {
|
||||
return Ok(None);
|
||||
};
|
||||
let object: Value =
|
||||
serde_json::from_str(value.expose()).map_err(|_| Error::PrimarySecret)?;
|
||||
let object = object.as_object().ok_or(Error::PrimarySecret)?;
|
||||
Ok(object.get(name).cloned().and_then(Secret::from_json))
|
||||
Ok(object.get(name).cloned().map(Secret::from_json))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn async_read_secret(&self, name: &str) -> Result<Option<SecretValue>, Error> {
|
||||
match self.client.get_secret_value().secret_id(name).send().await {
|
||||
Ok(response) => Ok(response.secret_string.map(SecretValue::new)),
|
||||
Ok(response) => response
|
||||
.secret_string
|
||||
.map(SecretValue::new)
|
||||
.map(Some)
|
||||
.ok_or(Error::MissingString),
|
||||
Err(error)
|
||||
if matches!(
|
||||
&error,
|
||||
|
|
@ -131,11 +132,14 @@ impl AwsSecretsManagerV2 {
|
|||
{
|
||||
Err(Error::Timeout)
|
||||
}
|
||||
Err(error) if request_preparation_failed(&error) => Err(Error::Read(Box::new(error))),
|
||||
Err(_) => {
|
||||
tracing::error!("AWS secret read failed");
|
||||
Err(error)
|
||||
if error
|
||||
.as_service_error()
|
||||
.is_some_and(|error| error.is_resource_not_found_exception()) =>
|
||||
{
|
||||
Ok(None)
|
||||
}
|
||||
Err(error) => Err(Error::Read(Box::new(error))),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -281,17 +285,3 @@ fn bootstrap_key(name: &str) -> bool {
|
|||
| AWS_BEDROCK_RUNTIME_ENDPOINT
|
||||
)
|
||||
}
|
||||
|
||||
fn request_preparation_failed(
|
||||
error: &aws_sdk_secretsmanager::error::SdkError<
|
||||
aws_sdk_secretsmanager::operation::get_secret_value::GetSecretValueError,
|
||||
>,
|
||||
) -> bool {
|
||||
matches!(
|
||||
error,
|
||||
aws_sdk_secretsmanager::error::SdkError::ConstructionFailure(_)
|
||||
) || std::iter::successors(Some(error as &(dyn std::error::Error + 'static)), |error| {
|
||||
error.source()
|
||||
})
|
||||
.any(|source| source.is::<aws_credential_types::provider::error::CredentialsError>())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -290,3 +290,23 @@ async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() {
|
|||
Err(Error::Timeout)
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::denied(400, "AccessDeniedException")]
|
||||
#[case::throttled(400, "ThrottlingException")]
|
||||
#[case::unavailable(503, "ServiceUnavailableException")]
|
||||
#[tokio::test]
|
||||
async fn service_failures_remain_errors(#[case] status: u16, #[case] code: &str) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue"))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(json!({"__type":code})))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
assert!(matches!(
|
||||
manager(&server, KeyManagementSettings::default())
|
||||
.async_read_secret("key")
|
||||
.await,
|
||||
Err(Error::Read(_))
|
||||
));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -26,7 +26,7 @@ pub struct GoogleSecretManager {
|
|||
credentials: Arc<GoogleCredentials>,
|
||||
endpoint: reqwest::Url,
|
||||
project: String,
|
||||
cache: Cache<String, Option<SecretValue>>,
|
||||
cache: Cache<String, SecretValue>,
|
||||
always_read: bool,
|
||||
}
|
||||
|
||||
|
|
@ -119,7 +119,7 @@ impl GoogleSecretManager {
|
|||
if !self.always_read
|
||||
&& let Some(cached) = self.cache.get(name).await
|
||||
{
|
||||
return Ok(cached.and_then(cached_secret));
|
||||
return Ok(Some(Secret::String(cached)));
|
||||
}
|
||||
let url = self
|
||||
.endpoint
|
||||
|
|
@ -138,32 +138,20 @@ impl GoogleSecretManager {
|
|||
.headers(self.credentials.request_headers().await?)
|
||||
.send()
|
||||
.await?;
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND {
|
||||
return Ok(None);
|
||||
}
|
||||
if response.status() != reqwest::StatusCode::OK {
|
||||
self.cache.insert(name.to_owned(), None).await;
|
||||
return Err(Error::Status(response.status().as_u16()));
|
||||
}
|
||||
let response: Response = response.json().await?;
|
||||
let Some(data) = response.payload.and_then(|payload| payload.data) else {
|
||||
self.cache.insert(name.to_owned(), None).await;
|
||||
return Err(Error::MissingPayload);
|
||||
};
|
||||
let filtered: String = data
|
||||
.chars()
|
||||
.filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '='))
|
||||
.collect();
|
||||
let bytes = STANDARD.decode(filtered)?;
|
||||
let bytes = STANDARD.decode(data)?;
|
||||
let plaintext = String::from_utf8(bytes).map_err(|_| Error::Utf8)?;
|
||||
let value = SecretValue::new(plaintext);
|
||||
self.cache
|
||||
.insert(name.to_owned(), Some(value.clone()))
|
||||
.await;
|
||||
self.cache.insert(name.to_owned(), value.clone()).await;
|
||||
Ok(Some(Secret::String(value)))
|
||||
}
|
||||
}
|
||||
|
||||
fn cached_secret(value: SecretValue) -> Option<Secret> {
|
||||
match serde_json::from_str(value.expose()) {
|
||||
Ok(json) => Secret::from_json(json),
|
||||
Err(_) => Some(Secret::String(value)),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration};
|
|||
|
||||
use base64::{Engine, engine::general_purpose::STANDARD};
|
||||
use litellm_secrets_google::{Error, GoogleSecretManager};
|
||||
use litellm_secrets_types::Secret;
|
||||
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{header, path},
|
||||
|
|
@ -55,32 +55,58 @@ async fn successful_reads_use_auth_latest_version_and_cache_including_empty_valu
|
|||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::not_found(ResponseTemplate::new(404))]
|
||||
#[case::missing_payload(
|
||||
ResponseTemplate::new(200).set_body_json(serde_json::json!({"payload":{}}))
|
||||
)]
|
||||
#[case::not_found(404, serde_json::json!({}))]
|
||||
#[case::unauthorized(401, serde_json::json!({}))]
|
||||
#[case::forbidden(403, serde_json::json!({}))]
|
||||
#[case::throttled(429, serde_json::json!({}))]
|
||||
#[case::unavailable(503, serde_json::json!({}))]
|
||||
#[case::missing_payload(200, serde_json::json!({"payload":{}}))]
|
||||
#[case::invalid_base64(200, serde_json::json!({"payload":{"data":"%%%"}}))]
|
||||
#[tokio::test]
|
||||
async fn negative_cache_returns_none_after_initial_error(#[case] response: ResponseTemplate) {
|
||||
async fn failed_or_missing_reads_are_not_cached(
|
||||
#[case] status: u16,
|
||||
#[case] body: serde_json::Value,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
let manager = manager(&server, false, Duration::from_secs(60));
|
||||
let failing = Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(body))
|
||||
.expect(1)
|
||||
.mount_as_scoped(&server)
|
||||
.await;
|
||||
let result = manager.get_secret_from_google_secret_manager("key").await;
|
||||
match status {
|
||||
404 => assert_eq!(result.unwrap(), None),
|
||||
200 => assert!(matches!(
|
||||
result,
|
||||
Err(Error::MissingPayload | Error::Base64(_))
|
||||
)),
|
||||
status => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)),
|
||||
}
|
||||
drop(failing);
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
))
|
||||
.respond_with(response)
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("recovered")}})),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, false, Duration::from_secs(60));
|
||||
assert!(matches!(
|
||||
manager.get_secret_from_google_secret_manager("key").await,
|
||||
Err(Error::Status(404) | Error::MissingPayload)
|
||||
));
|
||||
assert!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.is_none()
|
||||
);
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
Some("recovered")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
|
|
@ -130,21 +156,12 @@ fn google_manager_requires_host_license_and_project_configuration() {
|
|||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::boolean("true", Some(Secret::Bool(true)))]
|
||||
#[case::null("null", None)]
|
||||
#[case::string(
|
||||
"\"text\"",
|
||||
Some(Secret::String(litellm_secrets_types::SecretValue::new("text")))
|
||||
)]
|
||||
#[case::object(
|
||||
"{\"key\":1}",
|
||||
Secret::from_json(serde_json::json!({"key":1}))
|
||||
)]
|
||||
#[case("true")]
|
||||
#[case("null")]
|
||||
#[case("\"text\"")]
|
||||
#[case("{\"key\":1}")]
|
||||
#[tokio::test]
|
||||
async fn cached_values_preserve_python_json_conversion(
|
||||
#[case] raw: &str,
|
||||
#[case] expected: Option<Secret>,
|
||||
) {
|
||||
async fn cache_preserves_raw_values(#[case] raw: &str) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/key/versions/latest:access",
|
||||
|
|
@ -157,20 +174,15 @@ async fn cached_values_preserve_python_json_conversion(
|
|||
.mount(&server)
|
||||
.await;
|
||||
let manager = manager(&server, false, Duration::from_secs(60));
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
Some(raw)
|
||||
);
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap(),
|
||||
expected
|
||||
);
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
manager
|
||||
.get_secret_from_google_secret_manager("key")
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.as_str(),
|
||||
Some(raw)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -14,12 +14,11 @@ impl From<SecretValue> for Secret {
|
|||
}
|
||||
|
||||
impl Secret {
|
||||
pub fn from_json(value: serde_json::Value) -> Option<Self> {
|
||||
pub fn from_json(value: serde_json::Value) -> Self {
|
||||
match value {
|
||||
serde_json::Value::Null => None,
|
||||
serde_json::Value::String(value) => Some(Self::String(SecretValue::new(value))),
|
||||
serde_json::Value::Bool(value) => Some(Self::Bool(value)),
|
||||
value => Some(Self::Json(value)),
|
||||
serde_json::Value::String(value) => Self::String(SecretValue::new(value)),
|
||||
serde_json::Value::Bool(value) => Self::Bool(value),
|
||||
value => Self::Json(value),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -21,13 +21,10 @@ strum.workspace = true
|
|||
jsonwebtoken.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
tracing = "0.1"
|
||||
reqwest.workspace = true
|
||||
moka.workspace = true
|
||||
tokio = { workspace = true, features = ["fs"] }
|
||||
|
||||
rustpython-parser = { version = "0.4.0", default-features = false, features = ["num-bigint"] }
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
wiremock = "0.6.5"
|
||||
|
|
|
|||
11
litellm-rust/crates/secrets/README.md
Normal file
11
litellm-rust/crates/secrets/README.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# Secret resolution
|
||||
|
||||
Construct `SecretManagerState::new(backend, settings)` for a configured manager or use `SecretManagerState::default()` for environment lookups. The configured backend determines its provider identity. Write-only settings and names excluded by `hosted_keys` use the environment directly. `secret_manager_would_be_consulted` follows the same routing decision as resolution
|
||||
|
||||
`get_secret` returns `Ok(Some(value))` for a found value, `Ok(None)` when no source contains the value, and `Err(error)` when lookup fails. For managed names, resolution checks the manager, then the environment, then the caller's default. An empty string, `false`, or an explicitly stored JSON null is a found value
|
||||
|
||||
Backend failures propagate by default. To allow fallback during a backend failure, construct the resolver with `.with_failure_policy(FailurePolicy::EnvironmentFallback)`. It then tries the environment and default, in that order. If neither exists, the original error is returned. This policy applies to manager lookups. Explicit OIDC references retain their own authentication errors and never fall back to environment secrets under the reference name
|
||||
|
||||
`get_secret` preserves value types. `get_secret_str` accepts a string default and rejects boolean or JSON values with `Error::TypeMismatch`. `get_secret_bool` accepts a boolean default and converts strings containing `true` or `false`, ignoring surrounding whitespace and ASCII case. Other strings and JSON values produce `Error::TypeMismatch`. Conversion failures never activate fallback or replace a found value with the default
|
||||
|
||||
Provider payloads remain strings unless explicitly selecting a field from an AWS primary JSON secret. Google caches only successfully decoded string payloads, so reads have identical values and types before and after caching. Confirmed absence and failed reads are not cached. AWS resource-not-found responses and Google HTTP 404 responses indicate absence. Other provider errors remain errors, and successful responses without the required payload are malformed responses rather than missing secrets
|
||||
|
|
@ -1,5 +1,3 @@
|
|||
use crate::KeyManagementSystem;
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum Error {
|
||||
#[error("encrypted environment value is missing")]
|
||||
|
|
@ -8,10 +6,6 @@ pub enum Error {
|
|||
InvalidCiphertext,
|
||||
#[error("decrypted value is not UTF-8")]
|
||||
Utf8,
|
||||
#[error("secret manager backend is not compiled: {0:?}")]
|
||||
UnsupportedBackend(KeyManagementSystem),
|
||||
#[error("configured secret manager does not match its backend")]
|
||||
BackendMismatch,
|
||||
#[error("unsupported OIDC provider or missing build feature")]
|
||||
UnsupportedOidc,
|
||||
#[error("OIDC reference requires a provider and audience")]
|
||||
|
|
@ -28,8 +22,8 @@ pub enum Error {
|
|||
UnsafeOidcPath,
|
||||
#[error("OIDC file could not be read")]
|
||||
OidcFile,
|
||||
#[error("secret manager returned no secret")]
|
||||
MissingSecret,
|
||||
#[error("secret cannot be converted to {expected}")]
|
||||
TypeMismatch { expected: &'static str },
|
||||
#[cfg(feature = "aws")]
|
||||
#[error(transparent)]
|
||||
Aws(#[from] litellm_secrets_aws::Error),
|
||||
|
|
|
|||
|
|
@ -76,9 +76,8 @@ pub async fn get_secret_from_manager(
|
|||
#[cfg(feature = "google")]
|
||||
SecretManager::GoogleSecretManager(client) => client
|
||||
.get_secret_from_google_secret_manager(secret_name)
|
||||
.await?
|
||||
.map(Some)
|
||||
.ok_or(Error::MissingSecret),
|
||||
.await
|
||||
.map_err(Error::from),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ pub use litellm_secrets_types::{
|
|||
AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue,
|
||||
};
|
||||
pub use oidc::{OidcProvider, OidcReference, OidcResolver};
|
||||
pub use resolver::SecretResolver;
|
||||
pub use resolver::{FailurePolicy, SecretResolver};
|
||||
pub use state::{SecretManagerState, secret_manager_would_be_consulted};
|
||||
|
||||
#[cfg(feature = "aws")]
|
||||
|
|
|
|||
|
|
@ -88,8 +88,13 @@ pub struct OidcResolver {
|
|||
|
||||
impl Default for OidcResolver {
|
||||
fn default() -> Self {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(600))
|
||||
.connect_timeout(Duration::from_secs(5))
|
||||
.build()
|
||||
.expect("HTTP client configuration");
|
||||
Self::new(
|
||||
reqwest::Client::builder().timeout(Duration::from_secs(600)).connect_timeout(Duration::from_secs(5)).build().expect("HTTP client configuration"),
|
||||
client,
|
||||
reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"),
|
||||
)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,14 +2,21 @@ use std::sync::Arc;
|
|||
|
||||
use litellm_core_utils::settings::{Lookup, ProcessEnvironment};
|
||||
|
||||
use crate::state::{LookupTarget, normalize_secret_name};
|
||||
use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue};
|
||||
|
||||
use crate::state::{LookupTarget, normalize_secret_name};
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub enum FailurePolicy {
|
||||
#[default]
|
||||
Propagate,
|
||||
EnvironmentFallback,
|
||||
}
|
||||
|
||||
pub struct SecretResolver {
|
||||
state: Arc<SecretManagerState>,
|
||||
environment: Arc<dyn Lookup + Send + Sync>,
|
||||
oidc: OidcResolver,
|
||||
failure_policy: FailurePolicy,
|
||||
}
|
||||
|
||||
impl Default for SecretResolver {
|
||||
|
|
@ -32,13 +39,21 @@ impl SecretResolver {
|
|||
state,
|
||||
environment,
|
||||
oidc,
|
||||
failure_policy: FailurePolicy::default(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_failure_policy(self, failure_policy: FailurePolicy) -> Self {
|
||||
Self {
|
||||
failure_policy,
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_secret(
|
||||
&self,
|
||||
name: &str,
|
||||
_default_value: Option<Secret>,
|
||||
default_value: Option<Secret>,
|
||||
) -> Result<Option<Secret>, Error> {
|
||||
let name = normalize_secret_name(name);
|
||||
if name.starts_with("oidc/") {
|
||||
|
|
@ -46,32 +61,26 @@ impl SecretResolver {
|
|||
.oidc
|
||||
.resolve(name, self.environment.as_ref())
|
||||
.await
|
||||
.map(|value| value.map(Secret::String));
|
||||
.map(|value| value.map(Secret::String).or(default_value));
|
||||
}
|
||||
if !self.state.readable() {
|
||||
return Ok(self
|
||||
.environment
|
||||
.get(name)
|
||||
.map(|value| match str_to_bool(&value) {
|
||||
Some(value) => Secret::Bool(value),
|
||||
None => Secret::String(SecretValue::new(value)),
|
||||
}));
|
||||
let LookupTarget::Manager { backend, settings } = self.state.lookup_target(name) else {
|
||||
return Ok(self.environment_secret(name).or(default_value));
|
||||
};
|
||||
match crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref())
|
||||
.await
|
||||
{
|
||||
Ok(value) => Ok(value
|
||||
.or_else(|| self.environment_secret(name))
|
||||
.or(default_value)),
|
||||
Err(error) => match self.failure_policy {
|
||||
FailurePolicy::Propagate => Err(error),
|
||||
FailurePolicy::EnvironmentFallback => self
|
||||
.environment_secret(name)
|
||||
.or(default_value)
|
||||
.map(Some)
|
||||
.ok_or(error),
|
||||
},
|
||||
}
|
||||
let result = match self.state.lookup_target(name) {
|
||||
LookupTarget::Environment => Ok(self.environment_secret(name)),
|
||||
LookupTarget::Manager { backend, settings } => {
|
||||
crate::get_secret_from_manager(backend, name, settings, self.environment.as_ref())
|
||||
.await
|
||||
}
|
||||
};
|
||||
let value = match result {
|
||||
Ok(value) => value,
|
||||
Err(_) => {
|
||||
tracing::error!("secret manager lookup failed; falling back to environment");
|
||||
self.environment_secret(name)
|
||||
}
|
||||
};
|
||||
Ok(value.and_then(managed_secret))
|
||||
}
|
||||
|
||||
fn environment_secret(&self, name: &str) -> Option<Secret> {
|
||||
|
|
@ -84,12 +93,18 @@ impl SecretResolver {
|
|||
pub async fn get_secret_str(
|
||||
&self,
|
||||
name: &str,
|
||||
default_value: Option<Secret>,
|
||||
default_value: Option<SecretValue>,
|
||||
) -> Result<Option<SecretValue>, Error> {
|
||||
Ok(match self.get_secret(name, default_value).await? {
|
||||
Some(Secret::String(value)) => Some(value),
|
||||
Some(Secret::Bool(_) | Secret::Json(_)) | None => None,
|
||||
})
|
||||
match self
|
||||
.get_secret(name, default_value.map(Secret::String))
|
||||
.await?
|
||||
{
|
||||
Some(Secret::String(value)) => Ok(Some(value)),
|
||||
None => Ok(None),
|
||||
Some(Secret::Bool(_) | Secret::Json(_)) => {
|
||||
Err(Error::TypeMismatch { expected: "string" })
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn get_secret_bool(
|
||||
|
|
@ -97,44 +112,24 @@ impl SecretResolver {
|
|||
name: &str,
|
||||
default_value: Option<bool>,
|
||||
) -> Result<Option<bool>, Error> {
|
||||
Ok(
|
||||
match self
|
||||
.get_secret(name, default_value.map(Secret::Bool))
|
||||
.await?
|
||||
{
|
||||
Some(Secret::Bool(value)) => Some(value),
|
||||
Some(Secret::String(value)) => str_to_bool(value.expose()),
|
||||
Some(Secret::Json(_)) | None => None,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
fn str_to_bool(value: &str) -> Option<bool> {
|
||||
match value.trim().to_ascii_lowercase().as_str() {
|
||||
"true" => Some(true),
|
||||
"false" => Some(false),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn literal_bool(value: &str) -> Option<bool> {
|
||||
use rustpython_parser::{Parse, ast};
|
||||
match ast::Expr::parse(value.trim_start_matches([' ', '\t']), "<secret>").ok()? {
|
||||
ast::Expr::Constant(node) => match node.value {
|
||||
ast::Constant::Bool(value) => Some(value),
|
||||
_ => None,
|
||||
},
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn managed_secret(value: Secret) -> Option<Secret> {
|
||||
match value {
|
||||
Secret::String(value) => Some(match literal_bool(value.expose()) {
|
||||
Some(boolean) => Secret::Bool(boolean),
|
||||
None => Secret::String(value),
|
||||
}),
|
||||
Secret::Bool(_) | Secret::Json(_) => None,
|
||||
match self
|
||||
.get_secret(name, default_value.map(Secret::Bool))
|
||||
.await?
|
||||
{
|
||||
Some(Secret::Bool(value)) => Ok(Some(value)),
|
||||
Some(Secret::String(value)) => {
|
||||
match value.expose().trim().to_ascii_lowercase().as_str() {
|
||||
"true" => Ok(Some(true)),
|
||||
"false" => Ok(Some(false)),
|
||||
_ => Err(Error::TypeMismatch {
|
||||
expected: "boolean",
|
||||
}),
|
||||
}
|
||||
}
|
||||
Some(Secret::Json(_)) => Err(Error::TypeMismatch {
|
||||
expected: "boolean",
|
||||
}),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
use crate::{Error, KeyManagementSettings, KeyManagementSystem, SecretManager};
|
||||
use crate::{KeyManagementSettings, KeyManagementSystem, SecretManager};
|
||||
|
||||
pub(crate) enum LookupTarget<'a> {
|
||||
Environment,
|
||||
|
|
@ -14,73 +14,37 @@ pub(crate) fn normalize_secret_name(name: &str) -> &str {
|
|||
|
||||
#[derive(Clone, Default)]
|
||||
pub struct SecretManagerState {
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<KeyManagementSettings>,
|
||||
backend: Option<SecretManager>,
|
||||
manager: Option<(SecretManager, KeyManagementSettings)>,
|
||||
}
|
||||
|
||||
impl SecretManagerState {
|
||||
pub fn new(
|
||||
system: Option<KeyManagementSystem>,
|
||||
settings: Option<KeyManagementSettings>,
|
||||
backend: Option<SecretManager>,
|
||||
) -> Result<Self, Error> {
|
||||
if let Some(system) = system {
|
||||
let available = match system {
|
||||
KeyManagementSystem::Local => true,
|
||||
KeyManagementSystem::AwsKms | KeyManagementSystem::AwsSecretManager => {
|
||||
cfg!(feature = "aws")
|
||||
}
|
||||
KeyManagementSystem::GoogleKms | KeyManagementSystem::GoogleSecretManager => {
|
||||
cfg!(feature = "google")
|
||||
}
|
||||
KeyManagementSystem::AzureKeyVault
|
||||
| KeyManagementSystem::HashicorpVault
|
||||
| KeyManagementSystem::Cyberark
|
||||
| KeyManagementSystem::Custom => false,
|
||||
};
|
||||
if !available {
|
||||
return Err(Error::UnsupportedBackend(system));
|
||||
}
|
||||
if let Some(backend) = &backend
|
||||
&& system != backend.system()
|
||||
{
|
||||
return Err(Error::BackendMismatch);
|
||||
}
|
||||
pub fn new(backend: SecretManager, settings: KeyManagementSettings) -> Self {
|
||||
Self {
|
||||
manager: Some((backend, settings)),
|
||||
}
|
||||
Ok(Self {
|
||||
system,
|
||||
settings,
|
||||
backend,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn system(&self) -> Option<KeyManagementSystem> {
|
||||
self.system
|
||||
}
|
||||
pub fn settings(&self) -> Option<&KeyManagementSettings> {
|
||||
self.settings.as_ref()
|
||||
}
|
||||
pub fn backend(&self) -> Option<&SecretManager> {
|
||||
self.backend.as_ref()
|
||||
self.backend().map(SecretManager::system)
|
||||
}
|
||||
|
||||
pub(crate) fn readable(&self) -> bool {
|
||||
self.backend.is_some()
|
||||
&& self
|
||||
.settings
|
||||
.as_ref()
|
||||
.is_some_and(|settings| settings.access_mode.readable())
|
||||
pub fn settings(&self) -> Option<&KeyManagementSettings> {
|
||||
self.manager.as_ref().map(|(_, settings)| settings)
|
||||
}
|
||||
|
||||
pub fn backend(&self) -> Option<&SecretManager> {
|
||||
self.manager.as_ref().map(|(backend, _)| backend)
|
||||
}
|
||||
|
||||
pub(crate) fn lookup_target(&self, name: &str) -> LookupTarget<'_> {
|
||||
match (&self.backend, &self.settings) {
|
||||
(Some(backend), Some(settings))
|
||||
if settings.access_mode.readable()
|
||||
&& hosts_secret(settings, name)
|
||||
&& self
|
||||
.system
|
||||
.is_some_and(|system| system != KeyManagementSystem::Local) =>
|
||||
match &self.manager {
|
||||
Some((backend, settings))
|
||||
if backend.system() != KeyManagementSystem::Local
|
||||
&& settings.access_mode.readable()
|
||||
&& settings
|
||||
.hosted_keys
|
||||
.as_ref()
|
||||
.is_none_or(|keys| keys.iter().any(|key| key == name)) =>
|
||||
{
|
||||
LookupTarget::Manager { backend, settings }
|
||||
}
|
||||
|
|
@ -90,16 +54,6 @@ impl SecretManagerState {
|
|||
}
|
||||
|
||||
pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool {
|
||||
state.readable()
|
||||
&& state
|
||||
.settings
|
||||
.as_ref()
|
||||
.is_some_and(|settings| hosts_secret(settings, normalize_secret_name(name)))
|
||||
}
|
||||
|
||||
fn hosts_secret(settings: &KeyManagementSettings, name: &str) -> bool {
|
||||
settings
|
||||
.hosted_keys
|
||||
.as_ref()
|
||||
.is_none_or(|keys| keys.iter().any(|key| key == name))
|
||||
let name = normalize_secret_name(name);
|
||||
!name.starts_with("oidc/") && matches!(state.lookup_target(name), LookupTarget::Manager { .. })
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,18 +1,13 @@
|
|||
use std::sync::Arc;
|
||||
|
||||
use litellm_secrets::{
|
||||
AccessMode, KeyManagementSettings, KeyManagementSystem, OidcResolver, Secret, SecretManager,
|
||||
SecretManagerState, SecretResolver, SecretValue, secret_manager_would_be_consulted,
|
||||
Error, KeyManagementSettings, OidcResolver, Secret, SecretManager, SecretManagerState,
|
||||
SecretResolver, SecretValue, secret_manager_would_be_consulted,
|
||||
};
|
||||
|
||||
fn resolver(value: Option<&str>, readable: bool) -> SecretResolver {
|
||||
let state = if readable {
|
||||
SecretManagerState::new(
|
||||
Some(KeyManagementSystem::Local),
|
||||
Some(KeyManagementSettings::default()),
|
||||
Some(SecretManager::Local),
|
||||
)
|
||||
.unwrap()
|
||||
fn resolver(value: Option<&str>, configured: bool) -> SecretResolver {
|
||||
let state = if configured {
|
||||
SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default())
|
||||
} else {
|
||||
SecretManagerState::default()
|
||||
};
|
||||
|
|
@ -25,77 +20,107 @@ fn resolver(value: Option<&str>, readable: bool) -> SecretResolver {
|
|||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::lowercase_true("true", Some(true), None)]
|
||||
#[case::whitespace_lowercase_false(" FALSE ", Some(false), None)]
|
||||
#[case::python_true("True", Some(true), Some(true))]
|
||||
#[case::python_false("False", Some(false), Some(false))]
|
||||
#[case::parenthesized_python_true("(True)", None, Some(true))]
|
||||
#[case::commented_python_false("False # comment", None, Some(false))]
|
||||
#[case::integer("1", None, None)]
|
||||
#[case::yes("yes", None, None)]
|
||||
#[case::plain_string("secret", None, None)]
|
||||
#[case("true", Some(true))]
|
||||
#[case(" FALSE ", Some(false))]
|
||||
#[case("(True)", None)]
|
||||
#[case("False # comment", None)]
|
||||
#[case("1", None)]
|
||||
#[case("secret", None)]
|
||||
#[tokio::test]
|
||||
async fn boolean_conversion_preserves_local_and_manager_differences(
|
||||
async fn conversion_is_explicit_and_independent_of_manager_configuration(
|
||||
#[case] input: &str,
|
||||
#[case] local: Option<bool>,
|
||||
#[case] manager: Option<bool>,
|
||||
#[values(false, true)] readable: bool,
|
||||
#[case] boolean: Option<bool>,
|
||||
#[values(false, true)] configured: bool,
|
||||
) {
|
||||
let boolean = if readable { manager } else { local };
|
||||
let resolver = resolver(Some(input), readable);
|
||||
let expected = boolean
|
||||
.map(Secret::Bool)
|
||||
.unwrap_or_else(|| Secret::String(SecretValue::new(input)));
|
||||
let resolver = resolver(Some(input), configured);
|
||||
assert_eq!(
|
||||
resolver.get_secret("key", None).await.unwrap(),
|
||||
Some(expected)
|
||||
Some(Secret::String(SecretValue::new(input)))
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret_str("key", None)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|v| v.expose().to_owned()),
|
||||
boolean.is_none().then(|| input.to_owned())
|
||||
.unwrap()
|
||||
.expose(),
|
||||
input
|
||||
);
|
||||
match boolean {
|
||||
Some(value) => assert_eq!(
|
||||
resolver.get_secret_bool("key", None).await.unwrap(),
|
||||
Some(value)
|
||||
),
|
||||
None => assert!(matches!(
|
||||
resolver.get_secret_bool("key", Some(true)).await,
|
||||
Err(Error::TypeMismatch {
|
||||
expected: "boolean"
|
||||
})
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn defaults_apply_only_to_absence(#[values(false, true)] configured: bool) {
|
||||
let missing = resolver(None, configured);
|
||||
assert_eq!(missing.get_secret("key", None).await.unwrap(), None);
|
||||
assert_eq!(
|
||||
missing.get_secret_bool("key", Some(false)).await.unwrap(),
|
||||
Some(false)
|
||||
);
|
||||
assert_eq!(
|
||||
missing
|
||||
.get_secret_str("key", Some(SecretValue::new("default")))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
"default"
|
||||
);
|
||||
for value in [
|
||||
Secret::Bool(false),
|
||||
Secret::from_json(serde_json::json!({"key":1})),
|
||||
Secret::from_json(serde_json::Value::Null),
|
||||
] {
|
||||
assert_eq!(
|
||||
missing
|
||||
.get_secret("key", Some(value.clone()))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(value)
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
resolver(Some(""), configured)
|
||||
.get_secret_str("key", Some(SecretValue::new("default")))
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
""
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn manager_boolean_conversion_trims_whitespace() {
|
||||
async fn prefix_is_removed_once_and_local_manager_is_not_consulted() {
|
||||
let state = SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default());
|
||||
assert_eq!(
|
||||
resolver(Some(" true "), true)
|
||||
.get_secret_bool("key", None)
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(true)
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() {
|
||||
let missing = resolver(None, false);
|
||||
assert_eq!(
|
||||
missing
|
||||
.get_secret("missing", Some(Secret::Bool(true)))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
);
|
||||
assert_eq!(
|
||||
missing
|
||||
.get_secret_bool("missing", Some(true))
|
||||
.await
|
||||
.unwrap(),
|
||||
None
|
||||
state.system(),
|
||||
Some(litellm_secrets::KeyManagementSystem::Local)
|
||||
);
|
||||
assert!(!secret_manager_would_be_consulted(
|
||||
&state,
|
||||
"os.environ/os.environ/KEY"
|
||||
));
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(SecretManagerState::default()),
|
||||
Arc::new(|name: &str| (name == "KEY").then(|| "value".into())),
|
||||
Arc::new(state),
|
||||
Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())),
|
||||
OidcResolver::default(),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret_str("os.environ/KEY", None)
|
||||
.get_secret_str("os.environ/os.environ/KEY", None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
|
|
@ -104,177 +129,6 @@ async fn missing_values_ignore_defaults_and_prefix_is_removed_before_lookup() {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::all_keys(None)]
|
||||
#[case::no_keys(Some(Vec::new()))]
|
||||
#[case::allowlisted_key(Some(vec!["KEY".into()]))]
|
||||
fn manager_gating_requires_client_readable_settings_and_allowlisted_name(
|
||||
#[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)]
|
||||
access_mode: AccessMode,
|
||||
#[values(false, true)] client: bool,
|
||||
#[case] keys: Option<Vec<String>>,
|
||||
) {
|
||||
let expected = client
|
||||
&& access_mode.readable()
|
||||
&& keys
|
||||
.as_ref()
|
||||
.is_none_or(|keys| keys.iter().any(|key| key == "KEY"));
|
||||
let state = SecretManagerState::new(
|
||||
Some(KeyManagementSystem::Local),
|
||||
Some(KeyManagementSettings {
|
||||
access_mode,
|
||||
hosted_keys: keys,
|
||||
..Default::default()
|
||||
}),
|
||||
client.then_some(SecretManager::Local),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
secret_manager_would_be_consulted(&state, "os.environ/KEY"),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn manager_gating_requires_settings() {
|
||||
let no_settings = SecretManagerState::new(None, None, Some(SecretManager::Local)).unwrap();
|
||||
assert!(!secret_manager_would_be_consulted(&no_settings, "KEY"));
|
||||
}
|
||||
|
||||
#[cfg(feature = "aws")]
|
||||
#[rstest::rstest]
|
||||
#[case::missing_value(None, None)]
|
||||
#[case::lookup_error(Some("primary".to_owned()), Some("environment-value"))]
|
||||
#[tokio::test]
|
||||
async fn aws_missing_values_do_not_fallback_but_lookup_errors_do(
|
||||
#[case] primary: Option<String>,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
use litellm_secrets::aws::AwsSecretsManagerV2;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_partial_json};
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(body_partial_json(serde_json::json!({"SecretId":"KEY"})))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({})))
|
||||
.expect(u64::from(primary.is_none()))
|
||||
.mount(&server)
|
||||
.await;
|
||||
Mock::given(body_partial_json(serde_json::json!({"SecretId":"primary"})))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"SecretString":"invalid-json"})),
|
||||
)
|
||||
.expect(u64::from(primary.is_some()))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let endpoint = server.uri();
|
||||
let environment: Arc<dyn litellm_core_utils::settings::Lookup + Send + Sync> =
|
||||
Arc::new(move |name: &str| match name {
|
||||
"AWS_REGION_NAME" => Some("us-east-1".into()),
|
||||
"AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()),
|
||||
"AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()),
|
||||
"KEY" => Some("environment-value".into()),
|
||||
_ => None,
|
||||
});
|
||||
let settings = KeyManagementSettings {
|
||||
primary_secret_name: primary,
|
||||
..Default::default()
|
||||
};
|
||||
let manager = AwsSecretsManagerV2::load_aws_secret_manager(
|
||||
Some(true),
|
||||
settings.clone(),
|
||||
environment.clone(),
|
||||
)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
let state = SecretManagerState::new(
|
||||
Some(KeyManagementSystem::AwsSecretManager),
|
||||
Some(settings),
|
||||
Some(SecretManager::AwsSecretsManagerV2(manager)),
|
||||
)
|
||||
.unwrap();
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(state),
|
||||
environment.clone(),
|
||||
OidcResolver::default(),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret_str("os.environ/KEY", None)
|
||||
.await
|
||||
.unwrap()
|
||||
.map(|v| v.expose().to_owned())
|
||||
.as_deref(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(feature = "google")]
|
||||
#[rstest::rstest]
|
||||
#[case::hosted_filter(Some(Vec::new()), Some(KeyManagementSystem::GoogleSecretManager))]
|
||||
#[case::negative_cache(None, Some(KeyManagementSystem::GoogleSecretManager))]
|
||||
#[case::missing_system(None, None)]
|
||||
#[case::hosted_nested_prefix(Some(vec!["os.environ/KEY".into()]), Some(KeyManagementSystem::GoogleSecretManager))]
|
||||
#[tokio::test]
|
||||
async fn google_negative_cache_still_falls_back_and_hosted_filter_avoids_io(
|
||||
#[case] hosted_keys: Option<Vec<String>>,
|
||||
#[case] system: Option<KeyManagementSystem>,
|
||||
) {
|
||||
use litellm_secrets::google::GoogleSecretManager;
|
||||
use std::time::Duration;
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::path};
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(path(
|
||||
"/v1/projects/project/secrets/os%2Eenviron%2FKEY/versions/latest:access",
|
||||
))
|
||||
.respond_with(ResponseTemplate::new(404))
|
||||
.expect(u64::from(
|
||||
hosted_keys.as_ref().is_none_or(|keys| !keys.is_empty()) && system.is_some(),
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let environment: Arc<dyn litellm_core_utils::settings::Lookup + Send + Sync> =
|
||||
Arc::new(|name: &str| match name {
|
||||
"VERTEX_AI_API_KEY" => Some("token".into()),
|
||||
"os.environ/KEY" => Some("environment-value".into()),
|
||||
_ => None,
|
||||
});
|
||||
let manager = GoogleSecretManager::with_client(
|
||||
reqwest::Client::new(),
|
||||
server.uri().parse().unwrap(),
|
||||
"project".into(),
|
||||
environment.clone(),
|
||||
Some(Duration::from_secs(60)),
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
let settings = KeyManagementSettings {
|
||||
hosted_keys,
|
||||
..Default::default()
|
||||
};
|
||||
let state = SecretManagerState::new(
|
||||
system,
|
||||
Some(settings),
|
||||
Some(SecretManager::GoogleSecretManager(manager)),
|
||||
)
|
||||
.unwrap();
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(state),
|
||||
environment.clone(),
|
||||
OidcResolver::default(),
|
||||
);
|
||||
for _ in 0..2 {
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret_str("os.environ/os.environ/KEY", None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
"environment-value"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolver_future_can_run_on_a_tokio_worker() {
|
||||
let resolver = resolver(Some("worker-value"), false);
|
||||
|
|
@ -285,59 +139,230 @@ async fn resolver_future_can_run_on_a_tokio_worker() {
|
|||
assert_eq!(result.unwrap().expose(), "worker-value");
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::nested_true("((True)) # comment", Some(true))]
|
||||
#[case::commented_false("(False # comment\n)", Some(false))]
|
||||
#[case::boolean_expression("True and False", None)]
|
||||
#[case::string_literal("'True'", None)]
|
||||
#[case::tuple("(True,)", None)]
|
||||
#[case::unary_expression("not False", None)]
|
||||
#[case::multiple_expressions("True\nFalse", None)]
|
||||
#[case::incomplete_expression("(True", None)]
|
||||
#[tokio::test]
|
||||
async fn manager_boolean_literals_follow_python_syntax(
|
||||
#[case] input: &str,
|
||||
#[case] expected: Option<bool>,
|
||||
) {
|
||||
let value = resolver(Some(input), true)
|
||||
.get_secret("key", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
value,
|
||||
Some(
|
||||
expected
|
||||
.map(Secret::Bool)
|
||||
.unwrap_or_else(|| Secret::String(SecretValue::new(input)))
|
||||
#[cfg(feature = "aws")]
|
||||
mod aws {
|
||||
use super::*;
|
||||
use litellm_secrets::{AccessMode, FailurePolicy, aws::AwsSecretsManagerV2};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
|
||||
|
||||
fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState {
|
||||
let endpoint = server.uri();
|
||||
let environment = Arc::new(move |name: &str| match name {
|
||||
"AWS_REGION_NAME" => Some("us-east-1".into()),
|
||||
"AWS_ACCESS_KEY_ID" | "AWS_SECRET_ACCESS_KEY" => Some("test".into()),
|
||||
"AWS_BEDROCK_RUNTIME_ENDPOINT" => Some(endpoint.clone()),
|
||||
_ => None,
|
||||
});
|
||||
let manager =
|
||||
AwsSecretsManagerV2::load_aws_secret_manager(Some(true), settings.clone(), environment)
|
||||
.unwrap()
|
||||
.unwrap();
|
||||
SecretManagerState::new(SecretManager::AwsSecretsManagerV2(manager), settings)
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::missing(400, serde_json::json!({"__type":"ResourceNotFoundException"}), false)]
|
||||
#[case::denied(400, serde_json::json!({"__type":"AccessDeniedException"}), true)]
|
||||
#[case::malformed(200, serde_json::json!({}), true)]
|
||||
#[tokio::test]
|
||||
async fn failure_policy_preserves_errors_and_fallback_precedence(
|
||||
#[case] status: u16,
|
||||
#[case] body: serde_json::Value,
|
||||
#[case] fails: bool,
|
||||
#[values(FailurePolicy::Propagate, FailurePolicy::EnvironmentFallback)]
|
||||
policy: FailurePolicy,
|
||||
#[values(None, Some("environment"))] environment: Option<&'static str>,
|
||||
#[values(None, Some("default"))] default: Option<&str>,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(status).set_body_json(body))
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(state(&server, KeyManagementSettings::default())),
|
||||
Arc::new(move |_: &str| environment.map(str::to_owned)),
|
||||
OidcResolver::default(),
|
||||
)
|
||||
);
|
||||
.with_failure_policy(policy);
|
||||
let result = resolver
|
||||
.get_secret_str("KEY", default.map(SecretValue::new))
|
||||
.await;
|
||||
let fallback = environment.or(default);
|
||||
if fails && (policy == FailurePolicy::Propagate || fallback.is_none()) {
|
||||
assert!(matches!(result, Err(Error::Aws(_))));
|
||||
} else {
|
||||
assert_eq!(result.unwrap().as_ref().map(SecretValue::expose), fallback);
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[case::boolean(serde_json::json!(false))]
|
||||
#[case::object(serde_json::json!({"key":1}))]
|
||||
#[case::null(serde_json::Value::Null)]
|
||||
#[case::string(serde_json::json!("true"))]
|
||||
#[tokio::test]
|
||||
async fn typed_values_survive_resolution_and_accessors_reject_wrong_types(
|
||||
#[case] value: serde_json::Value,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_json(
|
||||
serde_json::json!({"SecretString":serde_json::json!({"KEY":value}).to_string()}),
|
||||
))
|
||||
.expect(3)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let settings = KeyManagementSettings {
|
||||
primary_secret_name: Some("primary".into()),
|
||||
..Default::default()
|
||||
};
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(state(&server, settings)),
|
||||
Arc::new(|_: &str| Some("fallback".into())),
|
||||
OidcResolver::default(),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret("KEY", Some(Secret::Bool(true)))
|
||||
.await
|
||||
.unwrap(),
|
||||
Some(Secret::from_json(value.clone()))
|
||||
);
|
||||
match &value {
|
||||
serde_json::Value::String(text) => assert_eq!(
|
||||
resolver
|
||||
.get_secret_str("KEY", None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
text
|
||||
),
|
||||
_ => assert!(matches!(
|
||||
resolver.get_secret_str("KEY", None).await,
|
||||
Err(Error::TypeMismatch { expected: "string" })
|
||||
)),
|
||||
}
|
||||
match value {
|
||||
serde_json::Value::Bool(boolean) => assert_eq!(
|
||||
resolver.get_secret_bool("KEY", None).await.unwrap(),
|
||||
Some(boolean)
|
||||
),
|
||||
serde_json::Value::String(_) => assert_eq!(
|
||||
resolver.get_secret_bool("KEY", None).await.unwrap(),
|
||||
Some(true)
|
||||
),
|
||||
_ => assert!(matches!(
|
||||
resolver.get_secret_bool("KEY", None).await,
|
||||
Err(Error::TypeMismatch {
|
||||
expected: "boolean"
|
||||
})
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest::rstest]
|
||||
#[tokio::test]
|
||||
async fn gating_prediction_matches_actual_lookup(
|
||||
#[values(AccessMode::ReadOnly, AccessMode::WriteOnly, AccessMode::ReadAndWrite)]
|
||||
access_mode: AccessMode,
|
||||
#[values(None, Some(vec![]), Some(vec!["KEY".into()]))] hosted_keys: Option<Vec<String>>,
|
||||
#[values("os.environ/KEY", "os.environ/oidc/env/KEY")] name: &str,
|
||||
) {
|
||||
let server = MockServer::start().await;
|
||||
let expected = name == "os.environ/KEY"
|
||||
&& access_mode.readable()
|
||||
&& hosted_keys
|
||||
.as_ref()
|
||||
.is_none_or(|keys| keys.iter().any(|key| key == "KEY"));
|
||||
Mock::given(method("POST"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_json(serde_json::json!({"SecretString":"remote"})),
|
||||
)
|
||||
.expect(u64::from(expected))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let state = state(
|
||||
&server,
|
||||
KeyManagementSettings {
|
||||
access_mode,
|
||||
hosted_keys,
|
||||
..Default::default()
|
||||
},
|
||||
);
|
||||
assert!(state.backend().is_some());
|
||||
assert_eq!(state.settings().unwrap().access_mode, access_mode);
|
||||
assert_eq!(secret_manager_would_be_consulted(&state, name), expected);
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(state),
|
||||
Arc::new(|_: &str| Some("environment".into())),
|
||||
OidcResolver::default(),
|
||||
);
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret_str(name, None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
if expected { "remote" } else { "environment" }
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "google")]
|
||||
#[rstest::rstest]
|
||||
#[case::missing(404)]
|
||||
#[case::failure(503)]
|
||||
#[tokio::test]
|
||||
async fn environment_prefix_is_removed_only_once_and_gating_uses_the_same_name() {
|
||||
let name = "os.environ/folder/os.environ/KEY";
|
||||
let state = SecretManagerState::new(
|
||||
Some(KeyManagementSystem::Local),
|
||||
Some(KeyManagementSettings {
|
||||
hosted_keys: Some(vec!["folder/os.environ/KEY".into()]),
|
||||
..Default::default()
|
||||
}),
|
||||
Some(SecretManager::Local),
|
||||
async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) {
|
||||
use litellm_secrets::{FailurePolicy, google::GoogleSecretManager};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method};
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.respond_with(ResponseTemplate::new(status))
|
||||
.expect(2)
|
||||
.mount(&server)
|
||||
.await;
|
||||
let environment: Arc<dyn litellm_core_utils::settings::Lookup + Send + Sync> =
|
||||
Arc::new(|name: &str| match name {
|
||||
"VERTEX_AI_API_KEY" => Some("token".into()),
|
||||
"KEY" => Some("environment".into()),
|
||||
_ => None,
|
||||
});
|
||||
let manager = GoogleSecretManager::with_client(
|
||||
reqwest::Client::new(),
|
||||
server.uri().parse().unwrap(),
|
||||
"project".into(),
|
||||
environment.clone(),
|
||||
None,
|
||||
false,
|
||||
)
|
||||
.unwrap();
|
||||
assert!(secret_manager_would_be_consulted(&state, name));
|
||||
let resolver = SecretResolver::new(
|
||||
Arc::new(state),
|
||||
Arc::new(|name: &str| (name == "folder/os.environ/KEY").then(|| "value".into())),
|
||||
OidcResolver::default(),
|
||||
let state = SecretManagerState::new(
|
||||
SecretManager::GoogleSecretManager(manager),
|
||||
KeyManagementSettings::default(),
|
||||
);
|
||||
let resolver = SecretResolver::new(Arc::new(state), environment, OidcResolver::default());
|
||||
let result = resolver.get_secret_str("KEY", None).await;
|
||||
if status == 404 {
|
||||
assert_eq!(result.unwrap().unwrap().expose(), "environment");
|
||||
} else {
|
||||
assert!(
|
||||
matches!(result, Err(Error::Google(litellm_secrets::google::Error::Status(actual))) if actual == status)
|
||||
);
|
||||
}
|
||||
assert_eq!(
|
||||
resolver
|
||||
.get_secret_str(name, None)
|
||||
.with_failure_policy(FailurePolicy::EnvironmentFallback)
|
||||
.get_secret_str("KEY", None)
|
||||
.await
|
||||
.unwrap()
|
||||
.unwrap()
|
||||
.expose(),
|
||||
"value"
|
||||
"environment"
|
||||
);
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue