From cfc651dcd5c4daa9f3e12a97b56e7da066680fd0 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 28 Oct 2026 07:32:41 -0700 Subject: [PATCH] fix(http): honor transport settings for Rust WebSockets (#47968) * fix(http): honor transport settings for Rust WebSockets The Responses WebSocket connected through tokio-tungstenite with its own native-roots TLS, so it ignored the configured CA bundle, verify flag, client certificate, proxies and timeouts. The handshake now goes through a pooled reqwest client (ClientVariant::WebSocket, HTTP/1.1, no redirects) and upgrades the response, so it shares the settings every other outbound request uses The upgrade response is checked for the Sec-WebSocket-Accept key, the Upgrade and Connection headers and a subprotocol that was offered Co-Authored-By: Claude Opus 5.5 * fix(http): preserve native WebSocket trust and optional protocols * fix(http): allow plain WebSockets without native roots --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 7 +- litellm-rust/Cargo.toml | 4 +- litellm-rust/crates/http/Cargo.toml | 1 + litellm-rust/crates/http/src/config.rs | 1 + litellm-rust/crates/http/src/pool.rs | 18 +- litellm-rust/crates/http/src/tls.rs | 6 + litellm-rust/crates/http/src/websocket.rs | 179 ++++-- litellm-rust/crates/http/tests/websocket.rs | 550 ++++++++++++++++-- .../crates/inference-responses/Cargo.toml | 1 + .../inference-responses/src/websocket.rs | 49 +- .../inference-responses/tests/responses.rs | 1 + .../src/routes/responses/websocket.rs | 12 +- 12 files changed, 699 insertions(+), 130 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index d37b65be7a3..71ed05c0adf 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3987,6 +3987,7 @@ dependencies = [ "tempfile", "thiserror 2.0.19", "tokio", + "tokio-rustls 0.26.4", "tokio-tungstenite", "tracing", "veil", @@ -7425,11 +7426,7 @@ checksum = "edc5f74e248dc973e0dbb7b74c7e0d6fcc301c694ff50049504004ef4d0cdcd9" dependencies = [ "futures-util", "log", - "rustls 0.23.42", - "rustls-native-certs", - "rustls-pki-types", "tokio", - "tokio-rustls 0.26.4", "tungstenite", ] @@ -7774,8 +7771,6 @@ dependencies = [ "httparse", "log", "rand 0.8.7", - "rustls 0.23.42", - "rustls-pki-types", "sha1 0.10.7", "thiserror 1.0.69", "utf-8", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 95f2da8c073..745b29dd0ff 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -100,10 +100,10 @@ rstest = "0.26.1" wiremock = "0.6.5" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } +rustls-native-certs = "0.8" rustify = "=0.7.0" rustify_derive = "=0.5.5" vaultrs = { version = "=0.8.0", default-features = false, features = ["rustls"] } -rustls-native-certs = "0.8" serde = { version = "1.0", features = ["derive"] } serde_json = { version = "1.0", features = ["float_roundtrip"] } serde_with = { version = "=3.21.0", default-features = false, features = ["std", "macros"] } @@ -118,7 +118,7 @@ thiserror = "2.0" tokenizers = { version = "0.23.1", default-features = false, features = ["onig"] } tiktoken-rs = "0.12.0" tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] } -tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } +tokio-tungstenite = { version = "0.24", default-features = false, features = ["handshake"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" flate2 = "1" diff --git a/litellm-rust/crates/http/Cargo.toml b/litellm-rust/crates/http/Cargo.toml index 3b46825ad9d..0552f9775fd 100644 --- a/litellm-rust/crates/http/Cargo.toml +++ b/litellm-rust/crates/http/Cargo.toml @@ -29,5 +29,6 @@ webpki-roots.workspace = true futures-util.workspace = true rcgen = "0.14.10" tempfile.workspace = true +tokio-rustls = { version = "0.26", default-features = false, features = ["ring"] } rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index 6a2d80e3f85..b37ea6559b8 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -16,6 +16,7 @@ pub enum Verify { Disabled, CaBundle(PathBuf), BuiltInRoots, + NativeRoots, } #[derive(Clone, Debug, PartialEq, Eq, Hash)] diff --git a/litellm-rust/crates/http/src/pool.rs b/litellm-rust/crates/http/src/pool.rs index dfcc9270fe9..bf892b09921 100644 --- a/litellm-rust/crates/http/src/pool.rs +++ b/litellm-rust/crates/http/src/pool.rs @@ -6,7 +6,12 @@ use std::{ use reqwest::dns::Resolve; -use crate::{client::Client, config::HttpClientConfig, error::Error, proxy::EnvironmentProxies}; +use crate::{ + client::Client, + config::{HttpClientConfig, Verify}, + error::Error, + proxy::EnvironmentProxies, +}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientVariant { @@ -14,6 +19,7 @@ pub enum ClientVariant { NoRedirect, Media, UnpinnedMedia, + WebSocket, } const CLIENT_TTL: Duration = Duration::from_secs(3600); @@ -59,6 +65,14 @@ impl HttpClientPool { client_certificate: None, ..config.clone() }, + ClientVariant::WebSocket => HttpClientConfig { + http2: false, + verify: match &config.verify { + Verify::BuiltInRoots => Verify::NativeRoots, + verify => verify.clone(), + }, + ..config.clone() + }, ClientVariant::Provider | ClientVariant::NoRedirect => config.clone(), }; let key = (effective, variant); @@ -91,7 +105,7 @@ impl HttpClientPool { ) -> reqwest::ClientBuilder { match variant { ClientVariant::Provider => builder, - ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia => { + ClientVariant::NoRedirect | ClientVariant::UnpinnedMedia | ClientVariant::WebSocket => { builder.redirect(reqwest::redirect::Policy::none()) } ClientVariant::Media => builder diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index 0b1c44bf6fe..a52ea31af59 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -195,6 +195,12 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), }), + Verify::NativeRoots => { + let native = rustls_native_certs::load_native_certs(); + let mut roots = RootCertStore::empty(); + roots.add_parsable_certificates(native.certs); + builder.with_root_certificates(roots) + } Verify::CaBundle(path) => { builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?) } diff --git a/litellm-rust/crates/http/src/websocket.rs b/litellm-rust/crates/http/src/websocket.rs index e904392152a..685c79cf3a1 100644 --- a/litellm-rust/crates/http/src/websocket.rs +++ b/litellm-rust/crates/http/src/websocket.rs @@ -1,62 +1,135 @@ -use std::{ - io, - sync::{Arc, OnceLock}, +use http::{ + HeaderMap, HeaderValue, StatusCode, + header::{ + CONNECTION, SEC_WEBSOCKET_ACCEPT, SEC_WEBSOCKET_KEY, SEC_WEBSOCKET_PROTOCOL, + SEC_WEBSOCKET_VERSION, UPGRADE, + }, }; - -use rustls::{ClientConfig, RootCertStore}; -use tokio::net::TcpStream; +use reqwest::Url; use tokio_tungstenite::{ - Connector, MaybeTlsStream, WebSocketStream, connect_async_tls_with_config, - tungstenite::{client::IntoClientRequest, error::TlsError, handshake::client::Response}, + WebSocketStream, + tungstenite::{ + handshake::{client::generate_key, derive_accept_key}, + protocol::Role, + }, }; -pub type UpstreamWebSocket = WebSocketStream>; +use crate::{client::Client, transport::Error}; -static TLS_CONFIG: OnceLock> = OnceLock::new(); - -fn build_tls_config() -> Result> { - let native = rustls_native_certs::load_native_certs(); - let mut store = RootCertStore::empty(); - let (added, _ignored) = store.add_parsable_certificates(native.certs); - if added == 0 { - return Err(Box::new(tokio_tungstenite::tungstenite::Error::Io( - io::Error::other(format!( - "no usable native root certificates: {:?}", - native.errors - )), - ))); - } - ClientConfig::builder_with_provider(Arc::new(rustls::crypto::ring::default_provider())) - .with_safe_default_protocol_versions() - .map(|builder| builder.with_root_certificates(store).with_no_client_auth()) - .map_err(|error| { - Box::new(tokio_tungstenite::tungstenite::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_tls_config()?); - Ok(Arc::clone(TLS_CONFIG.get_or_init(|| built))) -} +pub type UpstreamWebSocket = WebSocketStream; #[tracing::instrument(name = "litellm.websocket.handshake", level = "debug", skip_all)] -pub async fn connect_upstream( - request: R, -) -> Result<(UpstreamWebSocket, Response), Box> -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) +pub async fn connect( + client: &Client, + url: &str, + headers: HeaderMap, +) -> Result { + let offered_protocols = headers.get(SEC_WEBSOCKET_PROTOCOL).cloned(); + let key = generate_key(); + let response = client + .get(http_url(url)?) + .headers(handshake_headers(headers, &key)?) + .send() .await - .map_err(Box::new) + .map_err(Error::from_reqwest_before_dispatch)?; + let status = response.status(); + if status != StatusCode::SWITCHING_PROTOCOLS { + return Err(Error::Http { + status: status.as_u16(), + body: response.text().await.unwrap_or_default(), + }); + } + let expected = derive_accept_key(key.as_bytes()); + if response + .headers() + .get(SEC_WEBSOCKET_ACCEPT) + .map(HeaderValue::as_bytes) + != Some(expected.as_bytes()) + { + return Err(Error::Network( + "WebSocket handshake returned a wrong Sec-WebSocket-Accept".into(), + )); + } + validate_response_headers(response.headers(), offered_protocols.as_ref())?; + let upgraded = response.upgrade().await?; + Ok(WebSocketStream::from_raw_socket(upgraded, Role::Client, None).await) +} + +fn validate_response_headers( + headers: &HeaderMap, + offered_protocols: Option<&HeaderValue>, +) -> Result<(), Error> { + if !headers + .get(UPGRADE) + .is_some_and(|value| value.as_bytes().eq_ignore_ascii_case(b"websocket")) + { + return Err(Error::Network( + "WebSocket handshake returned an invalid Upgrade header".into(), + )); + } + if !headers.get_all(CONNECTION).iter().any(|value| { + value.to_str().is_ok_and(|value| { + value + .split(',') + .any(|token| token.trim().eq_ignore_ascii_case("upgrade")) + }) + }) { + return Err(Error::Network( + "WebSocket handshake returned an invalid Connection header".into(), + )); + } + let protocol_matches = match (offered_protocols, headers.get(SEC_WEBSOCKET_PROTOCOL)) { + (_, None) => true, + (Some(offered), Some(selected)) => offered + .to_str() + .ok() + .zip(selected.to_str().ok()) + .is_some_and(|(offered, selected)| { + offered + .split(',') + .any(|protocol| protocol.trim() == selected) + }), + _ => false, + }; + if !protocol_matches { + return Err(Error::Network( + "WebSocket handshake returned an invalid subprotocol".into(), + )); + } + Ok(()) +} + +fn http_url(url: &str) -> Result { + let mut url = Url::parse(url).map_err(not_dispatched)?; + let scheme = match url.scheme() { + "ws" | "http" => "http", + "wss" | "https" => "https", + other => { + return Err(not_dispatched(format!( + "unsupported WebSocket URL scheme {other:?}" + ))); + } + }; + url.set_scheme(scheme) + .map_err(|()| not_dispatched("WebSocket URL cannot be sent over HTTP"))?; + Ok(url) +} + +fn handshake_headers(headers: HeaderMap, key: &str) -> Result { + let key = HeaderValue::from_str(key).map_err(not_dispatched)?; + Ok([ + (CONNECTION, HeaderValue::from_static("Upgrade")), + (UPGRADE, HeaderValue::from_static("websocket")), + (SEC_WEBSOCKET_VERSION, HeaderValue::from_static("13")), + (SEC_WEBSOCKET_KEY, key), + ] + .into_iter() + .fold(headers, |mut headers, (name, value)| { + headers.insert(name, value); + headers + })) +} + +fn not_dispatched(error: impl std::fmt::Display) -> Error { + Error::Connect(error.to_string()) } diff --git a/litellm-rust/crates/http/tests/websocket.rs b/litellm-rust/crates/http/tests/websocket.rs index 2fbaa3994e8..94aba97a84d 100644 --- a/litellm-rust/crates/http/tests/websocket.rs +++ b/litellm-rust/crates/http/tests/websocket.rs @@ -1,59 +1,539 @@ +use std::{ + net::SocketAddr, + path::PathBuf, + sync::{Arc, Mutex}, +}; + use futures_util::{SinkExt, StreamExt}; -use litellm_http::websocket::connect_upstream; -use rstest::rstest; -use tokio::net::TcpListener; +use http::{HeaderMap, HeaderValue}; +use litellm_http::{ + ClientIdentity, ClientVariant, EnvironmentProxies, HttpClientConfig, HttpClientPool, + HttpSettings, Resolution, Verify, media::PublicDnsResolver, transport::Error, + websocket::connect, +}; +use rcgen::{BasicConstraints, CertificateParams, CertifiedIssuer, IsCa, KeyPair}; +use rstest::{Context, rstest}; +use rustls::{ + RootCertStore, ServerConfig, + crypto::ring, + pki_types::{CertificateDer, PrivateKeyDer}, + server::WebPkiClientVerifier, +}; +use tokio::{ + io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + net::{TcpListener, TcpStream}, +}; +use tokio_rustls::TlsAcceptor; use tokio_tungstenite::{ accept_hdr_async, tungstenite::{ Message, - client::IntoClientRequest, handshake::server::{Request, Response}, }, }; +struct Pki { + directory: tempfile::TempDir, + ca: CertifiedIssuer<'static, KeyPair>, + server_chain: Vec>, + server_key: Vec, +} + +impl Pki { + fn new() -> Self { + let mut ca_params = CertificateParams::new(Vec::::new()).unwrap(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca = CertifiedIssuer::self_signed(ca_params, KeyPair::generate().unwrap()).unwrap(); + let server_key = KeyPair::generate().unwrap(); + let server = CertificateParams::new(vec!["localhost".into()]) + .unwrap() + .signed_by(&server_key, &ca) + .unwrap(); + let client_key = KeyPair::generate().unwrap(); + let client = CertificateParams::new(vec!["client".into()]) + .unwrap() + .signed_by(&client_key, &ca) + .unwrap(); + let directory = tempfile::tempdir().unwrap(); + std::fs::write(directory.path().join("ca.pem"), ca.pem()).unwrap(); + std::fs::write( + directory.path().join("client.pem"), + format!("{}{}", client.pem(), client_key.serialize_pem()), + ) + .unwrap(); + Self { + directory, + server_chain: vec![server.der().clone()], + server_key: server_key.serialize_der(), + ca, + } + } + + fn path(&self, name: &str) -> PathBuf { + self.directory.path().join(name) + } + + fn acceptor(&self, require_client_certificate: bool) -> TlsAcceptor { + let provider = Arc::new(ring::default_provider()); + let builder = ServerConfig::builder_with_provider(Arc::clone(&provider)) + .with_safe_default_protocol_versions() + .unwrap(); + let builder = if require_client_certificate { + let mut roots = RootCertStore::empty(); + roots.add(self.ca.der().clone()).unwrap(); + builder.with_client_cert_verifier( + WebPkiClientVerifier::builder_with_provider(Arc::new(roots), provider) + .build() + .unwrap(), + ) + } else { + builder.with_no_client_auth() + }; + let mut config = builder + .with_single_cert( + self.server_chain.clone(), + PrivateKeyDer::try_from(self.server_key.clone()).unwrap(), + ) + .unwrap(); + config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()]; + TlsAcceptor::from(Arc::new(config)) + } +} + +fn config(verify: Verify) -> HttpClientConfig { + HttpClientConfig { + verify, + http2: true, + ..Resolution::from(&HttpSettings::default()).config + } +} + +fn proxied_through(proxy: SocketAddr) -> EnvironmentProxies { + EnvironmentProxies::from_environment(&move |name: &str| { + (name == "HTTPS_PROXY").then(|| format!("http://{proxy}")) + }) +} + +async fn open( + config: &HttpClientConfig, + url: &str, + headers: HeaderMap, +) -> Result { + let client = HttpClientPool::new(Arc::new(PublicDnsResolver)) + .client(config, ClientVariant::WebSocket) + .unwrap(); + connect(&client, url, headers).await +} + +async fn echo_once( + stream: S, + seen: Arc>>, +) { + #[expect( + clippy::result_large_err, + reason = "tungstenite requires an unboxed handshake error response" + )] + let Ok(mut socket) = accept_hdr_async(stream, |request: &Request, response: Response| { + *seen.lock().unwrap() = Some(request.clone()); + Ok(response) + }) + .await + else { + return; + }; + if let Some(Ok(message)) = socket.next().await { + let _ = socket.send(message).await; + } +} + +async fn serve_tls( + pki: &Pki, + require_client_certificate: bool, +) -> (SocketAddr, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let acceptor = pki.acceptor(require_client_certificate); + let seen = Arc::new(Mutex::new(None)); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + while let Ok((stream, _)) = listener.accept().await { + if let Ok(stream) = acceptor.accept(stream).await { + echo_once(stream, Arc::clone(&recorded)).await; + } + } + }); + (address, seen) +} + +async fn read_head(stream: &mut TcpStream) -> String { + let mut head = Vec::new(); + while !head.ends_with(b"\r\n\r\n") { + let mut byte = [0u8; 1]; + if stream.read(&mut byte).await.unwrap() == 0 { + break; + } + head.push(byte[0]); + } + String::from_utf8(head).unwrap() +} + +async fn serve_tunnel_proxy(target: SocketAddr) -> (SocketAddr, Arc>>) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let requests = Arc::new(Mutex::new(Vec::new())); + let seen = Arc::clone(&requests); + tokio::spawn(async move { + while let Ok((mut client, _)) = listener.accept().await { + let head = read_head(&mut client).await; + seen.lock() + .unwrap() + .push(head.lines().next().unwrap_or_default().to_owned()); + let mut upstream = TcpStream::connect(target).await.unwrap(); + client + .write_all(b"HTTP/1.1 200 Connection established\r\n\r\n") + .await + .unwrap(); + let _ = tokio::io::copy_bidirectional(&mut client, &mut upstream).await; + } + }); + (address, requests) +} + +async fn serve_raw(response: &'static str) -> SocketAddr { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + read_head(&mut stream).await; + stream.write_all(response.as_bytes()).await.unwrap(); + let _ = stream.read(&mut [0u8; 1]).await; + }); + address +} + +async fn round_trip(socket: &mut litellm_http::websocket::UpstreamWebSocket) { + socket.send(Message::Text("ping".into())).await.unwrap(); + assert_eq!( + socket.next().await.unwrap().unwrap(), + Message::Text("ping".into()) + ); +} + #[rstest] #[case::without_query("/responses")] #[case::with_query("/responses?model=test-model")] #[tokio::test] -async fn connects_with_caller_headers_and_exchanges_frames(#[case] path: &'static str) { +async fn plain_websocket_sends_caller_headers_and_path(#[case] path: &'static str) { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); - let server = tokio::spawn(async move { + let seen = Arc::new(Mutex::new(None)); + let recorded = Arc::clone(&seen); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + echo_once(stream, recorded).await; + }); + let headers = HeaderMap::from_iter([( + http::header::AUTHORIZATION, + HeaderValue::from_static("Bearer test-key"), + )]); + + let mut socket = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}{path}"), + headers, + ) + .await + .unwrap(); + + round_trip(&mut socket).await; + let request = seen.lock().unwrap().take().unwrap(); + assert_eq!(request.uri().path_and_query().unwrap().as_str(), path); + assert_eq!(request.headers()["authorization"], "Bearer test-key"); +} + +#[rstest] +#[case::ca_bundle_trusts_the_private_ca(|pki: &Pki| Verify::CaBundle(pki.path("ca.pem")), true)] +#[case::ssl_verify_off_skips_verification(|_: &Pki| Verify::Disabled, true)] +#[case::built_in_roots_reject_the_private_ca(|_: &Pki| Verify::BuiltInRoots, false)] +#[tokio::test] +async fn wss_honors_ssl_verify(#[case] verify: fn(&Pki) -> Verify, #[case] trusted: bool) { + let pki = Pki::new(); + let (address, _) = serve_tls(&pki, false).await; + + let result = open( + &config(verify(&pki)), + &format!("wss://localhost:{}/responses", address.port()), + HeaderMap::new(), + ) + .await; + + match result { + Ok(mut socket) if trusted => round_trip(&mut socket).await, + Err(Error::Connect(_)) if !trusted => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[case::plain_without_roots(false, false, true)] +#[case::tls_with_roots(true, true, true)] +#[case::tls_without_roots(true, false, false)] +#[tokio::test] +async fn websocket_uses_system_roots( + #[context] context: Context, + #[case] tls: bool, + #[case] trusted_roots: bool, + #[case] connects: bool, +) { + const CHILD_URL: &str = "LITELLM_TEST_SYSTEM_ROOTS_WSS_URL"; + if let Ok(url) = std::env::var(CHILD_URL) { + let result = open(&config(Verify::BuiltInRoots), &url, HeaderMap::new()).await; + match result { + Ok(mut socket) if connects => round_trip(&mut socket).await, + Err(Error::Connect(_)) if !connects => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } + return; + } + + let pki = Pki::new(); + let url = if tls { + let (address, _) = serve_tls(&pki, false).await; + format!("wss://localhost:{}/responses", address.port()) + } else { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + echo_once(stream, Arc::new(Mutex::new(None))).await; + }); + format!("ws://{address}/responses") + }; + let ca = if trusted_roots { + pki.path("ca.pem") + } else { + let empty = pki.path("empty.pem"); + std::fs::write(&empty, "").unwrap(); + empty + }; + let test = format!( + "{}::case_{}_{}", + context.name, + context.case.unwrap() + 1, + context.description.unwrap() + ); + let output = tokio::task::spawn_blocking(move || { + std::process::Command::new(std::env::current_exe().unwrap()) + .args(["--exact", &test, "--nocapture"]) + .env(CHILD_URL, url) + .env("SSL_CERT_FILE", ca) + .env("SSL_CERT_DIR", "") + .output() + .unwrap() + }) + .await + .unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + assert!( + output.status.success() && stdout.contains("1 passed"), + "system-root WebSocket child failed:\n{stdout}\n{}", + String::from_utf8_lossy(&output.stderr) + ); +} + +#[rstest] +#[case::with_certificate(true)] +#[case::without_certificate(false)] +#[tokio::test] +async fn wss_presents_the_configured_client_certificate(#[case] configured: bool) { + let pki = Pki::new(); + let (address, _) = serve_tls(&pki, true).await; + let config = HttpClientConfig { + client_certificate: configured.then(|| ClientIdentity::Pem(pki.path("client.pem"))), + ..config(Verify::CaBundle(pki.path("ca.pem"))) + }; + + let result = open( + &config, + &format!("wss://localhost:{}/responses", address.port()), + HeaderMap::new(), + ) + .await; + + match result { + Ok(mut socket) if configured => round_trip(&mut socket).await, + Err(_) if !configured => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } +} + +#[rstest] +#[tokio::test] +async fn wss_tunnels_through_the_https_proxy() { + let pki = Pki::new(); + let (target, _) = serve_tls(&pki, false).await; + let (proxy, requests) = serve_tunnel_proxy(target).await; + let config = HttpClientConfig { + proxies: proxied_through(proxy), + ..config(Verify::CaBundle(pki.path("ca.pem"))) + }; + + let mut socket = open( + &config, + &format!("wss://localhost:{}/responses", target.port()), + HeaderMap::new(), + ) + .await + .unwrap(); + + round_trip(&mut socket).await; + assert_eq!( + *requests.lock().unwrap(), + [format!("CONNECT localhost:{} HTTP/1.1", target.port())] + ); +} + +#[rstest] +#[tokio::test] +async fn refused_upgrade_reports_the_status_and_body() { + let address = + serve_raw("HTTP/1.1 403 Forbidden\r\nContent-Length: 6\r\nConnection: close\r\n\r\ndenied") + .await; + + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + HeaderMap::new(), + ) + .await; + + assert_eq!( + result.map(drop), + Err(Error::Http { + status: 403, + body: "denied".into() + }) + ); +} + +#[rstest] +#[tokio::test] +async fn upgrade_with_a_wrong_accept_key_is_rejected() { + let address = serve_raw( + "HTTP/1.1 101 Switching Protocols\r\nConnection: Upgrade\r\nUpgrade: websocket\r\nSec-WebSocket-Accept: wrong\r\n\r\n", + ) + .await; + + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + HeaderMap::new(), + ) + .await; + + assert!(matches!(result, Err(Error::Network(_)))); +} + +#[rstest] +#[case::non_websocket_scheme("ftp://localhost/responses")] +#[case::unparseable("not a url")] +#[tokio::test] +async fn invalid_url_is_rejected_before_dispatch(#[case] url: &str) { + let result = open(&config(Verify::BuiltInRoots), url, HeaderMap::new()).await; + + assert!(matches!(result, Err(Error::Connect(_)))); +} + +#[rstest] +#[case::missing_upgrade(None, Some("Upgrade"), None, None, false)] +#[case::wrong_upgrade(Some("http"), Some("Upgrade"), None, None, false)] +#[case::missing_connection(Some("websocket"), None, None, None, false)] +#[case::wrong_connection(Some("websocket"), Some("keep-alive"), None, None, false)] +#[case::unsolicited_protocol(Some("websocket"), Some("Upgrade"), None, Some("other"), false)] +#[case::unoffered_protocol( + Some("websocket"), + Some("Upgrade"), + Some("events"), + Some("other"), + false +)] +#[case::declined_optional_protocol(Some("websocket"), Some("Upgrade"), Some("events"), None, true)] +#[case::matching_protocol( + Some("websocket"), + Some("Upgrade"), + Some("events"), + Some("events"), + true +)] +#[case::multiple_offers( + Some("websocket"), + Some("Upgrade"), + Some("events, realtime"), + Some("realtime"), + true +)] +#[case::case_insensitive_upgrade(Some("WebSocket"), Some("uPgRaDe"), None, None, true)] +#[case::connection_tokens(Some("websocket"), Some("keep-alive, Upgrade"), None, None, true)] +#[tokio::test] +async fn upgrade_validates_headers_and_negotiated_protocol( + #[case] upgrade: Option<&'static str>, + #[case] connection: Option<&'static str>, + #[case] offered: Option<&'static str>, + #[case] selected: Option<&'static str>, + #[case] accepted: bool, +) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); #[expect( clippy::result_large_err, reason = "tungstenite requires an unboxed handshake error response" )] - let mut socket = accept_hdr_async(stream, move |request: &Request, response: Response| { - assert_eq!(request.uri().path_and_query().unwrap().as_str(), path); - assert_eq!(request.headers()["authorization"], "Bearer test-key"); + let result = accept_hdr_async(stream, move |_: &Request, mut response: Response| { + response.headers_mut().remove("upgrade"); + response.headers_mut().remove("connection"); + if let Some(value) = upgrade { + response + .headers_mut() + .insert("upgrade", HeaderValue::from_static(value)); + } + if let Some(value) = connection { + response + .headers_mut() + .insert("connection", HeaderValue::from_static(value)); + } + if let Some(value) = selected { + response + .headers_mut() + .insert("sec-websocket-protocol", HeaderValue::from_static(value)); + } Ok(response) }) - .await - .unwrap(); - assert_eq!( - socket.next().await.unwrap().unwrap(), - Message::Text("request".into()) - ); - socket.send(Message::Text("response".into())).await.unwrap(); - socket.close(None).await.unwrap(); + .await; + let Ok(mut socket) = result else { + return; + }; + if let Some(Ok(message)) = socket.next().await { + let _ = socket.send(message).await; + } }); - let mut request = format!("ws://{address}{path}") - .into_client_request() - .unwrap(); - request - .headers_mut() - .insert("authorization", "Bearer test-key".parse().unwrap()); - let (mut socket, response) = connect_upstream(request).await.unwrap(); - assert_eq!(response.status(), 101); - socket.send(Message::Text("request".into())).await.unwrap(); - assert_eq!( - socket.next().await.unwrap().unwrap(), - Message::Text("response".into()) - ); - assert!(matches!( - socket.next().await.unwrap().unwrap(), - Message::Close(_) - )); - server.await.unwrap(); + let headers = offered.map_or_else(HeaderMap::new, |value| { + HeaderMap::from_iter([( + http::header::SEC_WEBSOCKET_PROTOCOL, + HeaderValue::from_static(value), + )]) + }); + let result = open( + &config(Verify::BuiltInRoots), + &format!("ws://{address}/responses"), + headers, + ) + .await; + + match result { + Ok(mut socket) if accepted => round_trip(&mut socket).await, + Err(Error::Network(_)) if !accepted => {} + other => panic!("unexpected handshake outcome: {:?}", other.map(drop)), + } } diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 793483ee76d..3dcd9976e0a 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-host-native.workspace = true diff --git a/litellm-rust/crates/inference-responses/src/websocket.rs b/litellm-rust/crates/inference-responses/src/websocket.rs index 213a9a95831..c2ca89ecbaf 100644 --- a/litellm-rust/crates/inference-responses/src/websocket.rs +++ b/litellm-rust/crates/inference-responses/src/websocket.rs @@ -1,12 +1,14 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use futures_util::{SinkExt, StreamExt}; -use litellm_http::websocket::{UpstreamWebSocket, connect_upstream}; +use litellm_http::{ + Client, + websocket::{self, UpstreamWebSocket}, +}; use tokio::sync::Mutex; use tokio_tungstenite::tungstenite::{ Message, - client::IntoClientRequest, - http::{HeaderName, HeaderValue}, + http::{HeaderMap, HeaderName, HeaderValue}, }; use super::Error; @@ -24,42 +26,33 @@ impl ResponsesWebSocketConnection { fields(outcome) )] pub async fn connect_url( + client: &Client, url: &str, headers: &HashMap, timeout: Option, ) -> Result { litellm_inference::diagnostic::operation("litellm.websocket.connect_url", async { - let mut request = url.into_client_request().map_err(|error| { - Error::Transport(litellm_http::transport::Error::Network(error.to_string())) - })?; - for (name, value) in headers { - let header_name = name - .parse::() - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - let header_value = HeaderValue::from_str(value) - .map_err(|error| Error::InvalidRequest(error.to_string().into()))?; - request.headers_mut().insert(header_name, header_value); - } - let connect = connect_upstream(request); - let result = match timeout { + let headers = headers + .iter() + .map(|(name, value)| { + Ok(( + name.parse::() + .map_err(|error| Error::InvalidRequest(error.to_string().into()))?, + HeaderValue::from_str(value) + .map_err(|error| Error::InvalidRequest(error.to_string().into()))?, + )) + }) + .collect::>()?; + let connect = websocket::connect(client, url, headers); + let socket = match timeout { Some(timeout) => tokio::time::timeout(timeout, connect).await.map_err(|_| { Error::Transport(litellm_http::transport::Error::Network( "Responses WebSocket connection timed out".into(), )) })?, None => connect.await, - }; - let (socket, _) = result.map_err(|error| match *error { - tokio_tungstenite::tungstenite::Error::Http(response) => { - Error::Transport(litellm_http::transport::Error::Http { - status: response.status().as_u16(), - body: String::new(), - }) - } - other => { - Error::Transport(litellm_http::transport::Error::Network(other.to_string())) - } - })?; + } + .map_err(Error::Transport)?; Ok(Self { socket: Arc::new(Mutex::new(Some(socket))), }) diff --git a/litellm-rust/crates/inference-responses/tests/responses.rs b/litellm-rust/crates/inference-responses/tests/responses.rs index ae07a71dabd..6d39d55d33c 100644 --- a/litellm-rust/crates/inference-responses/tests/responses.rs +++ b/litellm-rust/crates/inference-responses/tests/responses.rs @@ -538,6 +538,7 @@ async fn websocket_operations_trace_outcomes_without_capturing_frames_or_credent .logger() .instrument(async { let connection = ResponsesWebSocketConnection::connect_url( + &litellm_http::Client::for_test(), &format!("ws://{address}/responses"), &std::collections::HashMap::from([( "authorization".into(), diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs index 615a838a597..c89185dc21f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/websocket.rs @@ -1,7 +1,9 @@ -use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; use pyo3::prelude::*; use serde_json::Value; +use litellm_http::ClientVariant; +use litellm_inference_responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection; + use crate::{ errors::route_error_to_pyerr, marshal::{marshal_headers, optional_timeout}, @@ -25,10 +27,12 @@ impl ResponsesWebSocketConnection { ) -> PyResult> { let headers = marshal_headers(headers)?; let timeout = optional_timeout(timeout_seconds); + let client = crate::http::host_client(py, ClientVariant::WebSocket)?; crate::execution::run_async_value(py, async move { - let inner = RustResponsesWebSocketConnection::connect_url(&url, &headers, timeout) - .await - .map_err(route_error_to_pyerr)?; + let inner = + RustResponsesWebSocketConnection::connect_url(&client, &url, &headers, timeout) + .await + .map_err(route_error_to_pyerr)?; Ok(ResponsesWebSocketConnection { inner }) }) }