diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 70d828b5c8b..278fa7c425f 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -127,6 +127,13 @@ jobs: cargo check -p litellm-python-bridge --locked --no-default-features --features "abi3${features:+,$features}" done + - name: Test secret manager feature combinations + run: | + cargo test -p litellm-auth-gcp --locked --no-default-features + for features in '' aws google aws,google; do + cargo test -p litellm-secrets --locked --no-default-features --features "$features" + done + rust-wheel: runs-on: ubuntu-latest timeout-minutes: 30 diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index af3a31ddbfa..8c35a0be0b4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -82,6 +82,16 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03918c3dbd7701a85c6b9887732e2921175f26c350b4563841d0958c21d57e6d" +[[package]] +name = "assert-json-diff" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e4f2b81832e72834d7518d8487a0396a28cc408186a2e8854c0f98011faf12" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "async-compression" version = "0.4.46" @@ -191,9 +201,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.8.1" +version = "1.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7816e98ee912159f45d307e5ee6bfea4a335a55aee15f7f3e32f81a6f3000f1d" +checksum = "25b43ad47adc2517efe3d706559d94b97e50e80e0321b3cadbc9f77cee88adcd" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -214,6 +224,58 @@ dependencies = [ "uuid", ] +[[package]] +name = "aws-sdk-kms" +version = "1.120.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6b0fe38fee2ba5b6cd24d32d08365b314ae1adea123649e061a7eb6300b6f5b" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + +[[package]] +name = "aws-sdk-secretsmanager" +version = "1.117.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d32d781b34ab083e0dc54b4c68fc5e89ddf35e97d97bdcb9386d21325c14767" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "http 0.2.12", + "http 1.4.2", + "regex-lite", + "tracing", +] + [[package]] name = "aws-sdk-sts" version = "1.108.0" @@ -243,9 +305,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.5.1" +version = "1.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "723c2234ad7511ceef63eab016b7ba6ff7c55590fefb96fa8467af014a07309f" +checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -308,9 +370,9 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.2.0" +version = "1.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e" +checksum = "7bd25384a4e437aa8d8f339afad4b69e786b936a7cb10db668a7aaf66717b1a8" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -338,9 +400,9 @@ dependencies = [ [[package]] name = "aws-smithy-json" -version = "0.63.0" +version = "0.63.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3dc65a121adb4b33729919fcfa14fa36fb33c1555a8f06bb0e2188dbfdc1d9ef" +checksum = "3385d469edbe8b60cc72002784652b5efca39178192aa9cc4b44c9875c6bdc18" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-schema", @@ -368,9 +430,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime" -version = "1.12.0" +version = "1.14.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bea94a9ff8464016338c851e24b472d7131c388c88898a502e781815b2ee6045" +checksum = "3296253d3a91b3f938a3f2bcce4daebadcb4aa4228153fcec90f9d23532b4484" dependencies = [ "aws-smithy-async", "aws-smithy-http", @@ -394,9 +456,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.13.0" +version = "1.16.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ed1ebe6e0a95ea84570225f5a8208dec4b8f77e61a9b0d6f51773fcb4612f0" +checksum = "6d881a7b7ad179fd6611680c9de89f716fb00ab40299a9a7b8c6913e8f7511a8" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", @@ -423,9 +485,9 @@ dependencies = [ [[package]] name = "aws-smithy-schema" -version = "0.2.0" +version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" +checksum = "e8f395d93304280b64b7632fea798d177e74897fe7f063416ce627cd6fa24829" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -434,9 +496,9 @@ dependencies = [ [[package]] name = "aws-smithy-types" -version = "1.6.1" +version = "1.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8" +checksum = "0b791f3ac597193fe1d08b82366986eb1f5bc31f2ac6c194c0855276116c76cd" dependencies = [ "base64-simd", "bytes", @@ -472,9 +534,9 @@ dependencies = [ [[package]] name = "aws-types" -version = "1.4.0" +version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e957a6c6dbce82b7a91f44231c09273159703769f447cbe85e854dfe9cf67f86" +checksum = "209f3a6d82a6e9e5f94abbed94c7a26e1c052341002bf57a5fb5481f625896fc" dependencies = [ "aws-credential-types", "aws-smithy-async", @@ -580,6 +642,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.1" @@ -632,6 +700,9 @@ name = "bytes" version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" +dependencies = [ + "serde", +] [[package]] name = "bytes-utils" @@ -1065,6 +1136,55 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be1e0bca6c3637f992fc1cc7cbc52a78c1ef6db076dbf1059c4323d6a2048376" +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" + +[[package]] +name = "defmt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2953bfe4f93bbd20cc71198842756f77d161884c99ebbabc41d80231ded88d1" +dependencies = [ + "bitflags 1.3.2", + "defmt-macros", +] + +[[package]] +name = "defmt-macros" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bad9c72e7ca2137e0dc3813245a0d282fd6daad32fd800af018306a9169b5fe8" +dependencies = [ + "defmt-parser", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "defmt-parser" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10d60334b3b2e7c9d91ef8150abfb6fa4c1c39ebbcf4a81c2e346aad939fee3e" +dependencies = [ + "thiserror 2.0.19", +] + [[package]] name = "deranged" version = "0.5.8" @@ -1440,6 +1560,224 @@ version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" +[[package]] +name = "google-cloud-auth" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff461519b1a948200f163574be072753bcfb462a323f0eb426629d89872dd685" +dependencies = [ + "async-trait", + "base64 0.23.1", + "bytes", + "google-cloud-gax", + "hex", + "hmac", + "http 1.4.2", + "jiff", + "reqwest 0.13.5", + "rustc_version", + "rustls 0.23.42", + "rustls-pki-types", + "serde", + "serde_json", + "sha2 0.11.0", + "thiserror 2.0.19", + "time", + "tokio", + "url", +] + +[[package]] +name = "google-cloud-gax" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5615cff28ee59cfe52fbb4c11b8b1e77f650296e2ea4f4c2b7757ac6b19e752" +dependencies = [ + "bytes", + "futures", + "google-cloud-rpc", + "google-cloud-wkt", + "http 1.4.2", + "pin-project", + "rand 0.10.2", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", +] + +[[package]] +name = "google-cloud-gax-internal" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2766757d877a7a8ac23da9884cb0e3f10ed9b75a0ce59801ce6b19bf9d5819e" +dependencies = [ + "bytes", + "futures", + "google-cloud-auth", + "google-cloud-gax", + "google-cloud-rpc", + "google-cloud-wkt", + "h2 0.4.15", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "lazy_static", + "opentelemetry", + "opentelemetry-semantic-conventions", + "opentelemetry_sdk", + "percent-encoding", + "pin-project", + "prost", + "prost-types", + "reqwest 0.13.5", + "rustc_version", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tokio-stream", + "tonic", + "tonic-prost", + "tower", + "tracing", + "tracing-opentelemetry", +] + +[[package]] +name = "google-cloud-iam-v1" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f962b40234b1531e6ef73f7558871c96e117231e962086c98804323fb8d2c82" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-type", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-kms-v1" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d0f6eab19254d9abd98035cd54e3f2522d2c49abf9d93bf5425ee738d198c15" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-iam-v1", + "google-cloud-location", + "google-cloud-longrunning", + "google-cloud-lro", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-location" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "280d5acdba8fcb1232c0719ed788d85b7e362b82cbb425b7050d3ce46f075ede" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-longrunning" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c0363c5389ffda2b55cd8a86eef4b19a3481a48467c91dc5d77f626a9572766" +dependencies = [ + "async-trait", + "bytes", + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", + "tracing", +] + +[[package]] +name = "google-cloud-lro" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47af3deef75c14a2983c430898d960c765bddbcc9f9188ca0563108e9227cfe7" +dependencies = [ + "google-cloud-gax", + "google-cloud-gax-internal", + "google-cloud-longrunning", + "google-cloud-rpc", + "google-cloud-wkt", + "serde", + "tokio", + "tracing", +] + +[[package]] +name = "google-cloud-rpc" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2162c08a89118130979ba261080e960e44cdcb2d6e2ab8ca9b1da245285d353" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-type" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63acc3a92a85f96bab021c3a3e29b53bbacc97651e1b524d4c2991960a63eb82" +dependencies = [ + "bytes", + "google-cloud-wkt", + "serde", + "serde_json", + "serde_with", +] + +[[package]] +name = "google-cloud-wkt" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fccf98cfd5481a5f5a285181ab0c62123d7d47cd2bb7299448440649349e4e7" +dependencies = [ + "base64 0.22.1", + "bytes", + "serde", + "serde_json", + "serde_with", + "thiserror 2.0.19", + "time", + "url", +] + [[package]] name = "h2" version = "0.3.27" @@ -1507,6 +1845,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17592d60ebacc7d5e169f4663c5f84f9161cc90328abcfe8456f41e4dfcb284" + [[package]] name = "hex" version = "0.4.3" @@ -1636,6 +1980,7 @@ dependencies = [ "http 1.4.2", "http-body 1.1.0", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -1675,6 +2020,19 @@ dependencies = [ "webpki-roots", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.10.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.20" @@ -1884,6 +2242,43 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jiff" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ab1baf72f08796de0260609515130699b890ac25f30e610ad894bc5856cafdb" +dependencies = [ + "defmt", + "jiff-core", + "jiff-static", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", +] + +[[package]] +name = "jiff-core" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e52fe76043ccecc9005d2305ebaadf7d7fc0cc89ca6baa10a94d6bc68c7128c" +dependencies = [ + "defmt", + "log", +] + +[[package]] +name = "jiff-static" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378268a1116ad67ae6228701118ac9f491d78fda38a40a1f1a9e1348de6f7212" +dependencies = [ + "jiff-core", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "jni" version = "0.22.4" @@ -1954,6 +2349,21 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "jsonwebtoken" +version = "11.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e75fe14a82d81e5f5af639997db37d8b96045938a7ac6ab18cdbe1c7467e05e1" +dependencies = [ + "base64 0.22.1", + "getrandom 0.2.17", + "js-sys", + "serde", + "serde_json", + "signature", + "zeroize", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -1976,11 +2386,10 @@ checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" name = "litellm-auth" version = "0.1.0" dependencies = [ - "serde", - "subtle", - "thiserror 2.0.19", - "tokio", - "veil", + "litellm-auth-aws", + "litellm-auth-azure", + "litellm-auth-gcp", + "litellm-auth-types", ] [[package]] @@ -1993,7 +2402,7 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", - "litellm-auth", + "litellm-auth-types", "litellm-http", "moka", "reqwest 0.12.28", @@ -2009,7 +2418,7 @@ version = "0.1.0" dependencies = [ "azure_core", "azure_identity", - "litellm-auth", + "litellm-auth-types", "moka", "rstest", "serde_json", @@ -2024,13 +2433,26 @@ name = "litellm-auth-gcp" version = "0.1.0" dependencies = [ "gcp_auth", - "litellm-auth", + "google-cloud-auth", + "http 1.4.2", + "litellm-auth-types", "moka", "serde_json", "sha2 0.10.9", "tokio", ] +[[package]] +name = "litellm-auth-types" +version = "0.1.0" +dependencies = [ + "serde", + "subtle", + "thiserror 2.0.19", + "tokio", + "veil", +] + [[package]] name = "litellm-cache" version = "0.1.0" @@ -2242,6 +2664,88 @@ dependencies = [ "tokio-tungstenite", ] +[[package]] +name = "litellm-secrets" +version = "0.1.0" +dependencies = [ + "aws-sdk-kms", + "base64 0.22.1", + "google-cloud-auth", + "google-cloud-kms-v1", + "jsonwebtoken", + "litellm-core-utils", + "litellm-secrets-aws", + "litellm-secrets-google", + "litellm-secrets-types", + "moka", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "strum", + "tempfile", + "thiserror 2.0.19", + "tokio", + "wiremock", +] + +[[package]] +name = "litellm-secrets-aws" +version = "0.1.0" +dependencies = [ + "aws-credential-types", + "aws-sdk-kms", + "aws-sdk-secretsmanager", + "base64 0.22.1", + "litellm-auth-aws", + "litellm-core-utils", + "litellm-secrets-types", + "rstest", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", + "veil", + "wiremock", +] + +[[package]] +name = "litellm-secrets-google" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", + "google-cloud-auth", + "google-cloud-gax", + "google-cloud-kms-v1", + "litellm-auth-gcp", + "litellm-auth-types", + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "percent-encoding", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + +[[package]] +name = "litellm-secrets-types" +version = "0.1.0" +dependencies = [ + "litellm-auth-types", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "veil", +] + [[package]] name = "litellm-token-counter" version = "0.1.0" @@ -2477,6 +2981,16 @@ dependencies = [ "autocfg", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2489,7 +3003,7 @@ version = "6.5.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" dependencies = [ - "bitflags", + "bitflags 2.13.1", "libc", "once_cell", "onig_sys", @@ -2517,6 +3031,42 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" +[[package]] +name = "opentelemetry" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0142c63252a9e054e68a4c61a5778f7b14f576274d593f8ce883d191a099682" +dependencies = [ + "futures-core", + "futures-sink", + "js-sys", + "pin-project-lite", + "thiserror 2.0.19", + "tracing", +] + +[[package]] +name = "opentelemetry-semantic-conventions" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c913ac17a6c451661ee255f4625d143e51647ae78ebd969b75e41c4442f4fe47" + +[[package]] +name = "opentelemetry_sdk" +version = "0.32.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b59f80e1ac4d5ff7a2db8fb6c80badb7f0f3f858211fba08dd9aaec750894f9" +dependencies = [ + "futures-channel", + "futures-executor", + "futures-util", + "opentelemetry", + "percent-encoding", + "portable-atomic", + "rand 0.9.5", + "thiserror 2.0.19", +] + [[package]] name = "outref" version = "0.5.2" @@ -2652,6 +3202,15 @@ version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" +[[package]] +name = "portable-atomic-util" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715" +dependencies = [ + "portable-atomic", +] + [[package]] name = "potential_utf" version = "0.1.5" @@ -2702,7 +3261,7 @@ checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" dependencies = [ "bit-set", "bit-vec", - "bitflags", + "bitflags 2.13.1", "num-traits", "rand 0.9.5", "rand_chacha 0.9.0", @@ -2713,6 +3272,38 @@ dependencies = [ "unarray", ] +[[package]] +name = "prost" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "prost-types" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" +dependencies = [ + "prost", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -3039,7 +3630,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.1", ] [[package]] @@ -3170,6 +3761,9 @@ dependencies = [ "rustls 0.23.42", "rustls-pki-types", "rustls-platform-verifier", + "serde", + "serde_json", + "serde_urlencoded", "sync_wrapper", "tokio", "tokio-rustls 0.26.4", @@ -3259,7 +3853,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "errno", "libc", "linux-raw-sys", @@ -3285,6 +3879,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3c54fcab019b409d04215d3a17cb438fd7fbf192ee61461f20f4fe18704bc138" dependencies = [ "aws-lc-rs", + "log", "once_cell", "ring", "rustls-pki-types", @@ -3452,7 +4047,7 @@ version = "3.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" dependencies = [ - "bitflags", + "bitflags 2.13.1", "core-foundation", "core-foundation-sys", "libc", @@ -3612,6 +4207,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -3628,6 +4232,15 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -3859,6 +4472,15 @@ dependencies = [ "syn 3.0.0", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "tiktoken-rs" version = "0.12.0" @@ -4019,6 +4641,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + [[package]] name = "tokio-tungstenite" version = "0.24.0" @@ -4078,6 +4712,44 @@ dependencies = [ "winnow", ] +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "base64 0.22.1", + "bytes", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "hyper 1.10.1", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "rustls-native-certs", + "sync_wrapper", + "tokio", + "tokio-rustls 0.26.4", + "tokio-stream", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -4086,11 +4758,15 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.14.0", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", + "tracing", ] [[package]] @@ -4100,7 +4776,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "async-compression", - "bitflags", + "bitflags 2.13.1", "bytes", "futures-core", "futures-util", @@ -4157,6 +4833,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", ] [[package]] @@ -4169,6 +4846,31 @@ dependencies = [ "tracing", ] +[[package]] +name = "tracing-opentelemetry" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adbc64cba7137545b8044cb1fe9814f7aacf3c6b5f9b45be8bb5db538befdb26" +dependencies = [ + "js-sys", + "opentelemetry", + "tracing", + "tracing-core", + "tracing-subscriber", + "web-time", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "sharded-slab", + "thread_local", + "tracing-core", +] + [[package]] name = "try-lock" version = "0.2.5" @@ -4338,6 +5040,12 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "veil" version = "0.3.0" @@ -4714,6 +5422,29 @@ dependencies = [ "memchr", ] +[[package]] +name = "wiremock" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08db1edfb05d9b3c1542e521aea074442088292f00b5f28e435c714a98f85031" +dependencies = [ + "assert-json-diff", + "base64 0.22.1", + "deadpool", + "futures", + "http 1.4.2", + "http-body-util", + "hyper 1.10.1", + "hyper-util", + "log", + "once_cell", + "regex", + "serde", + "serde_json", + "tokio", + "url", +] + [[package]] name = "wit-bindgen" version = "0.57.1" @@ -4807,6 +5538,20 @@ name = "zeroize" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] [[package]] name = "zerotrie" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 4250fa81d26..2f6f5feb4ad 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -14,9 +14,14 @@ litellm-host = { path = "crates/host" } litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framing = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } +litellm-auth-types = { path = "crates/auth-types" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } +litellm-secrets = { path = "crates/secrets" } +litellm-secrets-types = { path = "crates/secrets-types" } +litellm-secrets-aws = { path = "crates/secrets-aws" } +litellm-secrets-google = { path = "crates/secrets-google" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } @@ -31,6 +36,8 @@ litellm-host-python = { path = "crates/host-python" } bytes = "1" http = "1" +google-cloud-auth = { version = "1.16.0", default-features = false } +jsonwebtoken = { version = "11.1.0", default-features = false } hyper-util = { version = "0.1.20", default-features = false, features = ["client-proxy"] } proptest = "1.7.0" pyo3 = "0.29.2" diff --git a/litellm-rust/crates/auth-aws/Cargo.toml b/litellm-rust/crates/auth-aws/Cargo.toml index 1f27c7bc990..1a35af48574 100644 --- a/litellm-rust/crates/auth-aws/Cargo.toml +++ b/litellm-rust/crates/auth-aws/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true litellm-http.workspace = true moka = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/auth-aws/src/constants.rs b/litellm-rust/crates/auth-aws/src/constants.rs index be215cc9016..9e7c6bfab43 100644 --- a/litellm-rust/crates/auth-aws/src/constants.rs +++ b/litellm-rust/crates/auth-aws/src/constants.rs @@ -3,6 +3,8 @@ pub const AWS_SECRET_ACCESS_KEY: &str = "AWS_SECRET_ACCESS_KEY"; pub const AWS_SESSION_TOKEN: &str = "AWS_SESSION_TOKEN"; pub const AWS_REGION_NAME: &str = "AWS_REGION_NAME"; pub const AWS_REGION: &str = "AWS_REGION"; +pub const AWS_DEFAULT_REGION: &str = "AWS_DEFAULT_REGION"; +pub const AWS_BEDROCK_RUNTIME_ENDPOINT: &str = "AWS_BEDROCK_RUNTIME_ENDPOINT"; pub const AWS_SESSION_NAME: &str = "AWS_SESSION_NAME"; pub const AWS_PROFILE_NAME: &str = "AWS_PROFILE_NAME"; pub const AWS_ROLE_NAME: &str = "AWS_ROLE_NAME"; diff --git a/litellm-rust/crates/auth-aws/src/error.rs b/litellm-rust/crates/auth-aws/src/error.rs index f80fbce456e..d4c6ae8cf6c 100644 --- a/litellm-rust/crates/auth-aws/src/error.rs +++ b/litellm-rust/crates/auth-aws/src/error.rs @@ -22,7 +22,7 @@ pub enum Error { AwsMissingWebIdentityCredentials, } -impl From for litellm_auth::Error { +impl From for litellm_auth_types::Error { fn from(error: Error) -> Self { Self::ProviderAuthentication(error.to_string()) } @@ -34,11 +34,11 @@ mod tests { #[test] fn converts_to_shared_auth_error_without_losing_context() { - let error = litellm_auth::Error::from(Error::AwsProfile("profile not found".into())); + let error = litellm_auth_types::Error::from(Error::AwsProfile("profile not found".into())); assert_eq!( error, - litellm_auth::Error::ProviderAuthentication( + litellm_auth_types::Error::ProviderAuthentication( "AWS profile credentials failed: profile not found".into() ) ); diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index 8099506d2e5..1fd9d39d113 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true moka.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs index ab9ffc719df..cd16b27f66d 100644 --- a/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs +++ b/litellm-rust/crates/auth-azure/src/credential_provider_cache.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use azure_core::credentials::TokenCredential; use moka::future::Cache; -use litellm_auth::Error; +use litellm_auth_types::Error; #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub(crate) struct AzureCredentialProviderCacheKey { diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 5f913a8ad01..d635e559641 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -12,8 +12,8 @@ use azure_identity::{ }; use sha2::{Digest, Sha256}; -use litellm_auth::Error; -use litellm_auth::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::Error; +use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -484,7 +484,7 @@ mod tests { use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, SecretValue, Sourced}; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) @@ -649,7 +649,7 @@ mod tests { assert!(matches!( error, - litellm_auth::Error::MixedAzureCredentialSources + litellm_auth_types::Error::MixedAzureCredentialSources )); } @@ -679,7 +679,10 @@ mod tests { authority, )) .unwrap_err(); - assert!(matches!(error, litellm_auth::Error::InvalidAzureAuthority)); + assert!(matches!( + error, + litellm_auth_types::Error::InvalidAzureAuthority + )); } } } diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 4e18cbb89aa..4d564b6e68a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,5 +1,5 @@ -use litellm_auth::Error; -use litellm_auth::{ +use litellm_auth_types::Error; +use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, SecretValue, Sourced, TokenProviderHandle, }; @@ -451,9 +451,9 @@ mod tests { }; use crate::native::ValidatedAzureRequest; use crate::types::AzureAuthInputs; - use litellm_auth::Error; - use litellm_auth::ResolvedCredential; - use litellm_auth::{ + use litellm_auth_types::Error; + use litellm_auth_types::ResolvedCredential; + use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, }; @@ -661,8 +661,8 @@ mod tests { #[derive(Debug)] struct CallerToken(&'static str); - impl litellm_auth::TokenProvider for CallerToken { - fn acquire(&self) -> litellm_auth::TokenFuture<'_> { + impl litellm_auth_types::TokenProvider for CallerToken { + fn acquire(&self) -> litellm_auth_types::TokenFuture<'_> { Box::pin(async move { Ok(ResolvedCredential::AccessToken { token: SecretValue::new(self.0), @@ -675,7 +675,7 @@ mod tests { fn caller_inputs(token: &'static str) -> AzureAuthInputs { let params = json!({"azure_ad_token": "static-token"}); AzureAuthInputs { - azure_ad_token_provider: Some(litellm_auth::TokenProviderHandle::new(Arc::new( + azure_ad_token_provider: Some(litellm_auth_types::TokenProviderHandle::new(Arc::new( CallerToken(token), ))), ..AzureAuthInputs::from_optional_params(params.as_object().unwrap()).unwrap() diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index 87e883a6a54..d5a00f09751 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use litellm_auth::{ +use litellm_auth_types::{ CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, }; use serde_json::{Map, Value}; @@ -126,7 +126,7 @@ fn source_for(sources: &BTreeMap, name: &str) -> InputSourc mod tests { use std::collections::BTreeMap; - use litellm_auth::{InputSource, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; use serde_json::json; use super::{AzureAuthInputs, AzureCredentialType, ConfigValue}; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index f24582db13e..0c6258a193c 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -5,8 +5,11 @@ edition.workspace = true license.workspace = true repository.workspace = true +[features] +google-sdk = ["dep:google-cloud-auth", "dep:http"] + [dependencies] -litellm-auth.workspace = true +litellm-auth-types.workspace = true moka.workspace = true serde_json.workspace = true @@ -14,3 +17,5 @@ sha2.workspace = true tokio.workspace = true gcp_auth = "0.12.7" +google-cloud-auth = { workspace = true, optional = true } +http = { workspace = true, optional = true } diff --git a/litellm-rust/crates/auth-gcp/src/lib.rs b/litellm-rust/crates/auth-gcp/src/lib.rs index bf619fee144..8aeddae9efc 100644 --- a/litellm-rust/crates/auth-gcp/src/lib.rs +++ b/litellm-rust/crates/auth-gcp/src/lib.rs @@ -1,13 +1,18 @@ use std::{collections::BTreeMap, future::Future, path::Path, pin::Pin, sync::Arc}; use gcp_auth::{CustomServiceAccount, TokenProvider}; -use litellm_auth::{ +use litellm_auth_types::{ CredentialPlacement, Error, InputSource, SecretValue, Sourced, http::apply_credential, }; use moka::future::Cache; use serde_json::{Map, Value}; use sha2::{Digest, Sha256}; +#[cfg(feature = "google-sdk")] +mod sdk; +#[cfg(feature = "google-sdk")] +pub use sdk::GoogleCredentials; + const CLOUD_PLATFORM_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; const GOOGLE_OAUTH_TOKEN_ENDPOINT: &str = "https://oauth2.googleapis.com/token"; const GOOGLE_APPLICATION_CREDENTIALS_ENV: &str = "GOOGLE_APPLICATION_CREDENTIALS"; @@ -26,19 +31,31 @@ pub struct VertexConfig { } impl VertexConfig { + pub fn new( + credentials: Option>, + project_id: Option, + location: Option, + ) -> Self { + Self { + credentials: credentials.filter(|value| !value.value().expose().trim().is_empty()), + project_id: project_id.filter(|value| !value.trim().is_empty()), + location: location.filter(|value| !value.trim().is_empty()), + } + } + pub fn from_sourced_optional_params( params: &Map, sources: &BTreeMap, ) -> Result { - Ok(Self { - credentials: optional_credentials( + Ok(Self::new( + optional_credentials( params, sources, &["vertex_credentials", "vertex_ai_credentials"], )?, - project_id: optional_string(params, &["vertex_project", "vertex_ai_project"])?, - location: optional_string(params, &["vertex_location", "vertex_ai_location"])?, - }) + optional_string(params, &["vertex_project", "vertex_ai_project"])?, + optional_string(params, &["vertex_location", "vertex_ai_location"])?, + )) } pub fn or_configured(self, project_id: Option<&str>, location: Option<&str>) -> Self { @@ -469,6 +486,39 @@ mod tests { assert_eq!(config.location(), Some("alias-location")); } + #[test] + fn typed_config_preserves_source_and_empty_value_fallback() { + let configured = VertexConfig::new( + Some(Sourced::new( + SecretValue::new("inline-json"), + InputSource::Request, + )), + Some("project".into()), + Some("location".into()), + ); + assert!(matches!( + credential_source(&configured, &|_| Some("environment-json".into())), + CredentialSource::Inline(value) if value.expose() == "inline-json" + )); + let empty = VertexConfig::new( + Some(Sourced::new(SecretValue::new(" "), InputSource::Request)), + Some(" ".into()), + Some(" ".into()), + ); + assert!(matches!( + credential_source(&empty, &|_| None), + CredentialSource::Adc + )); + assert_eq!( + get_vertex_ai_project(&empty, &|_| Some("env-project".into())).as_deref(), + Some("env-project") + ); + assert_eq!( + get_vertex_ai_location(&empty, &|_| Some("env-location".into())).as_deref(), + Some("env-location") + ); + } + #[test] fn project_and_location_prefer_input_then_environment() { let configured = diff --git a/litellm-rust/crates/auth-gcp/src/sdk.rs b/litellm-rust/crates/auth-gcp/src/sdk.rs new file mode 100644 index 00000000000..566df291b7d --- /dev/null +++ b/litellm-rust/crates/auth-gcp/src/sdk.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +use google_cloud_auth::credentials::{CacheableResource, CredentialsProvider, EntityTag}; +use google_cloud_auth::errors::CredentialsError; +use http::{Extensions, HeaderMap, HeaderName, HeaderValue}; +use litellm_auth_types::Error; + +use crate::{VertexAuth, VertexConfig}; + +type EnvironmentLookup = dyn Fn(&str) -> Option + Send + Sync; + +pub struct GoogleCredentials { + auth: VertexAuth, + config: VertexConfig, + environment: Arc, +} + +impl GoogleCredentials { + pub fn new(config: VertexConfig, environment: Arc) -> Self { + Self { + auth: VertexAuth::default(), + config, + environment, + } + } + + pub async fn request_headers(&self) -> Result { + let response = self + .auth + .validate_environment(Vec::new(), None, &self.config, &|name| { + (self.environment)(name) + }) + .await?; + response + .headers + .into_iter() + .map(|(key, value)| { + let name = + HeaderName::from_bytes(key.as_bytes()).map_err(|_| Error::InvalidHeader)?; + let value = HeaderValue::from_str(&value).map_err(|_| Error::InvalidHeader)?; + Ok((name, value)) + }) + .collect() + } +} + +impl CredentialsProvider for GoogleCredentials { + async fn headers( + &self, + _: Extensions, + ) -> Result, CredentialsError> { + self.request_headers() + .await + .map(|data| CacheableResource::New { + entity_tag: EntityTag::new(), + data, + }) + .map_err(|_| CredentialsError::from_msg(false, "Google authentication failed")) + } + + async fn universe_domain(&self) -> Option { + None + } +} + +impl std::fmt::Debug for GoogleCredentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GoogleCredentials").finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn sdk_and_http_credentials_share_token_resolution_and_redaction() { + let credentials = GoogleCredentials::new( + VertexConfig::new(None, Some("project".into()), None), + Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private-token".into())), + ); + let direct = credentials.request_headers().await.unwrap(); + let CacheableResource::New { data, .. } = + credentials.headers(Extensions::new()).await.unwrap() + else { + panic!("first request did not return headers"); + }; + assert_eq!(direct, data); + assert_eq!(data[http::header::AUTHORIZATION], "Bearer private-token"); + assert!(!format!("{credentials:?}").contains("private-token")); + } + + #[tokio::test] + async fn invalid_token_headers_return_a_redacted_sdk_error() { + let credentials = GoogleCredentials::new( + VertexConfig::new(None, Some("project".into()), None), + Arc::new(|name| (name == "VERTEX_AI_API_KEY").then(|| "private\nvalue".into())), + ); + assert_eq!( + credentials.request_headers().await.unwrap_err(), + Error::InvalidHeader + ); + let error = credentials.headers(Extensions::new()).await.unwrap_err(); + assert!(!format!("{error:?}").contains("private")); + } +} diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml new file mode 100644 index 00000000000..cd65412127d --- /dev/null +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "litellm-auth-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +subtle.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +tokio.workspace = true diff --git a/litellm-rust/crates/auth/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs similarity index 98% rename from litellm-rust/crates/auth/src/credential.rs rename to litellm-rust/crates/auth-types/src/credential.rs index 8ed1867622a..a5d14a43b71 100644 --- a/litellm-rust/crates/auth/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,9 +5,7 @@ use std::sync::Arc; use veil::Redact; -use crate::Error; - -use super::{ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { diff --git a/litellm-rust/crates/auth/src/error.rs b/litellm-rust/crates/auth-types/src/error.rs similarity index 100% rename from litellm-rust/crates/auth/src/error.rs rename to litellm-rust/crates/auth-types/src/error.rs diff --git a/litellm-rust/crates/auth/src/http.rs b/litellm-rust/crates/auth-types/src/http.rs similarity index 92% rename from litellm-rust/crates/auth/src/http.rs rename to litellm-rust/crates/auth-types/src/http.rs index dd87d00e70f..0cb5839f965 100644 --- a/litellm-rust/crates/auth/src/http.rs +++ b/litellm-rust/crates/auth-types/src/http.rs @@ -40,9 +40,6 @@ pub fn apply_credential( ) } -/// How the upstream call is authenticated. API-key strategies become headers -/// in `prepare`; SigV4 covers the serialized body, so it is applied where the -/// outbound request is built. #[derive(Clone, Debug, PartialEq, Eq)] pub enum RequestAuth { Header { diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs new file mode 100644 index 00000000000..9d399249c05 --- /dev/null +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -0,0 +1,57 @@ +#![forbid(unsafe_code)] + +mod credential; +mod error; +pub mod http; +mod policy; +mod secret; +mod token; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum InputSource { + Request, + #[default] + Deployment, + Environment, +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct Sourced { + value: T, + source: InputSource, +} + +impl Sourced { + pub fn new(value: T, source: InputSource) -> Self { + Self { value, source } + } + + pub fn value(&self) -> &T { + &self.value + } + + pub fn source(&self) -> InputSource { + self.source + } + + pub fn into_value(self) -> T { + self.value + } + + pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { + Sourced::new(map(self.value), self.source) + } +} + +pub use credential::{ + CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, + CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, +}; +pub use error::Error; +pub use http::{CredentialPlacement, RequestAuth}; +pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; +pub use secret::SecretValue; +pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs similarity index 96% rename from litellm-rust/crates/auth/src/policy.rs rename to litellm-rust/crates/auth-types/src/policy.rs index 4a1f5eeecf9..4c5c0365f0b 100644 --- a/litellm-rust/crates/auth/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -1,7 +1,5 @@ -use crate::Error; - -use super::http::apply_credential; -use super::{CredentialPlacement, ResolvedCredential}; +use crate::http::apply_credential; +use crate::{CredentialPlacement, Error, ResolvedCredential}; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CredentialPlanKind { diff --git a/litellm-rust/crates/auth/src/secret.rs b/litellm-rust/crates/auth-types/src/secret.rs similarity index 100% rename from litellm-rust/crates/auth/src/secret.rs rename to litellm-rust/crates/auth-types/src/secret.rs diff --git a/litellm-rust/crates/auth/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs similarity index 95% rename from litellm-rust/crates/auth/src/token.rs rename to litellm-rust/crates/auth-types/src/token.rs index 94da5f259fb..4175641ce10 100644 --- a/litellm-rust/crates/auth/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,9 +5,7 @@ use std::time::SystemTime; use veil::Redact; -use crate::Error; - -use super::secret::SecretValue; +use crate::{Error, SecretValue}; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index 128a05c1a25..ee4900ebcc2 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -5,11 +5,14 @@ edition.workspace = true license.workspace = true repository.workspace = true -[dependencies] -serde.workspace = true -subtle.workspace = true -thiserror.workspace = true -veil.workspace = true +[features] +default = [] +aws = ["dep:litellm-auth-aws"] +azure = ["dep:litellm-auth-azure"] +gcp = ["dep:litellm-auth-gcp"] -[dev-dependencies] -tokio.workspace = true +[dependencies] +litellm-auth-types.workspace = true +litellm-auth-aws = { workspace = true, optional = true } +litellm-auth-azure = { workspace = true, optional = true } +litellm-auth-gcp = { workspace = true, optional = true } diff --git a/litellm-rust/crates/auth/src/lib.rs b/litellm-rust/crates/auth/src/lib.rs index c8d73c239b0..622a5b2d58b 100644 --- a/litellm-rust/crates/auth/src/lib.rs +++ b/litellm-rust/crates/auth/src/lib.rs @@ -1,55 +1,10 @@ -mod credential; -mod error; -pub mod http; -mod policy; -mod secret; -mod token; +#![forbid(unsafe_code)] -use serde::{Deserialize, Serialize}; +pub use litellm_auth_types::*; -#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)] -#[serde(rename_all = "snake_case")] -pub enum InputSource { - Request, - #[default] - Deployment, - Environment, -} - -#[derive(Clone, Debug, Default, Eq, PartialEq)] -pub struct Sourced { - value: T, - source: InputSource, -} - -impl Sourced { - pub fn new(value: T, source: InputSource) -> Self { - Self { value, source } - } - - pub fn value(&self) -> &T { - &self.value - } - - pub fn source(&self) -> InputSource { - self.source - } - - pub fn into_value(self) -> T { - self.value - } - - pub fn map(self, map: impl FnOnce(T) -> U) -> Sourced { - Sourced::new(map(self.value), self.source) - } -} - -pub use credential::{ - CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialPlan, - CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, -}; -pub use error::Error; -pub use http::{CredentialPlacement, RequestAuth}; -pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; -pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +#[cfg(feature = "aws")] +pub use litellm_auth_aws as aws; +#[cfg(feature = "azure")] +pub use litellm_auth_azure as azure; +#[cfg(feature = "gcp")] +pub use litellm_auth_gcp as gcp; diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs new file mode 100644 index 00000000000..f1092b15def --- /dev/null +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -0,0 +1,33 @@ +use litellm_auth::{ + CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, + ProviderAuthPolicy, ResolvedCredential, SecretValue, +}; + +const RULES: &[CredentialRule] = &[CredentialRule { + kind: CredentialPlanKind::Static, + placement: CredentialPlacement::Header("x-api-key"), +}]; + +#[test] +fn facade_applies_shared_auth_policy() { + let policy = ProviderAuthPolicy { + rules: RULES, + accepted_existing_headers: &["x-api-key"], + existing_header_behavior: ExistingHeaderBehavior::Preserve, + scope: None, + audience: None, + }; + + let headers = policy + .apply( + Vec::new(), + CredentialPlanKind::Static, + &ResolvedCredential::Static(SecretValue::new("secret")), + ) + .expect("facade policy applies"); + + assert_eq!( + headers, + vec![("x-api-key".to_string(), "secret".to_string())] + ); +} diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml new file mode 100644 index 00000000000..b1dc5b33cda --- /dev/null +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-secrets-aws" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-aws.workspace = true +litellm-secrets-types.workspace = true +litellm-core-utils.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tracing = "0.1" +veil.workspace = true +aws-sdk-kms = "1.120.0" +aws-sdk-secretsmanager = "1.117.0" +aws-credential-types = "1.3.0" + +[dev-dependencies] +base64.workspace = true +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-aws/src/auth.rs b/litellm-rust/crates/secrets-aws/src/auth.rs new file mode 100644 index 00000000000..954cfa2f8fd --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/auth.rs @@ -0,0 +1,79 @@ +use std::sync::Arc; + +use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; +use litellm_auth_aws::{ + AwsAuthConfig, + constants::{AWS_DEFAULT_REGION, AWS_REGION, AWS_REGION_NAME}, + resolve_credentials, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; + +use crate::Error; + +#[derive(Clone)] +pub(crate) struct Credentials { + config: AwsAuthConfig, + environment: Arc, +} + +impl Credentials { + pub(crate) fn new( + settings: &KeyManagementSettings, + environment: Arc, + ) -> Self { + Self { + config: AwsAuthConfig { + region_name: region(settings, environment.as_ref()).ok(), + role_name: settings.aws_role_name.clone(), + session_name: settings.aws_session_name.clone(), + external_id: settings + .aws_external_id + .as_ref() + .map(|v| v.expose().to_owned()), + profile_name: settings.aws_profile_name.clone(), + web_identity_token: settings + .aws_web_identity_token + .as_ref() + .map(|v| v.expose().to_owned()), + sts_endpoint: settings.aws_sts_endpoint.clone(), + ..Default::default() + }, + environment, + } + } +} + +impl ProvideCredentials for Credentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::new(async { + resolve_credentials(self.config.clone(), &|name| self.environment.get(name)) + .await + .map_err(|_| { + CredentialsError::provider_error("secret manager authentication failed") + }) + }) + } +} + +pub(crate) fn region( + settings: &KeyManagementSettings, + environment: &dyn Lookup, +) -> Result { + settings + .aws_region_name + .clone() + .or_else(|| environment.get(AWS_REGION_NAME)) + .or_else(|| environment.get(AWS_REGION)) + .or_else(|| environment.get(AWS_DEFAULT_REGION)) + .ok_or(Error::MissingRegion) +} + +impl std::fmt::Debug for Credentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Credentials").finish_non_exhaustive() + } +} diff --git a/litellm-rust/crates/secrets-aws/src/error.rs b/litellm-rust/crates/secrets-aws/src/error.rs new file mode 100644 index 00000000000..23595397a13 --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/error.rs @@ -0,0 +1,31 @@ +use aws_sdk_secretsmanager::error::SdkError; + +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("AWS authentication failed")] + Auth(#[from] #[redact] litellm_auth_aws::Error), + #[error("AWS region is not configured")] + MissingRegion, + #[error("KMS response has no plaintext")] + MissingPlaintext, + #[error("AWS request timed out")] + Timeout, + #[error("AWS KMS decrypt failed")] + Decrypt(#[from] #[redact] Box>), + #[error("AWS Secrets Manager read failed")] + Read(#[from] #[redact] Box>), + #[error("AWS Secrets Manager create failed")] + Create(#[from] #[redact] Box>), + #[error("AWS Secrets Manager update failed")] + Put(#[from] #[redact] Box>), + #[error("AWS Secrets Manager delete failed")] + Delete(#[from] #[redact] Box>), + #[error("AWS Secrets Manager replication failed")] + Replicate(#[from] #[redact] Box>), + #[error("AWS Secrets Manager response has no string payload")] + MissingString, + #[error("primary secret is not a JSON object")] + PrimarySecret, + #[error(transparent)] + Operation(#[from] litellm_secrets_types::Error), +} diff --git a/litellm-rust/crates/secrets-aws/src/kms.rs b/litellm-rust/crates/secrets-aws/src/kms.rs new file mode 100644 index 00000000000..a66b1c4d2fe --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/kms.rs @@ -0,0 +1,63 @@ +use litellm_auth_aws::constants::AWS_REGION_NAME; +use std::sync::Arc; + +use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Region}, + primitives::Blob, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::KeyManagementSettings; + +use crate::{Error, auth}; + +#[derive(Clone)] +pub struct AwsKms { + client: Client, +} + +impl AwsKms { + pub fn new(client: Client) -> Self { + Self { client } + } + + pub async fn decrypt(&self, ciphertext: Vec) -> Result, Error> { + let response = self + .client + .decrypt() + .ciphertext_blob(Blob::new(ciphertext)) + .send() + .await + .map_err(|error| Error::Decrypt(Box::new(error)))?; + Ok(response + .plaintext + .ok_or(Error::MissingPlaintext)? + .into_inner()) + } +} + +pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> { + environment + .get(AWS_REGION_NAME) + .map(|_| ()) + .ok_or(Error::MissingRegion) +} + +pub fn load_aws_kms( + use_aws_kms: Option, + settings: &KeyManagementSettings, + environment: Arc, +) -> Result, Error> { + if use_aws_kms != Some(true) { + return Ok(None); + } + if settings.aws_region_name.is_none() { + validate_environment(environment.as_ref())?; + } + let config = aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(auth::region(settings, environment.as_ref())?)) + .credentials_provider(auth::Credentials::new(settings, environment)) + .build(); + Ok(Some(AwsKms::new(Client::from_conf(config)))) +} diff --git a/litellm-rust/crates/secrets-aws/src/lib.rs b/litellm-rust/crates/secrets-aws/src/lib.rs new file mode 100644 index 00000000000..d17eb38c7eb --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/lib.rs @@ -0,0 +1,10 @@ +#![forbid(unsafe_code)] + +mod auth; +mod error; +pub mod kms; +pub mod secret_manager; + +pub use error::Error; +pub use kms::{AwsKms, load_aws_kms}; +pub use secret_manager::{AwsSecretWriteSettings, AwsSecretsManagerV2, RotationResponse}; diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs new file mode 100644 index 00000000000..493cb1d2e8f --- /dev/null +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -0,0 +1,287 @@ +use litellm_auth_aws::constants::AWS_BEDROCK_RUNTIME_ENDPOINT; +use std::{collections::BTreeMap, sync::Arc}; + +use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Region}, + operation::{ + create_secret::CreateSecretOutput, delete_secret::DeleteSecretOutput, + put_secret_value::PutSecretValueOutput, + replicate_secret_to_regions::ReplicateSecretToRegionsOutput, + }, + types::{ReplicaRegionType, Tag}, +}; +use litellm_auth_aws::constants::{ + AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY, +}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{ + BaseSecretManager, KeyManagementSettings, Secret, SecretValue, async_rotate_secret, +}; +use serde_json::Value; + +use crate::{Error, auth}; + +#[derive(Clone)] +pub struct AwsSecretsManagerV2 { + client: Client, + write_settings: AwsSecretWriteSettings, +} + +#[derive(Clone, Debug, Default)] +pub struct AwsSecretWriteSettings { + pub kms_key_id: Option, + pub tags: Option>, + pub replica_regions: Option>, +} + +impl From<&KeyManagementSettings> for AwsSecretWriteSettings { + fn from(settings: &KeyManagementSettings) -> Self { + Self { + kms_key_id: settings.kms_key_id.clone(), + tags: settings.tags.clone(), + replica_regions: settings.replica_regions.clone(), + } + } +} + +#[derive(Debug)] +pub enum RotationResponse { + Created(CreateSecretOutput), + Updated(PutSecretValueOutput), +} + +impl AwsSecretsManagerV2 { + pub fn new(client: Client, write_settings: AwsSecretWriteSettings) -> Self { + Self { + client, + write_settings, + } + } + + pub fn load_aws_secret_manager( + use_aws_secret_manager: Option, + settings: KeyManagementSettings, + environment: Arc, + ) -> Result, Error> { + if use_aws_secret_manager != Some(true) { + return Ok(None); + } + let builder = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(auth::region(&settings, environment.as_ref())?)) + .credentials_provider(auth::Credentials::new(&settings, environment.clone())); + let config = match environment.get(AWS_BEDROCK_RUNTIME_ENDPOINT) { + Some(url) => builder + .endpoint_url(url.replace("bedrock-runtime", "secretsmanager")) + .build(), + None => builder.build(), + }; + Ok(Some(Self::new( + Client::from_conf(config), + (&settings).into(), + ))) + } + + pub async fn read_secret_for_resolver( + &self, + name: &str, + primary_name: Option<&str>, + environment: &(dyn Lookup + Sync), + ) -> Result, Error> { + if bootstrap_key(name) { + return Ok(environment + .get(name) + .map(SecretValue::new) + .map(Secret::String)); + } + match primary_name.filter(|name| !name.is_empty()) { + None => self + .async_read_secret(name) + .await + .map(|value| value.map(Secret::String)), + Some(primary) => { + let value = if bootstrap_key(primary) { + environment.get(primary).map(SecretValue::new) + } else { + self.async_read_secret(primary).await? + }; + 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().map(Secret::from_json)) + } + } + } + + pub async fn async_read_secret(&self, name: &str) -> Result, Error> { + match self.client.get_secret_value().secret_id(name).send().await { + Ok(response) => response + .secret_string + .map(SecretValue::new) + .map(Some) + .ok_or(Error::MissingString), + Err(error) + if matches!( + &error, + aws_sdk_secretsmanager::error::SdkError::TimeoutError(_) + ) || matches!(&error, aws_sdk_secretsmanager::error::SdkError::DispatchFailure(failure) if failure.is_timeout()) => + { + Err(Error::Timeout) + } + 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))), + } + } + + pub async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + let response = self + .client + .create_secret() + .name(name) + .secret_string(value.expose()) + .set_description(description.filter(|v| !v.is_empty()).map(str::to_owned)) + .set_kms_key_id( + self.write_settings + .kms_key_id + .clone() + .filter(|v| !v.is_empty()), + ) + .set_tags(self.write_settings.tags.as_ref().map(|tags| { + tags.iter() + .map(|(key, value)| Tag::builder().key(key).value(value).build()) + .collect() + })) + .send() + .await + .map_err(|error| Error::Create(Box::new(error)))?; + if let Some(regions) = &self.write_settings.replica_regions + && !regions.is_empty() + && self.async_replicate_secret(name, regions).await.is_err() + { + tracing::warn!("secret created but replication failed"); + } + Ok(response) + } + + pub async fn async_replicate_secret( + &self, + name: &str, + regions: &[String], + ) -> Result, Error> { + if regions.is_empty() { + return Ok(None); + } + self.client + .replicate_secret_to_regions() + .secret_id(name) + .set_add_replica_regions(Some( + regions + .iter() + .map(|region| ReplicaRegionType::builder().region(region).build()) + .collect(), + )) + .send() + .await + .map(Some) + .map_err(|error| Error::Replicate(Box::new(error))) + } + + pub async fn async_put_secret_value( + &self, + name: &str, + value: &SecretValue, + ) -> Result { + self.client + .put_secret_value() + .secret_id(name) + .secret_string(value.expose()) + .send() + .await + .map_err(|error| Error::Put(Box::new(error))) + } + + pub async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.client + .delete_secret() + .secret_id(name) + .recovery_window_in_days(recovery_window_in_days) + .send() + .await + .map_err(|error| Error::Delete(Box::new(error))) + } + + pub async fn async_rotate_secret( + &self, + current_name: &str, + new_name: &str, + value: &SecretValue, + ) -> Result { + if current_name == new_name { + return self + .async_put_secret_value(current_name, value) + .await + .map(RotationResponse::Updated); + } + async_rotate_secret(self, current_name, new_name, value) + .await + .map(RotationResponse::Created) + } +} + +impl BaseSecretManager for AwsSecretsManagerV2 { + type Error = Error; + type WriteResponse = CreateSecretOutput; + type DeleteResponse = DeleteSecretOutput; + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + self.async_read_secret(name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + self.async_write_secret(name, value, description).await + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.async_delete_secret(name, recovery_window_in_days) + .await + } +} + +fn bootstrap_key(name: &str) -> bool { + matches!( + name, + AWS_ACCESS_KEY_ID + | AWS_SECRET_ACCESS_KEY + | AWS_REGION_NAME + | AWS_REGION + | AWS_BEDROCK_RUNTIME_ENDPOINT + ) +} diff --git a/litellm-rust/crates/secrets-aws/tests/kms.rs b/litellm-rust/crates/secrets-aws/tests/kms.rs new file mode 100644 index 00000000000..39a50297551 --- /dev/null +++ b/litellm-rust/crates/secrets-aws/tests/kms.rs @@ -0,0 +1,59 @@ +use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, +}; +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_aws::AwsKms; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, header}, +}; + +#[tokio::test] +async fn kms_decrypt_calls_the_sdk_without_applying_lookup_policy() { + let server = MockServer::start().await; + let plaintext = " private-value\n"; + Mock::given(header("x-amz-target", "TrentService.Decrypt")) + .and(body_json( + serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"Plaintext": STANDARD.encode(plaintext)})), + ) + .expect(1) + .mount(&server) + .await; + let client = Client::from_conf( + aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + let manager = AwsKms::new(client); + assert_eq!( + manager.decrypt(b"encrypted".to_vec()).await.unwrap(), + plaintext.as_bytes() + ); +} + +#[test] +fn disabled_kms_loader_does_not_require_environment_configuration() { + use litellm_secrets_aws::load_aws_kms; + use litellm_secrets_types::KeyManagementSettings; + use std::sync::Arc; + for enabled in [None, Some(false)] { + assert!( + load_aws_kms( + enabled, + &KeyManagementSettings::default(), + Arc::new(|_: &str| None) + ) + .unwrap() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs new file mode 100644 index 00000000000..a410767cb5a --- /dev/null +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -0,0 +1,312 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use aws_sdk_secretsmanager::{ + Client, + config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, +}; +use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse}; +use litellm_secrets_types::{KeyManagementSettings, SecretValue}; +use serde_json::json; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header}, +}; + +fn manager(server: &MockServer, settings: KeyManagementSettings) -> AwsSecretsManagerV2 { + let client = Client::from_conf( + aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(), + ); + AwsSecretsManagerV2::new(client, (&settings).into()) +} + +#[rstest::rstest] +#[case::string_value("KEY", Some("value"))] +#[case::missing_value("missing", None)] +#[case::non_string_value("BOOL", None)] +#[tokio::test] +async fn primary_lookup_preserves_read_semantics( + #[case] name: &str, + #[case] expected: Option<&str>, +) { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.GetSecretValue")) + .and(body_partial_json(json!({"SecretId":"primary"}))) + .respond_with( + ResponseTemplate::new(200).set_body_json( + json!({"SecretString":json!({"KEY":"value", "BOOL":true}).to_string()}), + ), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, KeyManagementSettings::default()); + assert_eq!( + manager + .read_secret_for_resolver(name, Some("primary"), &|_: &str| None) + .await + .unwrap() + .and_then(|v| v.as_str().map(str::to_owned)) + .as_deref(), + expected + ); +} + +#[rstest::rstest] +#[case::access_key("AWS_ACCESS_KEY_ID")] +#[case::secret_access_key("AWS_SECRET_ACCESS_KEY")] +#[case::region_name("AWS_REGION_NAME")] +#[case::region("AWS_REGION")] +#[case::bedrock_endpoint("AWS_BEDROCK_RUNTIME_ENDPOINT")] +#[tokio::test] +async fn bootstrap_keys_bypass_primary_lookup(#[case] name: &str) { + let server = MockServer::start().await; + let manager = manager(&server, KeyManagementSettings::default()); + assert_eq!( + manager + .read_secret_for_resolver(name, Some("primary"), &|_: &str| Some("bootstrap".into())) + .await + .unwrap() + .unwrap() + .as_str() + .unwrap(), + "bootstrap" + ); +} + +#[tokio::test] +async fn failed_read_returns_none_but_invalid_primary_json_is_an_error() { + let server = MockServer::start().await; + Mock::given(body_partial_json(json!({"SecretId":"missing"}))) + .respond_with( + ResponseTemplate::new(400).set_body_json(json!({"__type":"ResourceNotFoundException"})), + ) + .mount(&server) + .await; + Mock::given(body_partial_json(json!({"SecretId":"invalid"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"SecretString":"not-json"}))) + .mount(&server) + .await; + let manager = manager(&server, KeyManagementSettings::default()); + assert!( + manager + .async_read_secret("missing") + .await + .unwrap() + .is_none() + ); + assert!(matches!( + manager + .read_secret_for_resolver("KEY", Some("invalid"), &|_: &str| None) + .await, + Err(Error::PrimarySecret) + )); +} + +#[tokio::test] +async fn same_name_rotation_uses_put_and_returns_its_response() { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.PutSecretValue")) + .and(body_partial_json( + json!({"SecretId":"key", "SecretString":"replacement"}), + )) + .respond_with( + ResponseTemplate::new(200).set_body_json(json!({"Name":"key", "VersionId":"version"})), + ) + .expect(1) + .mount(&server) + .await; + let response = manager(&server, KeyManagementSettings::default()) + .async_rotate_secret("key", "key", &SecretValue::new("replacement")) + .await + .unwrap(); + match response { + RotationResponse::Updated(output) => assert_eq!(output.version_id(), Some("version")), + _ => panic!("rotation created a second secret"), + } + assert_eq!(server.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn renamed_rotation_reads_creates_verifies_then_deletes() { + let server = MockServer::start().await; + let step = AtomicUsize::new(0); + Mock::given(wiremock::matchers::method("POST")) + .respond_with(move |request: &wiremock::Request| { + let body: serde_json::Value = request.body_json().unwrap(); + let action = request + .headers + .get("x-amz-target") + .unwrap() + .to_str() + .unwrap(); + match step.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert_eq!(action, "secretsmanager.GetSecretValue"); + assert_eq!(body["SecretId"], "old"); + ResponseTemplate::new(200).set_body_json(json!({"SecretString":"old-value"})) + } + 1 => { + assert_eq!(action, "secretsmanager.CreateSecret"); + assert_eq!(body["Name"], "new"); + assert_eq!(body["Description"], "Rotated from old"); + assert_eq!(body["SecretString"], "replacement"); + ResponseTemplate::new(200).set_body_json(json!({"Name":"new"})) + } + 2 => { + assert_eq!(action, "secretsmanager.GetSecretValue"); + assert_eq!(body["SecretId"], "new"); + ResponseTemplate::new(200).set_body_json(json!({"SecretString":"replacement"})) + } + 3 => { + assert_eq!(action, "secretsmanager.DeleteSecret"); + assert_eq!(body["SecretId"], "old"); + assert_eq!(body["RecoveryWindowInDays"], 7); + ResponseTemplate::new(200).set_body_json(json!({"Name":"old"})) + } + _ => panic!("unexpected request"), + } + }) + .expect(4) + .mount(&server) + .await; + assert!(matches!( + manager(&server, KeyManagementSettings::default()) + .async_rotate_secret("old", "new", &SecretValue::new("replacement")) + .await + .unwrap(), + RotationResponse::Created(_) + )); +} + +#[tokio::test] +async fn creation_passes_tags_and_kms_and_survives_replication_failure() { + let server = MockServer::start().await; + Mock::given(header("x-amz-target", "secretsmanager.CreateSecret")) + .and(body_partial_json(json!({"Name":"key", "SecretString":"value", "KmsKeyId":"kms-key", "Tags":[{"Key":"stage", "Value":"test"}]}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"Name":"key"}))).expect(1).mount(&server).await; + Mock::given(header( + "x-amz-target", + "secretsmanager.ReplicateSecretToRegions", + )) + .and(body_partial_json( + json!({"SecretId":"key", "AddReplicaRegions":[{"Region":"replica-region"}]}), + )) + .respond_with( + ResponseTemplate::new(400).set_body_json(json!({"__type":"InvalidRequestException"})), + ) + .expect(1) + .mount(&server) + .await; + let settings = KeyManagementSettings { + kms_key_id: Some("kms-key".into()), + tags: Some(std::collections::BTreeMap::from([( + "stage".into(), + "test".into(), + )])), + replica_regions: Some(vec!["replica-region".into()]), + ..Default::default() + }; + let manager = manager(&server, settings); + assert_eq!( + manager + .async_write_secret("key", &SecretValue::new("value"), None) + .await + .unwrap() + .name(), + Some("key") + ); + assert!( + manager + .async_replicate_secret("key", &[]) + .await + .unwrap() + .is_none() + ); +} + +#[tokio::test] +async fn credential_failures_are_not_swallowed_as_missing_secrets() { + use aws_credential_types::provider::{ProvideCredentials, error::CredentialsError, future}; + #[derive(Debug)] + struct FailedCredentials; + impl ProvideCredentials for FailedCredentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::ready(Err(CredentialsError::provider_error( + "private-auth-detail", + ))) + } + } + let server = MockServer::start().await; + let config = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(FailedCredentials) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .build(); + let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); + let error = manager.async_read_secret("key").await.unwrap_err(); + assert!(!format!("{error:?}").contains("private-auth-detail")); + assert!(matches!(error, Error::Read(_))); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn read_timeout_is_an_error_and_cannot_be_mistaken_for_missing() { + use std::time::Duration; + let server = MockServer::start().await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_delay(Duration::from_secs(1)) + .set_body_json(json!({"SecretString":"late"})), + ) + .mount(&server) + .await; + let config = aws_sdk_secretsmanager::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .retry_config(RetryConfig::disabled()) + .timeout_config( + aws_sdk_secretsmanager::config::timeout::TimeoutConfig::builder() + .operation_timeout(Duration::from_millis(30)) + .build(), + ) + .build(); + let manager = AwsSecretsManagerV2::new(Client::from_conf(config), Default::default()); + assert!(matches!( + manager.async_read_secret("key").await, + 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(_)) + )); +} diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml new file mode 100644 index 00000000000..daecf20ff9e --- /dev/null +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "litellm-secrets-google" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-gcp = { workspace = true, features = ["google-sdk"] } +litellm-secrets-types.workspace = true +litellm-auth-types.workspace = true +litellm-core-utils.workspace = true +base64.workspace = true +serde_json.workspace = true +thiserror.workspace = true +moka.workspace = true +veil.workspace = true +google-cloud-kms-v1 = "1.14.0" +google-cloud-gax = { version = "1.14.0", default-features = false } +percent-encoding = "2.3" +serde.workspace = true +reqwest.workspace = true + +[dev-dependencies] +google-cloud-auth.workspace = true +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" diff --git a/litellm-rust/crates/secrets-google/src/auth.rs b/litellm-rust/crates/secrets-google/src/auth.rs new file mode 100644 index 00000000000..45fc99d8d5f --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/auth.rs @@ -0,0 +1,21 @@ +use std::sync::Arc; + +use litellm_auth_gcp::{GoogleCredentials, VertexConfig}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::SecretValue; + +pub(crate) fn credentials( + project: Option, + credentials: Option, + environment: Arc, +) -> GoogleCredentials { + GoogleCredentials::new( + VertexConfig::new( + credentials.map(|value| Sourced::new(value, InputSource::Environment)), + project, + None, + ), + Arc::new(move |name| environment.get(name)), + ) +} diff --git a/litellm-rust/crates/secrets-google/src/error.rs b/litellm-rust/crates/secrets-google/src/error.rs new file mode 100644 index 00000000000..a94cc8c2de6 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/error.rs @@ -0,0 +1,43 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("Google KMS client configuration failed")] + Client( + #[from] + #[redact] + google_cloud_gax::client_builder::Error, + ), + #[error("Google authentication failed")] + Auth( + #[from] + #[redact] + litellm_auth_types::Error, + ), + #[error("Google KMS request failed")] + Kms( + #[from] + #[redact] + google_cloud_gax::error::Error, + ), + #[error("Google Secret Manager HTTP request failed")] + Http( + #[from] + #[redact] + reqwest::Error, + ), + #[error("Google Secret Manager returned HTTP {0}")] + Status(u16), + #[error("Google Secret Manager returned no payload")] + MissingPayload, + #[error("required environment variable is missing: {0}")] + MissingEnvironment(&'static str), + #[error("invalid refresh interval")] + RefreshInterval, + #[error("payload is not valid base64")] + Base64(#[from] base64::DecodeError), + #[error("decrypted value is not UTF-8")] + Utf8, + #[error("invalid Google Secret Manager endpoint")] + Endpoint, + #[error("Google Secret Manager requires an enterprise license")] + EnterpriseRequired, +} diff --git a/litellm-rust/crates/secrets-google/src/kms.rs b/litellm-rust/crates/secrets-google/src/kms.rs new file mode 100644 index 00000000000..3a247edaa35 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/kms.rs @@ -0,0 +1,67 @@ +use std::sync::Arc; + +use google_cloud_kms_v1::client::KeyManagementService; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::SecretValue; + +use crate::{Error, auth}; + +const GOOGLE_APPLICATION_CREDENTIALS: &str = "GOOGLE_APPLICATION_CREDENTIALS"; +const GOOGLE_KMS_RESOURCE_NAME: &str = "GOOGLE_KMS_RESOURCE_NAME"; + +#[derive(Clone)] +pub struct GoogleKms { + client: KeyManagementService, + resource_name: String, +} + +impl GoogleKms { + pub fn new(client: KeyManagementService, resource_name: String) -> Self { + Self { + client, + resource_name, + } + } + + pub async fn decrypt(&self, ciphertext: Vec) -> Result, Error> { + let response = self + .client + .decrypt() + .set_name(&self.resource_name) + .set_ciphertext(ciphertext) + .send() + .await?; + Ok(response.plaintext.to_vec()) + } +} + +pub fn validate_environment(environment: &dyn Lookup) -> Result<(), Error> { + for key in [GOOGLE_APPLICATION_CREDENTIALS, GOOGLE_KMS_RESOURCE_NAME] { + if environment.get(key).is_none() { + return Err(Error::MissingEnvironment(key)); + } + } + Ok(()) +} + +pub async fn load_google_kms( + use_google_kms: Option, + environment: Arc, +) -> Result, Error> { + if use_google_kms != Some(true) { + return Ok(None); + } + validate_environment(environment.as_ref())?; + let credentials = environment + .get(GOOGLE_APPLICATION_CREDENTIALS) + .ok_or(Error::MissingEnvironment(GOOGLE_APPLICATION_CREDENTIALS))?; + let resource_name = environment + .get(GOOGLE_KMS_RESOURCE_NAME) + .ok_or(Error::MissingEnvironment(GOOGLE_KMS_RESOURCE_NAME))?; + let credentials = auth::credentials(None, Some(SecretValue::new(credentials)), environment); + let client = KeyManagementService::builder() + .with_credentials(credentials) + .build() + .await?; + Ok(Some(GoogleKms::new(client, resource_name))) +} diff --git a/litellm-rust/crates/secrets-google/src/lib.rs b/litellm-rust/crates/secrets-google/src/lib.rs new file mode 100644 index 00000000000..a664f11a018 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/lib.rs @@ -0,0 +1,10 @@ +#![forbid(unsafe_code)] + +mod auth; +mod error; +pub mod kms; +pub mod secret_manager; + +pub use error::Error; +pub use kms::{GoogleKms, load_google_kms}; +pub use secret_manager::GoogleSecretManager; diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs new file mode 100644 index 00000000000..3c34d9cbcc4 --- /dev/null +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -0,0 +1,157 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{Secret, SecretValue}; +use moka::future::Cache; +use serde::Deserialize; + +use litellm_auth_gcp::GoogleCredentials; + +use crate::{Error, auth}; + +const GOOGLE_SECRET_MANAGER_PROJECT_ID: &str = "GOOGLE_SECRET_MANAGER_PROJECT_ID"; +const GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL: &str = "GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL"; +const SECRET_MANAGER_REFRESH_INTERVAL: &str = "SECRET_MANAGER_REFRESH_INTERVAL"; +const GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER: &str = + "GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER"; +const GCS_PATH_SERVICE_ACCOUNT: &str = "GCS_PATH_SERVICE_ACCOUNT"; +const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(86400); +const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(600); +const CACHE_CAPACITY: u64 = 200; + +#[derive(Clone)] +pub struct GoogleSecretManager { + client: reqwest::Client, + credentials: Arc, + endpoint: reqwest::Url, + project: String, + cache: Cache, + always_read: bool, +} + +#[derive(Deserialize)] +struct Response { + payload: Option, +} + +#[derive(Deserialize)] +struct Payload { + data: Option, +} + +impl GoogleSecretManager { + pub fn with_client( + client: reqwest::Client, + endpoint: reqwest::Url, + project: String, + environment: Arc, + refresh_interval: Option, + always_read: bool, + ) -> Result { + let credentials = auth::credentials( + Some(project.clone()), + environment + .get(GCS_PATH_SERVICE_ACCOUNT) + .map(SecretValue::new), + environment, + ); + let ttl = refresh_interval + .filter(|ttl| !ttl.is_zero()) + .unwrap_or(DEFAULT_CACHE_TTL); + let cache = Cache::builder() + .max_capacity(CACHE_CAPACITY) + .time_to_live(ttl) + .build(); + Ok(Self { + client, + credentials: Arc::new(credentials), + endpoint, + project, + cache, + always_read, + }) + } + + pub fn new( + environment: Arc, + enterprise_enabled: bool, + ) -> Result { + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let project = environment + .get(GOOGLE_SECRET_MANAGER_PROJECT_ID) + .ok_or(Error::MissingEnvironment(GOOGLE_SECRET_MANAGER_PROJECT_ID))?; + let ttl = environment + .get(GOOGLE_SECRET_MANAGER_REFRESH_INTERVAL) + .filter(|v| !v.is_empty()) + .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .transpose()? + .unwrap_or( + environment + .get(SECRET_MANAGER_REFRESH_INTERVAL) + .map(|v| v.parse::().map_err(|_| Error::RefreshInterval)) + .transpose()? + .unwrap_or(DEFAULT_REFRESH_INTERVAL.as_secs() as i64), + ); + let always_read = environment + .get(GOOGLE_SECRET_MANAGER_ALWAYS_READ_SECRET_MANAGER) + .is_some_and(|v| v.eq_ignore_ascii_case("true")); + Self::with_client( + reqwest::Client::new(), + reqwest::Url::parse("https://secretmanager.googleapis.com").expect("static URL"), + project, + environment, + Some(if ttl < 0 { + Duration::from_nanos(1) + } else { + Duration::from_secs(ttl as u64) + }), + always_read, + ) + } + + pub async fn get_secret_from_google_secret_manager( + &self, + name: &str, + ) -> Result, Error> { + if !self.always_read + && let Some(cached) = self.cache.get(name).await + { + return Ok(Some(Secret::String(cached))); + } + let url = self + .endpoint + .join(&format!( + "/v1/projects/{}/secrets/{}/versions/latest:access", + percent_encoding::utf8_percent_encode( + &self.project, + percent_encoding::NON_ALPHANUMERIC + ), + percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) + )) + .map_err(|_| Error::Endpoint)?; + let response = self + .client + .get(url) + .headers(self.credentials.request_headers().await?) + .send() + .await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if response.status() != reqwest::StatusCode::OK { + 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 { + return Err(Error::MissingPayload); + }; + 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(), value.clone()).await; + Ok(Some(Secret::String(value))) + } +} diff --git a/litellm-rust/crates/secrets-google/tests/kms.rs b/litellm-rust/crates/secrets-google/tests/kms.rs new file mode 100644 index 00000000000..667ecd268c8 --- /dev/null +++ b/litellm-rust/crates/secrets-google/tests/kms.rs @@ -0,0 +1,49 @@ +use base64::{Engine, engine::general_purpose::STANDARD}; +use google_cloud_kms_v1::client::KeyManagementService; +use litellm_secrets_google::GoogleKms; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, path}, +}; + +#[tokio::test] +async fn google_kms_decrypts_using_the_configured_resource() { + let server = MockServer::start().await; + let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key"; + Mock::given(path(format!("/v1/{resource}:decrypt"))) + .and(body_json( + serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = KeyManagementService::builder() + .with_endpoint(server.uri()) + .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build()) + .with_retry_policy(google_cloud_gax::retry_policy::NeverRetry) + .build() + .await + .unwrap(); + let manager = GoogleKms::new(client, resource.into()); + assert_eq!( + manager.decrypt(b"encrypted".to_vec()).await.unwrap(), + b" value\n" + ); +} + +#[tokio::test] +async fn disabled_google_kms_loader_does_not_require_environment_configuration() { + use std::sync::Arc; + for enabled in [None, Some(false)] { + assert!( + litellm_secrets_google::load_google_kms(enabled, Arc::new(|_: &str| None)) + .await + .unwrap() + .is_none() + ); + } +} diff --git a/litellm-rust/crates/secrets-google/tests/secret_manager.rs b/litellm-rust/crates/secrets-google/tests/secret_manager.rs new file mode 100644 index 00000000000..b3b1d29e62c --- /dev/null +++ b/litellm-rust/crates/secrets-google/tests/secret_manager.rs @@ -0,0 +1,188 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_google::{Error, GoogleSecretManager}; + +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, path}, +}; + +fn manager(server: &MockServer, always_read: bool, ttl: Duration) -> GoogleSecretManager { + GoogleSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "project".into(), + Arc::new(|name: &str| (name == "VERTEX_AI_API_KEY").then(|| "token".into())), + Some(ttl), + always_read, + ) + .unwrap() +} + +#[rstest::rstest] +#[case::nonempty("private-value")] +#[case::empty("")] +#[tokio::test] +async fn successful_reads_use_auth_latest_version_and_cache_including_empty_values( + #[case] value: &str, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .and(header("authorization", "Bearer token")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(value)}})), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str() + .unwrap(), + value + ); + } +} + +#[rstest::rstest] +#[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 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( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("recovered")}})), + ) + .expect(1) + .mount(&server) + .await; + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some("recovered") + ); + } +} + +#[rstest::rstest] +#[case::always_read(true, Duration::from_secs(60))] +#[case::expired_cache(false, Duration::from_millis(1))] +#[tokio::test] +async fn always_read_and_expired_cache_fetch_again( + #[case] always_read: bool, + #[case] ttl: Duration, +) { + let server = MockServer::start().await; + Mock::given(path( + "/v1/projects/project/secrets/key/versions/latest:access", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode("value")}})), + ) + .expect(2) + .mount(&server) + .await; + let manager = manager(&server, always_read, ttl); + for _ in 0..2 { + tokio::time::sleep(Duration::from_millis(5)).await; + assert!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .is_some() + ); + } +} + +#[test] +fn google_manager_requires_host_license_and_project_configuration() { + assert!(matches!( + GoogleSecretManager::new(Arc::new(|_: &str| None), false), + Err(Error::EnterpriseRequired) + )); + assert!(matches!( + GoogleSecretManager::new(Arc::new(|_: &str| None), true), + Err(Error::MissingEnvironment( + "GOOGLE_SECRET_MANAGER_PROJECT_ID" + )) + )); +} + +#[rstest::rstest] +#[case("true")] +#[case("null")] +#[case("\"text\"")] +#[case("{\"key\":1}")] +#[tokio::test] +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", + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"payload":{"data":STANDARD.encode(raw)}})), + ) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, false, Duration::from_secs(60)); + for _ in 0..2 { + assert_eq!( + manager + .get_secret_from_google_secret_manager("key") + .await + .unwrap() + .unwrap() + .as_str(), + Some(raw) + ); + } +} diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml new file mode 100644 index 00000000000..acd29746722 --- /dev/null +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "litellm-secrets-types" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-types.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs new file mode 100644 index 00000000000..d71bce64221 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs @@ -0,0 +1,58 @@ +use crate::{Error, SecretValue}; + +pub fn validate_secret_name(name: &str) -> Result<(), Error> { + if name.split('/').any(|segment| segment == "..") + || name + .chars() + .any(|c| c.is_control() || matches!(c, '\u{2028}' | '\u{2029}')) + { + return Err(Error::UnsafeSecretName); + } + Ok(()) +} + +#[expect( + async_fn_in_trait, + reason = "closed backend dispatch does not require Send bounds on generic rotation" +)] +pub trait BaseSecretManager { + type Error: From; + type WriteResponse; + type DeleteResponse; + + async fn async_read_secret(&self, name: &str) -> Result, Self::Error>; + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result; + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result; +} + +pub async fn async_rotate_secret( + manager: &M, + current_name: &str, + new_name: &str, + value: &SecretValue, +) -> Result { + if manager.async_read_secret(current_name).await?.is_none() { + return Err(Error::CurrentSecretMissing.into()); + } + let response = manager + .async_write_secret( + new_name, + value, + Some(&format!("Rotated from {current_name}")), + ) + .await?; + if manager.async_read_secret(new_name).await?.is_none() { + return Err(Error::NewSecretMissing.into()); + } + manager.async_delete_secret(current_name, 7).await?; + Ok(response) +} diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs new file mode 100644 index 00000000000..36d319311a3 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -0,0 +1,92 @@ +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; + +use crate::SecretValue; + +#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum KeyManagementSystem { + GoogleKms, + AzureKeyVault, + AwsSecretManager, + GoogleSecretManager, + HashicorpVault, + Cyberark, + Local, + AwsKms, + Custom, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AccessMode { + #[default] + ReadOnly, + WriteOnly, + ReadAndWrite, +} + +impl AccessMode { + pub fn readable(self) -> bool { + matches!(self, Self::ReadOnly | Self::ReadAndWrite) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(default)] +pub struct KeyManagementSettings { + pub hosted_keys: Option>, + pub store_virtual_keys: Option, + pub prefix_for_stored_virtual_keys: String, + pub access_mode: AccessMode, + pub primary_secret_name: Option, + pub description: Option, + pub tags: Option>, + pub kms_key_id: Option, + pub custom_secret_manager: Option, + pub aws_region_name: Option, + pub aws_role_name: Option, + pub aws_session_name: Option, + #[serde(serialize_with = "serialize_secret")] + pub aws_external_id: Option, + pub aws_profile_name: Option, + #[serde(serialize_with = "serialize_secret")] + pub aws_web_identity_token: Option, + pub aws_sts_endpoint: Option, + pub replica_regions: Option>, +} + +impl Default for KeyManagementSettings { + fn default() -> Self { + Self { + hosted_keys: None, + store_virtual_keys: Some(false), + prefix_for_stored_virtual_keys: "litellm/".into(), + access_mode: AccessMode::ReadOnly, + primary_secret_name: None, + description: None, + tags: None, + kms_key_id: None, + custom_secret_manager: None, + aws_region_name: None, + aws_role_name: None, + aws_session_name: None, + aws_external_id: None, + aws_profile_name: None, + aws_web_identity_token: None, + aws_sts_endpoint: None, + replica_regions: None, + } + } +} + +fn serialize_secret( + value: &Option, + serializer: S, +) -> Result { + value + .as_ref() + .map(SecretValue::expose) + .serialize(serializer) +} diff --git a/litellm-rust/crates/secrets-types/src/error.rs b/litellm-rust/crates/secrets-types/src/error.rs new file mode 100644 index 00000000000..cae9c7f4c69 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/error.rs @@ -0,0 +1,9 @@ +#[derive(Debug, thiserror::Error, PartialEq, Eq)] +pub enum Error { + #[error("secret name contains an unsafe path segment or control character")] + UnsafeSecretName, + #[error("current secret was not found")] + CurrentSecretMissing, + #[error("new secret could not be verified")] + NewSecretMissing, +} diff --git a/litellm-rust/crates/secrets-types/src/lib.rs b/litellm-rust/crates/secrets-types/src/lib.rs new file mode 100644 index 00000000000..0823ed13c06 --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/lib.rs @@ -0,0 +1,12 @@ +#![forbid(unsafe_code)] + +mod base_secret_manager; +mod config; +mod error; +mod value; + +pub use base_secret_manager::{BaseSecretManager, async_rotate_secret, validate_secret_name}; +pub use config::{AccessMode, KeyManagementSettings, KeyManagementSystem}; +pub use error::Error; +pub use litellm_auth_types::SecretValue; +pub use value::Secret; diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs new file mode 100644 index 00000000000..087537fb3eb --- /dev/null +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -0,0 +1,31 @@ +use crate::SecretValue; + +#[derive(Clone, PartialEq, Eq, veil::Redact)] +pub enum Secret { + String(SecretValue), + Bool(#[redact] bool), + Json(#[redact] serde_json::Value), +} + +impl From for Secret { + fn from(value: SecretValue) -> Self { + Self::String(value) + } +} + +impl Secret { + pub fn from_json(value: serde_json::Value) -> Self { + match value { + serde_json::Value::String(value) => Self::String(SecretValue::new(value)), + serde_json::Value::Bool(value) => Self::Bool(value), + value => Self::Json(value), + } + } + + pub fn as_str(&self) -> Option<&str> { + match self { + Self::String(value) => Some(value.expose()), + Self::Bool(_) | Self::Json(_) => None, + } + } +} diff --git a/litellm-rust/crates/secrets-types/tests/config.rs b/litellm-rust/crates/secrets-types/tests/config.rs new file mode 100644 index 00000000000..4a5f17bc68a --- /dev/null +++ b/litellm-rust/crates/secrets-types/tests/config.rs @@ -0,0 +1,60 @@ +use litellm_secrets_types::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +use serde_json::json; + +#[test] +fn config_preserves_defaults_nulls_and_serialized_names() { + let empty: KeyManagementSettings = serde_json::from_value(json!({})).unwrap(); + assert_eq!(empty, KeyManagementSettings::default()); + assert_eq!(empty.access_mode, AccessMode::ReadOnly); + assert_eq!(empty.store_virtual_keys, Some(false)); + assert_eq!(empty.prefix_for_stored_virtual_keys, "litellm/"); + let configured: KeyManagementSettings = serde_json::from_value(json!({ + "hosted_keys": [], "store_virtual_keys": null, "access_mode": "write_only", + "aws_web_identity_token": "private-token", "aws_external_id": "private-id", + "tags": {"stage": "test"}, "replica_regions": ["test-region"] + })) + .unwrap(); + assert!(!configured.access_mode.readable()); + assert_eq!(configured.store_virtual_keys, None); + assert_eq!(configured.hosted_keys.as_deref(), Some([].as_slice())); + assert!(!format!("{configured:?}").contains("private-")); + let serialized = serde_json::to_value(&configured).unwrap(); + assert_eq!(serialized["access_mode"], "write_only"); + assert_eq!(serialized["aws_web_identity_token"], "private-token"); + assert_eq!( + serde_json::from_value::(serialized).unwrap(), + configured + ); +} + +#[rstest::rstest] +#[case::aws_kms("aws_kms", KeyManagementSystem::AwsKms)] +#[case::aws_secret_manager("aws_secret_manager", KeyManagementSystem::AwsSecretManager)] +#[case::google_kms("google_kms", KeyManagementSystem::GoogleKms)] +#[case::google_secret_manager("google_secret_manager", KeyManagementSystem::GoogleSecretManager)] +#[case::azure_key_vault("azure_key_vault", KeyManagementSystem::AzureKeyVault)] +#[case::hashicorp_vault("hashicorp_vault", KeyManagementSystem::HashicorpVault)] +#[case::cyberark("cyberark", KeyManagementSystem::Cyberark)] +#[case::custom("custom", KeyManagementSystem::Custom)] +#[case::local("local", KeyManagementSystem::Local)] +fn key_management_system_serialization_round_trips( + #[case] name: &str, + #[case] system: KeyManagementSystem, +) { + assert_eq!( + serde_json::from_value::(json!(name)).unwrap(), + system + ); + assert_eq!(serde_json::to_value(system).unwrap(), name); +} + +#[test] +fn secret_debug_never_exposes_values() { + assert!( + !format!("{:?}", Secret::String(SecretValue::new("sensitive-value"))) + .contains("sensitive-value") + ); + assert!(!format!("{:?}", Secret::Bool(true)).contains("true")); +} diff --git a/litellm-rust/crates/secrets-types/tests/rotation.rs b/litellm-rust/crates/secrets-types/tests/rotation.rs new file mode 100644 index 00000000000..48a5304ece8 --- /dev/null +++ b/litellm-rust/crates/secrets-types/tests/rotation.rs @@ -0,0 +1,105 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; + +use litellm_secrets_types::{ + BaseSecretManager, Error, SecretValue, async_rotate_secret, validate_secret_name, +}; + +struct Manager { + step: AtomicUsize, + absent_at: Option, +} + +impl BaseSecretManager for Manager { + type Error = Error; + type WriteResponse = &'static str; + type DeleteResponse = (); + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + let step = self.step.fetch_add(1, Ordering::SeqCst); + assert_eq!(name, if step == 0 { "old" } else { "new" }); + Ok((self.absent_at != Some(step)).then(|| SecretValue::new("value"))) + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result { + assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 1); + assert_eq!(name, "new"); + assert_eq!(value.expose(), "replacement"); + assert_eq!(description, Some("Rotated from old")); + Ok("provider-response") + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result<(), Error> { + assert_eq!(self.step.fetch_add(1, Ordering::SeqCst), 3); + assert_eq!(name, "old"); + assert_eq!(recovery_window_in_days, 7); + Ok(()) + } +} + +#[tokio::test] +async fn rotation_verifies_before_deleting_and_returns_provider_response() { + let manager = Manager { + step: AtomicUsize::new(0), + absent_at: None, + }; + assert_eq!( + async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement")) + .await + .unwrap(), + "provider-response" + ); + assert_eq!(manager.step.load(Ordering::SeqCst), 4); +} + +#[rstest::rstest] +#[case::current_secret_missing(0, Error::CurrentSecretMissing, 1)] +#[case::new_secret_missing(2, Error::NewSecretMissing, 3)] +#[tokio::test] +async fn missing_old_or_new_value_stops_rotation_before_deletion( + #[case] absent_at: usize, + #[case] expected: Error, + #[case] calls: usize, +) { + let manager = Manager { + step: AtomicUsize::new(0), + absent_at: Some(absent_at), + }; + assert_eq!( + async_rotate_secret(&manager, "old", "new", &SecretValue::new("replacement")) + .await + .unwrap_err(), + expected + ); + assert_eq!(manager.step.load(Ordering::SeqCst), calls); +} + +#[rstest::rstest] +#[case::parent("..")] +#[case::parent_prefix("../x")] +#[case::parent_segment("x/../y")] +#[case::parent_suffix("x/..")] +#[case::line_feed("line\n")] +#[case::next_line("\u{85}")] +#[case::line_separator("\u{2028}")] +#[case::paragraph_separator("\u{2029}")] +fn names_reject_path_traversal_and_control_characters(#[case] name: &str) { + assert_eq!(validate_secret_name(name), Err(Error::UnsafeSecretName)); +} + +#[rstest::rstest] +#[case::embedded_double_dot("release-1.0..2")] +#[case::path_separator("folder/key")] +#[case::empty("")] +#[case::three_dots("...")] +fn names_allow_safe_values(#[case] name: &str) { + assert_eq!(validate_secret_name(name), Ok(())); +} diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml new file mode 100644 index 00000000000..a7e7ec80636 --- /dev/null +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "litellm-secrets" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[features] +default = [] +aws = ["dep:litellm-secrets-aws"] +google = ["dep:litellm-secrets-google"] + +[dependencies] +litellm-secrets-types.workspace = true +litellm-secrets-aws = { workspace = true, optional = true } +litellm-secrets-google = { workspace = true, optional = true } +litellm-core-utils.workspace = true +base64.workspace = true +serde.workspace = true +strum.workspace = true +jsonwebtoken.workspace = true +serde_json.workspace = true +thiserror.workspace = true +reqwest.workspace = true +moka.workspace = true +tokio = { workspace = true, features = ["fs"] } + +[dev-dependencies] +rstest.workspace = true +wiremock = "0.6.5" +tempfile = "3" +aws-sdk-kms = "1.120.0" +google-cloud-kms-v1 = "1.14.0" +google-cloud-auth.workspace = true diff --git a/litellm-rust/crates/secrets/README.md b/litellm-rust/crates/secrets/README.md new file mode 100644 index 00000000000..183a39e15bb --- /dev/null +++ b/litellm-rust/crates/secrets/README.md @@ -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 diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs new file mode 100644 index 00000000000..0c6e681b8aa --- /dev/null +++ b/litellm-rust/crates/secrets/src/error.rs @@ -0,0 +1,33 @@ +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("encrypted environment value is missing")] + MissingCiphertext, + #[error("ciphertext is not valid base64 for the configured manager")] + InvalidCiphertext, + #[error("decrypted value is not UTF-8")] + Utf8, + #[error("unsupported OIDC provider or missing build feature")] + UnsupportedOidc, + #[error("OIDC reference requires a provider and audience")] + InvalidOidc, + #[error("OIDC environment variable is missing")] + MissingEnvironment, + #[error("OIDC request failed")] + OidcHttp, + #[error("OIDC provider returned HTTP {0}")] + OidcStatus(u16), + #[error("OIDC response is invalid")] + OidcResponse, + #[error("OIDC file path must be absolute and within the credential allowlist")] + UnsafeOidcPath, + #[error("OIDC file could not be read")] + OidcFile, + #[error("secret cannot be converted to {expected}")] + TypeMismatch { expected: &'static str }, + #[cfg(feature = "aws")] + #[error(transparent)] + Aws(#[from] litellm_secrets_aws::Error), + #[cfg(feature = "google")] + #[error(transparent)] + Google(#[from] litellm_secrets_google::Error), +} diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs new file mode 100644 index 00000000000..943ffdf6158 --- /dev/null +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -0,0 +1,117 @@ +use litellm_core_utils::settings::Lookup; + +use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue}; + +#[derive(Clone)] +pub enum SecretManager { + Local, + #[cfg(feature = "aws")] + AwsKms(crate::aws::AwsKms), + #[cfg(feature = "aws")] + AwsSecretsManagerV2(crate::aws::AwsSecretsManagerV2), + #[cfg(feature = "google")] + GoogleKms(crate::google::GoogleKms), + #[cfg(feature = "google")] + GoogleSecretManager(crate::google::GoogleSecretManager), +} + +impl SecretManager { + pub fn system(&self) -> KeyManagementSystem { + match self { + Self::Local => KeyManagementSystem::Local, + #[cfg(feature = "aws")] + Self::AwsKms(_) => KeyManagementSystem::AwsKms, + #[cfg(feature = "aws")] + Self::AwsSecretsManagerV2(_) => KeyManagementSystem::AwsSecretManager, + #[cfg(feature = "google")] + Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, + #[cfg(feature = "google")] + Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + } + } +} + +pub async fn get_secret_from_manager( + client: &SecretManager, + secret_name: &str, + _settings: &KeyManagementSettings, + environment: &(dyn Lookup + Send + Sync), +) -> Result, Error> { + match client { + SecretManager::Local => Ok(environment + .get(secret_name) + .map(SecretValue::new) + .map(Secret::String)), + #[cfg(feature = "aws")] + SecretManager::AwsKms(client) => { + let ciphertext = environment + .get(secret_name) + .ok_or(Error::MissingCiphertext)?; + let plaintext = client + .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Permissive)?) + .await?; + let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + Ok(Some(Secret::String(SecretValue::new(value.trim())))) + } + #[cfg(feature = "google")] + SecretManager::GoogleKms(client) => { + let ciphertext = environment + .get(secret_name) + .ok_or(Error::MissingCiphertext)?; + let plaintext = client + .decrypt(decode_ciphertext(&ciphertext, Base64Mode::Canonical)?) + .await?; + let value = String::from_utf8(plaintext).map_err(|_| Error::Utf8)?; + Ok(Some(Secret::String(SecretValue::new(value)))) + } + #[cfg(feature = "aws")] + SecretManager::AwsSecretsManagerV2(client) => client + .read_secret_for_resolver( + secret_name, + _settings.primary_secret_name.as_deref(), + environment, + ) + .await + .map_err(Error::from), + #[cfg(feature = "google")] + SecretManager::GoogleSecretManager(client) => client + .get_secret_from_google_secret_manager(secret_name) + .await + .map_err(Error::from), + } +} + +#[cfg(any(feature = "aws", feature = "google"))] +#[derive(Clone, Copy)] +enum Base64Mode { + #[cfg(feature = "google")] + Canonical, + #[cfg(feature = "aws")] + Permissive, +} + +#[cfg(any(feature = "aws", feature = "google"))] +fn decode_ciphertext(value: &str, mode: Base64Mode) -> Result, Error> { + use base64::{Engine, engine::general_purpose::STANDARD}; + let canonical = match mode { + #[cfg(feature = "google")] + Base64Mode::Canonical => true, + #[cfg(feature = "aws")] + Base64Mode::Permissive => false, + }; + let encoded = if canonical { + value.to_owned() + } else { + value + .chars() + .filter(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '/' | '=')) + .collect() + }; + let ciphertext = STANDARD + .decode(&encoded) + .map_err(|_| Error::InvalidCiphertext)?; + if canonical && STANDARD.encode(&ciphertext) != encoded { + return Err(Error::InvalidCiphertext); + } + Ok(ciphertext) +} diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs new file mode 100644 index 00000000000..ff2e95f7b2f --- /dev/null +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -0,0 +1,21 @@ +#![forbid(unsafe_code)] + +mod error; +mod handler; +mod oidc; +mod resolver; +mod state; + +pub use error::Error; +pub use handler::{SecretManager, get_secret_from_manager}; +pub use litellm_secrets_types::{ + AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, +}; +pub use oidc::{OidcProvider, OidcReference, OidcResolver}; +pub use resolver::{FailurePolicy, SecretResolver}; +pub use state::{SecretManagerState, secret_manager_would_be_consulted}; + +#[cfg(feature = "aws")] +pub use litellm_secrets_aws as aws; +#[cfg(feature = "google")] +pub use litellm_secrets_google as google; diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs new file mode 100644 index 00000000000..fd477859bf6 --- /dev/null +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -0,0 +1,269 @@ +use std::{ + path::Path, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use jsonwebtoken::dangerous::insecure_decode_claims; +use litellm_core_utils::settings::Lookup; +use moka::future::Cache; +use serde::Deserialize; + +use crate::{Error, SecretValue}; + +const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540); +const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295); +const TOKEN_EXPIRY_MARGIN_SECONDS: f64 = 60.0; +const CIRCLE_OIDC_TOKEN: &str = "CIRCLE_OIDC_TOKEN"; +const CIRCLE_OIDC_TOKEN_V2: &str = "CIRCLE_OIDC_TOKEN_V2"; +const AZURE_FEDERATED_TOKEN_FILE: &str = "AZURE_FEDERATED_TOKEN_FILE"; +const ACTIONS_ID_TOKEN_REQUEST_URL: &str = "ACTIONS_ID_TOKEN_REQUEST_URL"; +const ACTIONS_ID_TOKEN_REQUEST_TOKEN: &str = "ACTIONS_ID_TOKEN_REQUEST_TOKEN"; +const OIDC_ALLOWED_CREDENTIAL_DIRS: &str = "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS"; +const DEFAULT_CREDENTIAL_DIRS: &str = "/var/run/secrets,/run/secrets"; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq, strum::EnumString, strum::AsRefStr)] +#[strum(serialize_all = "snake_case")] +pub enum OidcProvider { + Google, + #[strum(serialize = "circleci")] + CircleCi, + #[strum(serialize = "circleci_v2")] + CircleCiV2, + Github, + Azure, + File, + Env, + EnvPath, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct OidcReference<'a> { + pub provider: OidcProvider, + pub audience: &'a str, +} + +impl<'a> TryFrom<&'a str> for OidcReference<'a> { + type Error = Error; + + fn try_from(reference: &'a str) -> Result { + let (provider, audience) = reference + .strip_prefix("oidc/") + .and_then(|body| body.split_once('/')) + .ok_or(Error::InvalidOidc)?; + Ok(Self { + provider: provider.parse().map_err(|_| Error::UnsupportedOidc)?, + audience, + }) + } +} + +#[derive(Deserialize)] +struct OidcTokenClaims { + exp: Option, +} + +#[derive(Deserialize)] +#[serde(untagged)] +enum NumericDate { + Number(f64), + String(String), +} + +impl NumericDate { + fn seconds(self) -> Option { + match self { + Self::Number(value) => Some(value), + Self::String(value) => value.parse().ok(), + } + .filter(|value| value.is_finite()) + } +} + +pub struct OidcResolver { + client: reqwest::Client, + google_identity_endpoint: reqwest::Url, + cache: Cache, + clock: fn() -> SystemTime, +} + +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( + client, + reqwest::Url::parse("http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/identity").expect("static URL"), + ) + } +} + +impl OidcResolver { + pub fn new(client: reqwest::Client, google_identity_endpoint: reqwest::Url) -> Self { + Self { + client, + google_identity_endpoint, + cache: Cache::builder() + .max_capacity(200) + .time_to_live(GOOGLE_TOKEN_MAX_TTL) + .build(), + clock: SystemTime::now, + } + } + + pub fn with_clock(self, clock: fn() -> SystemTime) -> Self { + Self { clock, ..self } + } + + pub async fn resolve( + &self, + reference: &str, + environment: &(dyn Lookup + Send + Sync), + ) -> Result, Error> { + let OidcReference { provider, audience } = reference.try_into()?; + match provider { + OidcProvider::CircleCi => required_env(environment, CIRCLE_OIDC_TOKEN) + .map(SecretValue::new) + .map(Some), + OidcProvider::CircleCiV2 => required_env(environment, CIRCLE_OIDC_TOKEN_V2) + .map(SecretValue::new) + .map(Some), + OidcProvider::Env => required_env(environment, audience) + .map(SecretValue::new) + .map(Some), + OidcProvider::EnvPath => read_file(&required_env(environment, audience)?) + .await + .map(Some), + OidcProvider::File => read_allowed_file(audience, environment).await.map(Some), + OidcProvider::Azure => { + if let Some(path) = environment.get(AZURE_FEDERATED_TOKEN_FILE) { + return read_file(&path).await.map(Some); + } + Err(Error::UnsupportedOidc) + } + OidcProvider::Github => { + let url = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_URL)?; + let authorization = required_env(environment, ACTIONS_ID_TOKEN_REQUEST_TOKEN)?; + if let Some(value) = self.cached(reference).await { + return Ok(Some(value)); + } + let response = self + .client + .get(url) + .query(&[("audience", audience)]) + .bearer_auth(authorization) + .header("Accept", "application/json; api-version=2.0") + .send() + .await + .map_err(|_| Error::OidcHttp)?; + if response.status() != reqwest::StatusCode::OK { + return Err(Error::OidcStatus(response.status().as_u16())); + } + #[derive(Deserialize)] + struct Token { + value: Option, + } + let token: Token = response.json().await.map_err(|_| Error::OidcResponse)?; + if let Some(value) = &token.value { + self.cache + .insert( + reference.to_owned(), + (value.clone(), (self.clock)() + GITHUB_TOKEN_TTL), + ) + .await; + } + Ok(token.value) + } + OidcProvider::Google => { + if !cfg!(feature = "google") { + return Err(Error::UnsupportedOidc); + } + if let Some(value) = self.cached(reference).await { + return Ok(Some(value)); + } + let response = self + .client + .get(self.google_identity_endpoint.clone()) + .query(&[("audience", audience)]) + .header("Metadata-Flavor", "Google") + .send() + .await + .map_err(|_| Error::OidcHttp)?; + if response.status() != reqwest::StatusCode::OK { + return Err(Error::OidcStatus(response.status().as_u16())); + } + let token = response.text().await.map_err(|_| Error::OidcResponse)?; + let now = (self.clock)(); + let ttl = oidc_token_cache_ttl(&token, now, GOOGLE_TOKEN_MAX_TTL); + let value = SecretValue::new(token); + if let Some(ttl) = ttl.filter(|ttl| !ttl.is_zero()) { + self.cache + .insert(reference.to_owned(), (value.clone(), now + ttl)) + .await; + } + Ok(Some(value)) + } + } + } + + async fn cached(&self, reference: &str) -> Option { + self.cache + .get(reference) + .await + .and_then(|(value, expires)| ((self.clock)() < expires).then_some(value)) + } +} + +fn required_env(environment: &dyn Lookup, name: &str) -> Result { + environment.get(name).ok_or(Error::MissingEnvironment) +} + +async fn read_file(path: &str) -> Result { + tokio::fs::read_to_string(path) + .await + .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) + .map_err(|_| Error::OidcFile) +} + +async fn read_allowed_file( + path: &str, + environment: &(dyn Lookup + Sync), +) -> Result { + if !Path::new(path).is_absolute() { + return Err(Error::UnsafeOidcPath); + } + let resolved = tokio::fs::canonicalize(path) + .await + .map_err(|_| Error::OidcFile)?; + let allowed = environment + .get(OIDC_ALLOWED_CREDENTIAL_DIRS) + .filter(|v| !v.is_empty()) + .unwrap_or_else(|| DEFAULT_CREDENTIAL_DIRS.into()); + for directory in allowed.split(',').map(str::trim).filter(|d| !d.is_empty()) { + if let Ok(directory) = tokio::fs::canonicalize(directory).await + && resolved.starts_with(directory) + { + return tokio::fs::read_to_string(&resolved) + .await + .map(|value| SecretValue::new(value.replace("\r\n", "\n").replace('\r', "\n"))) + .map_err(|_| Error::OidcFile); + } + } + Err(Error::UnsafeOidcPath) +} + +fn oidc_token_cache_ttl(token: &str, now: SystemTime, max_ttl: Duration) -> Option { + let fallback = Some(max_ttl); + let Ok(claims) = insecure_decode_claims::(token) else { + return fallback; + }; + let Some(exp) = claims.exp.and_then(NumericDate::seconds) else { + return fallback; + }; + let seconds = exp.trunc() + - now.duration_since(UNIX_EPOCH).ok()?.as_secs() as f64 + - TOKEN_EXPIRY_MARGIN_SECONDS; + (seconds > 0.0).then(|| Duration::from_secs_f64(seconds.min(max_ttl.as_secs_f64()))) +} diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs new file mode 100644 index 00000000000..89439893852 --- /dev/null +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -0,0 +1,135 @@ +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}; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum FailurePolicy { + #[default] + Propagate, + EnvironmentFallback, +} + +pub struct SecretResolver { + state: Arc, + environment: Arc, + oidc: OidcResolver, + failure_policy: FailurePolicy, +} + +impl Default for SecretResolver { + fn default() -> Self { + Self::new( + Arc::new(SecretManagerState::default()), + Arc::new(ProcessEnvironment), + OidcResolver::default(), + ) + } +} + +impl SecretResolver { + pub fn new( + state: Arc, + environment: Arc, + oidc: OidcResolver, + ) -> Self { + Self { + 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, + ) -> Result, Error> { + let name = normalize_secret_name(name); + if name.starts_with("oidc/") { + return self + .oidc + .resolve(name, self.environment.as_ref()) + .await + .map(|value| value.map(Secret::String).or(default_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), + }, + } + } + + fn environment_secret(&self, name: &str) -> Option { + self.environment + .get(name) + .map(SecretValue::new) + .map(Secret::String) + } + + pub async fn get_secret_str( + &self, + name: &str, + default_value: Option, + ) -> Result, Error> { + 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( + &self, + name: &str, + default_value: Option, + ) -> Result, Error> { + 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), + } + } +} diff --git a/litellm-rust/crates/secrets/src/state.rs b/litellm-rust/crates/secrets/src/state.rs new file mode 100644 index 00000000000..7854d763ef2 --- /dev/null +++ b/litellm-rust/crates/secrets/src/state.rs @@ -0,0 +1,59 @@ +use crate::{KeyManagementSettings, KeyManagementSystem, SecretManager}; + +pub(crate) enum LookupTarget<'a> { + Environment, + Manager { + backend: &'a SecretManager, + settings: &'a KeyManagementSettings, + }, +} + +pub(crate) fn normalize_secret_name(name: &str) -> &str { + name.strip_prefix("os.environ/").unwrap_or(name) +} + +#[derive(Clone, Default)] +pub struct SecretManagerState { + manager: Option<(SecretManager, KeyManagementSettings)>, +} + +impl SecretManagerState { + pub fn new(backend: SecretManager, settings: KeyManagementSettings) -> Self { + Self { + manager: Some((backend, settings)), + } + } + + pub fn system(&self) -> Option { + self.backend().map(SecretManager::system) + } + + 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.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 } + } + _ => LookupTarget::Environment, + } + } +} + +pub fn secret_manager_would_be_consulted(state: &SecretManagerState, name: &str) -> bool { + let name = normalize_secret_name(name); + !name.starts_with("oidc/") && matches!(state.lookup_target(name), LookupTarget::Manager { .. }) +} diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs new file mode 100644 index 00000000000..a2cbbd843e1 --- /dev/null +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -0,0 +1,107 @@ +#[cfg(feature = "aws")] +#[tokio::test] +async fn aws_handler_reads_ciphertext_decodes_trims_and_redacts() { + use aws_sdk_kms::{ + Client, + config::{BehaviorVersion, Credentials, Region}, + }; + use base64::{Engine, engine::general_purpose::STANDARD}; + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, aws::AwsKms, get_secret_from_manager, + }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::body_json}; + + let server = MockServer::start().await; + Mock::given(body_json( + serde_json::json!({"CiphertextBlob": STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"Plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = Client::from_conf( + aws_sdk_kms::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new("us-east-1")) + .credentials_provider(Credentials::new("test", "test", None, None, "test")) + .endpoint_url(server.uri()) + .build(), + ); + let manager = SecretManager::AwsKms(AwsKms::new(client)); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|name: &str| { + assert_eq!(name, "KEY"); + Some(format!(" {}\n", STANDARD.encode("encrypted"))) + }) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some("value")); + assert!(!format!("{value:?}").contains("value")); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, + Err(Error::MissingCiphertext) + )); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some("abc".into())).await, + Err(Error::InvalidCiphertext) + )); +} + +#[cfg(feature = "google")] +#[tokio::test] +async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whitespace() { + use base64::{Engine, engine::general_purpose::STANDARD}; + use google_cloud_kms_v1::client::KeyManagementService; + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, get_secret_from_manager, google::GoogleKms, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_json, path}, + }; + + let server = MockServer::start().await; + let resource = "projects/project/locations/global/keyRings/ring/cryptoKeys/key"; + Mock::given(path(format!("/v1/{resource}:decrypt"))) + .and(body_json( + serde_json::json!({"ciphertext":STANDARD.encode("encrypted")}), + )) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"plaintext":STANDARD.encode(" value\n")})), + ) + .expect(1) + .mount(&server) + .await; + let client = KeyManagementService::builder() + .with_endpoint(server.uri()) + .with_credentials(google_cloud_auth::credentials::anonymous::Builder::new().build()) + .build() + .await + .unwrap(); + let manager = SecretManager::GoogleKms(GoogleKms::new(client, resource.into())); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| { + Some(STANDARD.encode("encrypted")) + }) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some(" value\n")); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| Some(format!( + " {}", + STANDARD.encode("encrypted") + ))) + .await, + Err(Error::InvalidCiphertext) + )); + assert!(matches!( + get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None).await, + Err(Error::MissingCiphertext) + )); +} diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs new file mode 100644 index 00000000000..b17e7de7f9d --- /dev/null +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -0,0 +1,295 @@ +use std::{collections::BTreeMap, sync::Arc}; + +use litellm_core_utils::settings::Lookup; +use litellm_secrets::{Error, OidcResolver, Secret, SecretManagerState, SecretResolver}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, method, path, query_param}, +}; + +fn environment(pairs: &[(&str, &str)]) -> Arc { + let values: BTreeMap = pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(); + Arc::new(move |name: &str| values.get(name).cloned()) +} + +#[rstest::rstest] +#[case::environment("oidc/env/TOKEN", "true")] +#[case::circleci("oidc/circleci/audience", "circle")] +#[case::circleci_v2("oidc/circleci_v2/audience", "circle-v2")] +#[tokio::test] +async fn environment_sources_resolve_expected_value( + #[case] reference: &str, + #[case] expected: &str, +) { + let env = environment(&[ + ("TOKEN", "true"), + ("CIRCLE_OIDC_TOKEN", "circle"), + ("CIRCLE_OIDC_TOKEN_V2", "circle-v2"), + ]); + assert_eq!( + OidcResolver::default() + .resolve(reference, env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + expected + ); +} + +#[tokio::test] +async fn environment_sources_bypass_boolean_conversion_and_defaults() { + let env = environment(&[("TOKEN", "true")]); + let oidc = OidcResolver::default(); + let resolver = SecretResolver::new(Arc::new(SecretManagerState::default()), env, oidc); + assert_eq!( + resolver + .get_secret_str("os.environ/oidc/env/TOKEN", None) + .await + .unwrap() + .unwrap() + .expose(), + "true" + ); + assert_eq!( + resolver + .get_secret_bool("oidc/env/TOKEN", None) + .await + .unwrap(), + Some(true) + ); + assert!(matches!( + resolver + .get_secret("oidc/env/MISSING", Some(Secret::Bool(true))) + .await, + Err(Error::MissingEnvironment) + )); + assert!(matches!( + resolver.get_secret("oidc/invalid", None).await, + Err(Error::InvalidOidc) + )); +} + +#[tokio::test] +async fn github_requests_are_authenticated_cached_and_revalidate_environment() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/token")) + .and(query_param("audience", "https://service/oidc/path")) + .and(header("authorization", "Bearer request-token")) + .and(header("accept", "application/json; api-version=2.0")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value":"identity-token"})), + ) + .expect(1) + .mount(&server) + .await; + let env = environment(&[ + ( + "ACTIONS_ID_TOKEN_REQUEST_URL", + &format!("{}/token", server.uri()), + ), + ("ACTIONS_ID_TOKEN_REQUEST_TOKEN", "request-token"), + ]); + let oidc = OidcResolver::default(); + for _ in 0..2 { + assert_eq!( + oidc.resolve("oidc/github/https://service/oidc/path", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "identity-token" + ); + } + assert!(matches!( + oidc.resolve( + "oidc/github/https://service/oidc/path", + environment(&[]).as_ref() + ) + .await, + Err(Error::MissingEnvironment) + )); +} + +#[tokio::test] +async fn file_allowlist_resolves_symlinks_while_environment_paths_remain_explicit() { + let allowed = tempfile::tempdir().unwrap(); + let outside = tempfile::tempdir().unwrap(); + let token = allowed.path().join("token"); + let private = outside.path().join("private"); + std::fs::write(&token, "token\r\n").unwrap(); + std::fs::write(&private, "outside").unwrap(); + let env = environment(&[ + ( + "LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS", + allowed.path().to_str().unwrap(), + ), + ("PATH_TOKEN", private.to_str().unwrap()), + ("AZURE_FEDERATED_TOKEN_FILE", token.to_str().unwrap()), + ]); + let oidc = OidcResolver::default(); + assert_eq!( + oidc.resolve(&format!("oidc/file/{}", token.display()), env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "token\n" + ); + assert!(matches!( + oidc.resolve("oidc/file/relative", env.as_ref()).await, + Err(Error::UnsafeOidcPath) + )); + assert!(matches!( + oidc.resolve(&format!("oidc/file/{}", private.display()), env.as_ref()) + .await, + Err(Error::UnsafeOidcPath) + )); + assert_eq!( + oidc.resolve("oidc/env_path/PATH_TOKEN", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "outside" + ); + assert_eq!( + oidc.resolve("oidc/azure/scope", env.as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + "token\n" + ); + #[cfg(unix)] + { + let link = allowed.path().join("link"); + std::os::unix::fs::symlink(&private, &link).unwrap(); + assert!(matches!( + oidc.resolve(&format!("oidc/file/{}", link.display()), env.as_ref()) + .await, + Err(Error::UnsafeOidcPath) + )); + } +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::at_refresh_boundary(serde_json::json!(1060), 2)] +#[case::beyond_refresh_boundary(serde_json::json!(1061), 1)] +#[case::already_expired(serde_json::json!(999), 2)] +#[case::string_expiry(serde_json::json!("999"), 2)] +#[case::fractional_expiry(serde_json::json!(1060.9), 2)] +#[case::negative_expiry(serde_json::json!(-1), 2)] +#[case::null_expiry(serde_json::Value::Null, 1)] +#[case::unreadable_expiry(serde_json::json!("invalid"), 1)] +#[case::nonfinite_expiry(serde_json::json!("NaN"), 1)] +#[tokio::test] +async fn google_expiry_caps_cache_and_preserves_audience( + #[case] expiry: serde_json::Value, + #[case] calls: u64, +) { + use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; + use std::time::{Duration, SystemTime, UNIX_EPOCH}; + fn now() -> SystemTime { + UNIX_EPOCH + Duration::from_secs(1000) + } + let server = MockServer::start().await; + let token = format!( + "{}.{}.signature", + URL_SAFE_NO_PAD.encode(serde_json::json!({"alg":"RS256","typ":"JWT"}).to_string()), + URL_SAFE_NO_PAD.encode(serde_json::json!({"exp":expiry}).to_string()) + ); + Mock::given(method("GET")) + .and(header("metadata-flavor", "Google")) + .and(query_param("audience", "https://service/oidc/path")) + .respond_with(ResponseTemplate::new(200).set_body_string(&token)) + .expect(calls) + .mount(&server) + .await; + let oidc = + OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()).with_clock(now); + for _ in 0..2 { + assert_eq!( + oidc.resolve( + "oidc/google/https://service/oidc/path", + environment(&[]).as_ref() + ) + .await + .unwrap() + .unwrap() + .expose(), + token + ); + } +} + +#[cfg(not(feature = "google"))] +#[tokio::test] +async fn google_oidc_requires_its_build_feature() { + assert!(matches!( + OidcResolver::default() + .resolve("oidc/google/audience", environment(&[]).as_ref()) + .await, + Err(Error::UnsupportedOidc) + )); +} + +#[tokio::test] +async fn azure_oidc_without_a_token_file_requires_an_unimplemented_backend() { + assert!(matches!( + OidcResolver::default() + .resolve("oidc/azure/scope", environment(&[]).as_ref()) + .await, + Err(Error::UnsupportedOidc) + )); +} + +#[rstest::rstest] +#[case::missing_prefix("env/TOKEN", false)] +#[case::missing_audience_separator("oidc/env", false)] +#[case::unknown_provider("oidc/unknown/TOKEN", true)] +#[tokio::test] +async fn invalid_references_fail_before_environment_lookup( + #[case] reference: &str, + #[case] unsupported: bool, +) { + let error = OidcResolver::default() + .resolve(reference, &|_: &str| { + panic!("invalid reference reached environment lookup") + }) + .await + .unwrap_err(); + assert!(matches!(error, Error::UnsupportedOidc) == unsupported); + assert!(matches!(error, Error::InvalidOidc) != unsupported); +} + +#[cfg(feature = "google")] +#[rstest::rstest] +#[case::opaque("opaque-token")] +#[case::missing_expiry("header.e30.signature")] +#[tokio::test] +async fn unreadable_expiry_keeps_python_cache_fallback(#[case] token: &str) { + let server = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_string(token)) + .expect(1) + .mount(&server) + .await; + let resolver = OidcResolver::new(reqwest::Client::new(), server.uri().parse().unwrap()); + for _ in 0..2 { + assert_eq!( + resolver + .resolve("oidc/google/audience", environment(&[]).as_ref()) + .await + .unwrap() + .unwrap() + .expose(), + token, + ); + } +} diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs new file mode 100644 index 00000000000..3a826092d72 --- /dev/null +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -0,0 +1,368 @@ +use std::sync::Arc; + +use litellm_secrets::{ + Error, KeyManagementSettings, OidcResolver, Secret, SecretManager, SecretManagerState, + SecretResolver, SecretValue, secret_manager_would_be_consulted, +}; + +fn resolver(value: Option<&str>, configured: bool) -> SecretResolver { + let state = if configured { + SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default()) + } else { + SecretManagerState::default() + }; + let value = value.map(str::to_owned); + SecretResolver::new( + Arc::new(state), + Arc::new(move |_: &str| value.clone()), + OidcResolver::default(), + ) +} + +#[rstest::rstest] +#[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 conversion_is_explicit_and_independent_of_manager_configuration( + #[case] input: &str, + #[case] boolean: Option, + #[values(false, true)] configured: bool, +) { + let resolver = resolver(Some(input), configured); + assert_eq!( + resolver.get_secret("key", None).await.unwrap(), + Some(Secret::String(SecretValue::new(input))) + ); + assert_eq!( + resolver + .get_secret_str("key", None) + .await + .unwrap() + .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 prefix_is_removed_once_and_local_manager_is_not_consulted() { + let state = SecretManagerState::new(SecretManager::Local, KeyManagementSettings::default()); + assert_eq!( + 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(state), + Arc::new(|name: &str| (name == "os.environ/KEY").then(|| "value".into())), + OidcResolver::default(), + ); + assert_eq!( + resolver + .get_secret_str("os.environ/os.environ/KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[tokio::test] +async fn resolver_future_can_run_on_a_tokio_worker() { + let resolver = resolver(Some("worker-value"), false); + let result = tokio::spawn(async move { resolver.get_secret_str("KEY", None).await }) + .await + .unwrap() + .unwrap(); + assert_eq!(result.unwrap().expose(), "worker-value"); +} + +#[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>, + #[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 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 = + 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(); + 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 + .with_failure_policy(FailurePolicy::EnvironmentFallback) + .get_secret_str("KEY", None) + .await + .unwrap() + .unwrap() + .expose(), + "environment" + ); +}