Merge pull request #39530 from BerriAI/litellm_fix_gateway_rustls_provider

fix(ai-gateway): dial upstream WebSockets over an explicit rustls provider

litellm-rust/Cargo.lock

@@ -1415,6 +1415,8 @@ dependencies = [
  "litellm-config",
  "litellm-core",
  "reqwest",
+ "rustls 0.23.42",
+ "rustls-native-certs",
  "serde",
  "serde_json",
  "sha2 0.10.9",

litellm-rust/Cargo.toml

@@ -28,6 +28,8 @@ pythonize = "0.29.0"
 rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "multipart", "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 = { version = "1.0", features = ["float_roundtrip"] }
 sha2 = "0.10"

litellm-rust/crates/ai-gateway/Cargo.toml

@@ -20,6 +20,10 @@ litellm-config.workspace = true
 # reqwest (rustls + json) is used by io/ocr and ships realtime logs to the
 # Python proxy callbacks API.
 reqwest.workspace = true
+# 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

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;

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)
@@ -284,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");

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(),
@@ -324,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");

litellm-rust/crates/ai-gateway/src/io/tls.rs

@@ -0,0 +1,80 @@
+//! 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 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::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::{
+    Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config,
+};
+
+static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
+
+fn build_config() -> Result<ClientConfig, Box<Error>> {
+    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<Arc<ClientConfig>, Box<Error>> {
+    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<R>(
+    request: R,
+) -> Result<(WebSocketStream<MaybeTlsStream<TcpStream>>, Response), Box<Error>>
+where
+    R: IntoClientRequest + Unpin,
+{
+    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 super::build_config;
+
+    #[test]
+    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());
+    }
+}

litellm-rust/crates/ai-gateway/tests/crypto_provider_wiring.rs

@@ -0,0 +1,48 @@
+//! 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;
+
+use litellm_ai_gateway::io::responses_ws::ResponsesWebSocketConnection;
+use tokio::net::TcpListener;
+
+async fn dead_tls_server() -> u16 {
+    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);
+        }
+    });
+
+    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(),
+        Some(Duration::from_secs(10)),
+    )
+    .await;
+
+    assert!(
+        result.is_err(),
+        "a plain TCP server cannot finish a TLS handshake"
+    );
+    assert!(
+        rustls::crypto::CryptoProvider::get_default().is_none(),
+        "the dial settles its provider on its own connector, not process-wide"
+    );
+}