mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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 <noreply@anthropic.com> * fix(http): preserve native WebSocket trust and optional protocols * fix(http): allow plain WebSockets without native roots --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yujonglee
GitHub
parent
6c0585b8f1
commit
cfc651dcd5
12 changed files with 699 additions and 130 deletions
7
litellm-rust/Cargo.lock
generated
7
litellm-rust/Cargo.lock
generated
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@ pub enum Verify {
|
|||
Disabled,
|
||||
CaBundle(PathBuf),
|
||||
BuiltInRoots,
|
||||
NativeRoots,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)?)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<MaybeTlsStream<TcpStream>>;
|
||||
use crate::{client::Client, transport::Error};
|
||||
|
||||
static TLS_CONFIG: OnceLock<Arc<ClientConfig>> = OnceLock::new();
|
||||
|
||||
fn build_tls_config() -> Result<ClientConfig, Box<tokio_tungstenite::tungstenite::Error>> {
|
||||
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<Arc<ClientConfig>, Box<tokio_tungstenite::tungstenite::Error>> {
|
||||
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<reqwest::Upgraded>;
|
||||
|
||||
#[tracing::instrument(name = "litellm.websocket.handshake", level = "debug", skip_all)]
|
||||
pub async fn connect_upstream<R>(
|
||||
request: R,
|
||||
) -> Result<(UpstreamWebSocket, Response), Box<tokio_tungstenite::tungstenite::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)
|
||||
pub async fn connect(
|
||||
client: &Client,
|
||||
url: &str,
|
||||
headers: HeaderMap,
|
||||
) -> Result<UpstreamWebSocket, Error> {
|
||||
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<Url, Error> {
|
||||
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<HeaderMap, Error> {
|
||||
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())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<CertificateDer<'static>>,
|
||||
server_key: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Pki {
|
||||
fn new() -> Self {
|
||||
let mut ca_params = CertificateParams::new(Vec::<String>::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<litellm_http::websocket::UpstreamWebSocket, Error> {
|
||||
let client = HttpClientPool::new(Arc::new(PublicDnsResolver))
|
||||
.client(config, ClientVariant::WebSocket)
|
||||
.unwrap();
|
||||
connect(&client, url, headers).await
|
||||
}
|
||||
|
||||
async fn echo_once<S: AsyncRead + AsyncWrite + Unpin>(
|
||||
stream: S,
|
||||
seen: Arc<Mutex<Option<Request>>>,
|
||||
) {
|
||||
#[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<Mutex<Option<Request>>>) {
|
||||
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<Mutex<Vec<String>>>) {
|
||||
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)),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<String, String>,
|
||||
timeout: Option<Duration>,
|
||||
) -> Result<Self, Error> {
|
||||
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::<HeaderName>()
|
||||
.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::<HeaderName>()
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?,
|
||||
HeaderValue::from_str(value)
|
||||
.map_err(|error| Error::InvalidRequest(error.to_string().into()))?,
|
||||
))
|
||||
})
|
||||
.collect::<Result<HeaderMap, Error>>()?;
|
||||
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))),
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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<Bound<'py, PyAny>> {
|
||||
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 })
|
||||
})
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue