diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..2ed998174e4 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1415,6 +1415,7 @@ dependencies = [ "litellm-core", "pyo3", "reqwest", + "rustls 0.23.42", "serde", "serde_json", "sha2 0.10.9", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a13dd4c04b0..15074082981 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -26,6 +26,7 @@ pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } rstest = "0.26.1" +rustls = { version = "0.23", default-features = false, features = ["ring", "std"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" sha2 = "0.10" diff --git a/litellm-rust/crates/ai-gateway/Cargo.toml b/litellm-rust/crates/ai-gateway/Cargo.toml index e3dbdf24ce6..414abc2356d 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -19,6 +19,9 @@ litellm-core = { workspace = true, features = ["bedrock-auth"] } # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the # Python proxy callbacks API. reqwest.workspace = true +# rustls is a direct dependency so `io::tls` can install a process-level +# crypto provider; see that module for why the graph needs one. +rustls.workspace = true # `sync` powers the bounded mpsc channel the realtime logger drains. tokio = { workspace = true, features = ["rt-multi-thread", "macros", "net", "time", "sync"] } tokio-tungstenite.workspace = true diff --git a/litellm-rust/crates/ai-gateway/src/io/mod.rs b/litellm-rust/crates/ai-gateway/src/io/mod.rs index cce56dd2121..7098d67993f 100644 --- a/litellm-rust/crates/ai-gateway/src/io/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/io/mod.rs @@ -3,3 +3,4 @@ pub mod ocr; pub mod realtime; pub mod realtime_pool; pub mod responses_ws; +pub(crate) mod tls; diff --git a/litellm-rust/crates/ai-gateway/src/io/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 662f7328982..53d87848342 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -23,10 +23,12 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; use litellm_core::providers::openai::realtime::transformation::OPENAI_REALTIME_CONFIG; +use crate::io::tls::connect_upstream; + /// Environment variable holding the OpenAI API key (last-resort fallback). const OPENAI_API_KEY_ENV: &str = "OPENAI_API_KEY"; @@ -84,7 +86,7 @@ pub(crate) async fn dial_upstream( .map_err(|err| Error::Auth(err.to_string()))?, ); - let (upstream, _response) = connect_async(request) + let (upstream, _response) = connect_upstream(request) .await .map_err(|err| Error::Network(err.to_string()))?; Ok(upstream) diff --git a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs index 0b01747b1a5..ee181791413 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -14,7 +14,9 @@ use tokio_tungstenite::tungstenite::Message; use tokio_tungstenite::tungstenite::client::IntoClientRequest; use tokio_tungstenite::tungstenite::http::HeaderValue; use tokio_tungstenite::tungstenite::http::header::{AUTHORIZATION, HeaderName}; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream}; + +use crate::io::tls::connect_upstream; use crate::constants::{ DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS, DEFAULT_RESPONSES_WS_IDLE_TIMEOUT_SECS, @@ -49,14 +51,14 @@ impl ResponsesWebSocketConnection { .map_err(|error| Error::InvalidRequest(error.to_string()))?; request.headers_mut().insert(header_name, header_value); } - let connect = connect_async(request); + let connect = connect_upstream(request); let result = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { Error::Network("Responses WebSocket connection timed out".to_string()) })?, None => connect.await, }; - let (socket, _) = result.map_err(|error| match error { + let (socket, _) = result.map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), @@ -138,13 +140,13 @@ async fn dial_upstream( ); let result = tokio::time::timeout( Duration::from_secs(DEFAULT_RESPONSES_WS_CONNECT_TIMEOUT_SECS), - connect_async(request), + connect_upstream(request), ) .await .map_err(|_| Error::Network("Responses WebSocket connection timed out".to_string()))?; result .map(|(socket, _)| socket) - .map_err(|error| match error { + .map_err(|error| match *error { tokio_tungstenite::tungstenite::Error::Http(response) => Error::Http { status: response.status().as_u16(), body: String::new(), diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs new file mode 100644 index 00000000000..e3754e44002 --- /dev/null +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -0,0 +1,67 @@ +//! Outbound WebSocket dials, with the rustls crypto provider settled first. +//! +//! `reqwest/rustls-tls` enables `rustls/ring` and `litellm-core`'s `bedrock-auth` +//! enables `rustls/aws-lc-rs`, so `ClientConfig::builder()` — which is how +//! `tokio-tungstenite` builds its TLS config — panics rather than guess between +//! them. reqwest and the AWS SDK pick a provider explicitly and never panic. +//! +//! Installing from the dial rather than from a `main` also covers the `cdylib` +//! the Python bridge loads, the tests, and the benches, none of which have one. +//! ring is what reqwest already falls back to, so installing it changes no +//! working path, and an embedder that installed its own provider first keeps it. + +use std::sync::Once; + +use tokio::net::TcpStream; +use tokio_tungstenite::tungstenite::Error; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::handshake::client::Response; +use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; + +static INSTALL_CRYPTO_PROVIDER: Once = Once::new(); + +pub(crate) fn ensure_crypto_provider() { + INSTALL_CRYPTO_PROVIDER.call_once(|| { + let _ = rustls::crypto::ring::default_provider().install_default(); + }); +} + +pub(crate) async fn connect_upstream( + request: R, +) -> Result<(WebSocketStream>, Response), Box> +where + R: IntoClientRequest + Unpin, +{ + ensure_crypto_provider(); + connect_async(request).await.map_err(Box::new) +} + +#[cfg(test)] +mod tests { + use super::ensure_crypto_provider; + + #[test] + fn client_config_builder_works_with_both_provider_features_enabled() { + ensure_crypto_provider(); + + assert!(rustls::crypto::CryptoProvider::get_default().is_some()); + + let config = rustls::ClientConfig::builder() + .with_root_certificates(rustls::RootCertStore::empty()) + .with_no_client_auth(); + + assert!(!config.crypto_provider().cipher_suites.is_empty()); + } + + #[test] + fn ensure_crypto_provider_is_idempotent() { + ensure_crypto_provider(); + let first = rustls::crypto::CryptoProvider::get_default().cloned(); + + ensure_crypto_provider(); + let second = rustls::crypto::CryptoProvider::get_default().cloned(); + + assert!(first.is_some()); + assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap())); + } +}