From 6463994f78161ee4448380a5e9fb2ad28729b98f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 17:29:10 -0700 Subject: [PATCH] fix(ai-gateway): dial upstream WebSockets over an explicit rustls provider Build one ClientConfig that names ring and loads the native roots once, and hand it to every tokio-tungstenite dial as its connector instead of installing a process-wide default from the dial path. Each of the three dial sites gets a wss:// test that reproduces the panic if its connector is dropped. --- litellm-rust/Cargo.lock | 1 + litellm-rust/Cargo.toml | 1 + litellm-rust/crates/ai-gateway/Cargo.toml | 5 +- .../crates/ai-gateway/src/io/realtime.rs | 27 +++++ .../crates/ai-gateway/src/io/responses_ws.rs | 23 ++++ litellm-rust/crates/ai-gateway/src/io/tls.rs | 113 ++++++++++-------- .../tests/crypto_provider_wiring.rs | 19 ++- 7 files changed, 134 insertions(+), 55 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2ed998174e4..d214e80d818 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1416,6 +1416,7 @@ dependencies = [ "pyo3", "reqwest", "rustls 0.23.42", + "rustls-native-certs", "serde", "serde_json", "sha2 0.10.9", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 2e2e8809b7f..d3d25cbda8d 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -27,6 +27,7 @@ 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", "tls12"] } +rustls-native-certs = "0.8" 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 414abc2356d..82bedd0c8a0 100644 --- a/litellm-rust/crates/ai-gateway/Cargo.toml +++ b/litellm-rust/crates/ai-gateway/Cargo.toml @@ -19,9 +19,10 @@ 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 and its root store are direct dependencies so `io::tls` can build the +# one TLS config the outbound dials use; see that module for why it has to. rustls.workspace = true +rustls-native-certs.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/realtime.rs b/litellm-rust/crates/ai-gateway/src/io/realtime.rs index 53d87848342..207c31dffa0 100644 --- a/litellm-rust/crates/ai-gateway/src/io/realtime.rs +++ b/litellm-rust/crates/ai-gateway/src/io/realtime.rs @@ -286,6 +286,33 @@ mod tests { serde_json::from_str(raw).expect("valid event json") } + /// The realtime dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = dial_upstream( + "gpt-realtime", + "sk-test", + Some(&format!("wss://127.0.0.1:{port}")), + ) + .await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + #[test] fn resolve_api_key_prefers_param_then_blank_falls_through() { assert_eq!(resolve_api_key(Some("sk-test")).unwrap(), "sk-test"); 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 ee181791413..9df3d0c6cc5 100644 --- a/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs +++ b/litellm-rust/crates/ai-gateway/src/io/responses_ws.rs @@ -326,6 +326,29 @@ mod tests { use tokio::net::TcpListener; use tokio_tungstenite::accept_async; + /// The Responses dial has to reach a `wss://` upstream without a process-wide + /// crypto provider installed, which is what dialing through `io::tls` buys. + #[tokio::test] + async fn dial_upstream_over_wss_reports_an_error_instead_of_panicking() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("bind a loopback port"); + let port = listener + .local_addr() + .expect("read the bound address") + .port(); + tokio::spawn(async move { + while let Ok((stream, _peer)) = listener.accept().await { + drop(stream); + } + }); + + let result = + dial_upstream("gpt-5", "sk-test", Some(&format!("wss://127.0.0.1:{port}"))).await; + + assert!(matches!(result, Err(Error::Network(_)))); + } + async fn websocket_base() -> (String, tokio::task::JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind"); let address = listener.local_addr().expect("local address"); diff --git a/litellm-rust/crates/ai-gateway/src/io/tls.rs b/litellm-rust/crates/ai-gateway/src/io/tls.rs index 96544adffb7..16fd11e2e79 100644 --- a/litellm-rust/crates/ai-gateway/src/io/tls.rs +++ b/litellm-rust/crates/ai-gateway/src/io/tls.rs @@ -1,29 +1,54 @@ -//! Outbound WebSocket dials, with the rustls crypto provider settled first. +//! Outbound WebSocket dials over a TLS config this crate builds once and owns. //! //! `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 whoever installs into this rustls build first still wins. +//! enables `rustls/aws-lc-rs`, so the bare `ClientConfig::builder()` that +//! `tokio-tungstenite` uses when handed no connector panics rather than guess +//! between them. Naming ring on a connector of our own settles that for these +//! dials without touching the process-wide default, and building the config +//! once keeps the platform trust store, which `tokio-tungstenite` would +//! otherwise re-read on every dial, off the dial path. -use std::sync::Once; +use std::io; +use std::sync::{Arc, OnceLock}; +use rustls::{ClientConfig, RootCertStore}; use tokio::net::TcpStream; use tokio_tungstenite::tungstenite::Error; use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::error::TlsError; use tokio_tungstenite::tungstenite::handshake::client::Response; -use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; +use tokio_tungstenite::{ + Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, +}; -static INSTALL_CRYPTO_PROVIDER: Once = Once::new(); +static TLS_CONFIG: OnceLock> = OnceLock::new(); -pub(crate) fn ensure_crypto_provider() { - INSTALL_CRYPTO_PROVIDER.call_once(|| { - let _ = rustls::crypto::ring::default_provider().install_default(); - }); +fn build_config() -> Result> { + let native = rustls_native_certs::load_native_certs(); + let roots = { + let mut store = RootCertStore::empty(); + let (added, _ignored) = store.add_parsable_certificates(native.certs); + if added == 0 { + return Err(Box::new(Error::Io(io::Error::other(format!( + "no usable native root certificates: {:?}", + native.errors + ))))); + } + store + }; + + ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) + .with_safe_default_protocol_versions() + .map(|builder| builder.with_root_certificates(roots).with_no_client_auth()) + .map_err(|error| Box::new(Error::Tls(TlsError::Rustls(error)))) +} + +fn tls_config() -> Result, Box> { + if let Some(config) = TLS_CONFIG.get() { + return Ok(Arc::clone(config)); + } + let built = Arc::new(build_config()?); + Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) } pub(crate) async fn connect_upstream( @@ -32,19 +57,25 @@ pub(crate) async fn connect_upstream( where R: IntoClientRequest + Unpin, { - ensure_crypto_provider(); - connect_async(request).await.map_err(Box::new) + let request = request.into_client_request().map_err(Box::new)?; + let connector = match request.uri().scheme_str() { + Some("wss") => Some(Connector::Rustls(tls_config()?)), + _ => None, + }; + connect_async_tls_with_config(request, None, false, connector) + .await + .map_err(Box::new) } #[cfg(test)] mod tests { - use rustls::crypto::CryptoProvider; + use rustls::CipherSuite; + use rustls::NamedGroup; + use rustls::crypto::{CryptoProvider, aws_lc_rs, ring}; - use super::ensure_crypto_provider; + use super::{Arc, build_config, tls_config}; - fn fingerprint( - provider: &CryptoProvider, - ) -> (Vec, Vec) { + fn fingerprint(provider: &CryptoProvider) -> (Vec, Vec) { ( provider .cipher_suites @@ -60,43 +91,31 @@ mod tests { } #[test] - fn client_config_builder_works_with_both_provider_features_enabled() { - ensure_crypto_provider(); - - assert!(CryptoProvider::get_default().is_some()); - - let config = rustls::ClientConfig::builder() - .with_root_certificates(rustls::RootCertStore::empty()) - .with_no_client_auth(); + fn builds_a_usable_config_with_both_provider_features_enabled() { + let config = build_config().expect("a client config"); assert!(!config.crypto_provider().cipher_suites.is_empty()); } #[test] - fn installs_ring_rather_than_aws_lc_rs() { - ensure_crypto_provider(); - - let installed = CryptoProvider::get_default().expect("a provider is installed"); + fn dials_with_ring_rather_than_aws_lc_rs() { + let config = build_config().expect("a client config"); assert_eq!( - fingerprint(installed), - fingerprint(&rustls::crypto::ring::default_provider()) + fingerprint(config.crypto_provider()), + fingerprint(&ring::default_provider()) ); assert_ne!( - fingerprint(installed), - fingerprint(&rustls::crypto::aws_lc_rs::default_provider()) + fingerprint(config.crypto_provider()), + fingerprint(&aws_lc_rs::default_provider()) ); } #[test] - fn ensure_crypto_provider_is_idempotent() { - ensure_crypto_provider(); - let first = CryptoProvider::get_default().cloned(); + fn the_trust_store_is_loaded_once_and_shared() { + let first = tls_config().expect("a client config"); + let second = tls_config().expect("a client config"); - ensure_crypto_provider(); - let second = CryptoProvider::get_default().cloned(); - - assert!(first.is_some()); - assert!(std::sync::Arc::ptr_eq(&first.unwrap(), &second.unwrap())); + assert!(Arc::ptr_eq(&first, &second)); } } diff --git a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs index db5698f8460..05f7d9610d5 100644 --- a/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs +++ b/litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs @@ -1,5 +1,6 @@ -//! Guards the wiring, not just the helper: the dial itself has to install the -//! rustls provider, in a test binary where nothing else has installed one. +//! Guards the wiring, not just the helper: a `wss://` dial through the public +//! API has to resolve its own crypto provider, in a test binary where nothing +//! has installed a process-wide one, and has to leave it uninstalled. use std::collections::HashMap; use std::time::Duration; @@ -7,8 +8,7 @@ use std::time::Duration; use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection; use tokio::net::TcpListener; -#[tokio::test] -async fn dialing_wss_returns_an_error_instead_of_panicking() { +async fn dead_tls_server() -> u16 { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("bind a loopback port"); @@ -23,6 +23,13 @@ async fn dialing_wss_returns_an_error_instead_of_panicking() { } }); + port +} + +#[tokio::test] +async fn dialing_wss_returns_an_error_instead_of_panicking() { + let port = dead_tls_server().await; + let result = ResponsesWebSocketConnection::connect_url( &format!("wss://127.0.0.1:{port}/"), &HashMap::new(), @@ -35,7 +42,7 @@ async fn dialing_wss_returns_an_error_instead_of_panicking() { "a plain TCP server cannot finish a TLS handshake" ); assert!( - rustls::crypto::CryptoProvider::get_default().is_some(), - "the dial is what installs the process-wide provider" + rustls::crypto::CryptoProvider::get_default().is_none(), + "the dial settles its provider on its own connector, not process-wide" ); }