From a75c982d9caf9dbe8238c2b902dbc99fe2f7d5d4 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 4 Nov 2026 18:39:00 +0000 Subject: [PATCH] fix(gateway): preserve required forwarding headers and separate upgrades --- litellm-rust/FORWARDING.md | 24 ++++++ .../crates/gateway-forward/src/forward.rs | 49 ++++++++---- .../crates/gateway-forward/src/rewrite.rs | 20 +++++ .../crates/gateway-forward/src/settings.rs | 1 + .../crates/gateway-forward/src/websocket.rs | 5 +- .../gateway-forward/tests/client_identity.rs | 14 +++- .../crates/gateway-forward/tests/exchange.rs | 78 +++++++++++++++++++ .../crates/gateway-forward/tests/framing.rs | 5 +- .../crates/gateway-forward/tests/websocket.rs | 8 +- 9 files changed, 180 insertions(+), 24 deletions(-) create mode 100644 litellm-rust/FORWARDING.md diff --git a/litellm-rust/FORWARDING.md b/litellm-rust/FORWARDING.md new file mode 100644 index 00000000000..58ac3d6300b --- /dev/null +++ b/litellm-rust/FORWARDING.md @@ -0,0 +1,24 @@ +# Python forwarding compatibility + +The opt-in forwarding gateway serves HTTP/1.1 and WebSockets in front of Python. Python still owns every route and policy. Its protocol contract and environment settings are in [gateway-forward/AGENTS.md](crates/gateway-forward/AGENTS.md) + +## Known differences from uvicorn with h11 + +These cases remain executable as ignored socket tests. They are not passing parity checks + +| Input | Direct Python | Through Rust | Regression test | +|---|---|---|---| +| Literal `#` in a request path or query | Keeps every byte | `http::Uri` truncates at `#`, potentially changing the route or query. Must be fixed before shipping | [forwards_an_unescaped_hash_in_the_target](crates/gateway-forward/tests/exchange.rs) | +| Header after a removed hop-by-hop field | Keeps the remaining header order | `HeaderMap::remove` can move a later header | [dropping_a_hop_by_hop_header_keeps_the_others_in_order](crates/gateway-forward/tests/exchange.rs) | +| Repeated header names interleaved with other names | Keeps the interleaving | Groups values by name, keeping their order within that name | [interleaved_repeated_headers_reach_python_interleaved](crates/gateway-forward/tests/exchange.rs) | +| Obsolete folded header, such as `x-a: b\r\n c` | Unfolds to `b c` | Rejects with 400 | [a_folded_header_reaches_python_unfolded](crates/gateway-forward/tests/exchange.rs) | +| Identical Content-Length list, such as `7, 7` | Accepts length 7 | Rejects with 400 | [a_content_length_list_of_one_value_reaches_python](crates/gateway-forward/tests/framing.rs) | +| HTTP/1.0 request with chunked transfer coding | Decodes the body | Rejects with 400 | [an_http10_chunked_upload_reaches_python](crates/gateway-forward/tests/framing.rs) | + +Clients can avoid the parser rejections by sending unfolded headers, one Content-Length value, and HTTP/1.1 for chunked uploads. Integrations must not depend on header ordering across different field names + +Run `cargo test -p litellm-gateway-forward --test exchange --test framing -- --ignored` from this directory to reproduce the differences + +## Connection header handling + +The gateway strips standard hop-by-hop fields and fields named by Connection, except routing, representation, credential, forwarding, and WebSocket handshake headers. A client cannot nominate those end-to-end fields for removal. The private `x-litellm-rust-peer` always replaces any client-supplied copy diff --git a/litellm-rust/crates/gateway-forward/src/forward.rs b/litellm-rust/crates/gateway-forward/src/forward.rs index 20a2af4a86b..69fb63ada7d 100644 --- a/litellm-rust/crates/gateway-forward/src/forward.rs +++ b/litellm-rust/crates/gateway-forward/src/forward.rs @@ -87,27 +87,48 @@ impl Forward { client: ClientSocket, mut request: Request, ) -> Result { - let downstream = websocket::take_handshake(&self.connections, &mut request)?; - let request = - rewrite::upstream_request(request, &self.upstream, client, downstream.is_some())?; - let response = self.client.request(request); - let response = match downstream { - Some(_) => self - .connections - .until_closed(response) - .await - .ok_or(Rejection::ShuttingDown)?, - None => response.await, + if let Some(downstream) = websocket::take_handshake(&self.connections, &mut request)? { + return self.forward_websocket(client, request, downstream).await; } - .map_err(Rejection::from_upstream)?; + let request = rewrite::upstream_request(request, &self.upstream, client, false)?; + let response = self + .client + .request(request) + .await + .map_err(Rejection::from_upstream)?; + http_response(response) + } + + async fn forward_websocket( + &self, + client: ClientSocket, + request: Request, + downstream: hyper::upgrade::OnUpgrade, + ) -> Result { + let request = rewrite::upstream_request(request, &self.upstream, client, true)?; + let response = self + .connections + .until_closed(self.client.request(request)) + .await + .ok_or(Rejection::ShuttingDown)? + .map_err(Rejection::from_upstream)?; if response.status() == StatusCode::SWITCHING_PROTOCOLS { return websocket::switch(&self.connections, downstream, response).await; } - Ok(rewrite::downstream_response(response) - .map(|body| Body::new(Relayed { body, failed: None }))) + http_response(response) } } +fn http_response(response: Response) -> Result { + if response.status() == StatusCode::SWITCHING_PROTOCOLS { + return Err(Rejection::UnrequestedUpgrade); + } + Ok( + rewrite::downstream_response(response) + .map(|body| Body::new(Relayed { body, failed: None })), + ) +} + /// Python's response body on its way to the client. struct Relayed { body: Incoming, diff --git a/litellm-rust/crates/gateway-forward/src/rewrite.rs b/litellm-rust/crates/gateway-forward/src/rewrite.rs index 058a4d3bac2..fa22bf19037 100644 --- a/litellm-rust/crates/gateway-forward/src/rewrite.rs +++ b/litellm-rust/crates/gateway-forward/src/rewrite.rs @@ -125,6 +125,7 @@ fn strip_hop_by_hop(headers: &mut HeaderMap) { .filter_map(|value| value.to_str().ok()) .flat_map(|value| value.split(',')) .filter_map(|name| header::HeaderName::from_bytes(name.trim().as_bytes()).ok()) + .filter(|name| !must_forward(name)) .collect(); for name in named { headers.remove(name); @@ -134,6 +135,25 @@ fn strip_hop_by_hop(headers: &mut HeaderMap) { } } +fn must_forward(name: &header::HeaderName) -> bool { + matches!( + name.as_str(), + "host" + | "authorization" + | "cookie" + | "set-cookie" + | "www-authenticate" + | "authentication-info" + | "api-key" + | "x-api-key" + | "litellm-api-key" + | "forwarded" + | PEER_HEADER + ) || name.as_str().starts_with("content-") + || name.as_str().starts_with("x-forwarded-") + || name.as_str().starts_with("sec-websocket-") +} + // Stripping drops the handshake's Connection and Upgrade, so each hop gets them back. fn restore_websocket_handshake(headers: &mut HeaderMap) { headers.insert(header::CONNECTION, HeaderValue::from_static("upgrade")); diff --git a/litellm-rust/crates/gateway-forward/src/settings.rs b/litellm-rust/crates/gateway-forward/src/settings.rs index 9fb5e73497f..de01a96c519 100644 --- a/litellm-rust/crates/gateway-forward/src/settings.rs +++ b/litellm-rust/crates/gateway-forward/src/settings.rs @@ -47,6 +47,7 @@ struct Environment { /// Everything forward mode runs with. The Python supervisor resolves each value from the settings /// Python actually runs with, so none of them is guessed or defaulted here. +/// Setting `LITELLM_FORWARD_UPSTREAM` requires every other `LITELLM_FORWARD_*` setting below. #[derive(Debug, Clone, PartialEq)] pub struct Settings { pub upstream: Upstream, diff --git a/litellm-rust/crates/gateway-forward/src/websocket.rs b/litellm-rust/crates/gateway-forward/src/websocket.rs index 3a5a2e91fed..84d89c75607 100644 --- a/litellm-rust/crates/gateway-forward/src/websocket.rs +++ b/litellm-rust/crates/gateway-forward/src/websocket.rs @@ -31,12 +31,9 @@ pub(crate) fn take_handshake( pub(crate) async fn switch( connections: &Connections, - downstream: Option, + downstream: OnUpgrade, mut response: Response, ) -> Result { - let Some(downstream) = downstream else { - return Err(Rejection::UnrequestedUpgrade); - }; if !rewrite::is_websocket_handshake(response.headers()) { return Err(Rejection::UnrequestedUpgrade); } diff --git a/litellm-rust/crates/gateway-forward/tests/client_identity.rs b/litellm-rust/crates/gateway-forward/tests/client_identity.rs index 1e31edb93c1..c484c89dc1d 100644 --- a/litellm-rust/crates/gateway-forward/tests/client_identity.rs +++ b/litellm-rust/crates/gateway-forward/tests/client_identity.rs @@ -16,20 +16,26 @@ fn ok() -> Response { } #[rstest] +#[case::ordinary("")] +#[case::connection_nominated( + "connection: X-Forwarded-For, X-Forwarded-Proto, Forwarded, X-Litellm-Rust-Peer\r\n" +)] #[tokio::test] -async fn python_sees_forwarding_headers_untouched_and_the_tcp_peer_rust_saw() { +async fn python_sees_forwarding_headers_untouched_and_the_tcp_peer_rust_saw( + #[case] connection: &str, +) { let (upstream, mut received) = recording(ok).await; let gateway = Gateway::start(upstream).await; let mut raw = gateway.connect().await; - raw.send( + raw.send(format!( "GET /sso/callback HTTP/1.1\r\n\ host: public.example\r\n\ x-forwarded-for: 192.0.2.1\r\n\ x-forwarded-for: 192.0.2.2\r\n\ x-forwarded-proto: https\r\n\ forwarded: for=192.0.2.1;proto=https\r\n\ - x-litellm-rust-peer: 10.1.2.3:1\r\n\r\n", - ) + x-litellm-rust-peer: 10.1.2.3:1\r\n{connection}\r\n", + )) .await; assert_eq!( within(raw.response_head()).await.unwrap().status, diff --git a/litellm-rust/crates/gateway-forward/tests/exchange.rs b/litellm-rust/crates/gateway-forward/tests/exchange.rs index ed805ea32d3..fdc20def97b 100644 --- a/litellm-rust/crates/gateway-forward/tests/exchange.rs +++ b/litellm-rust/crates/gateway-forward/tests/exchange.rs @@ -90,6 +90,84 @@ async fn forwards_the_request_as_the_client_sent_it_minus_hop_headers() { assert_eq!(request.body, "payload"); } +#[rstest] +#[case::host("HoSt", "public.example")] +#[case::length("Content-Length", "7")] +#[case::authorization("AUTHORIZATION", "Bearer test")] +#[case::cookie("Cookie", "session=test")] +#[case::api_key("X-Api-Key", "test")] +#[case::azure_key("Api-Key", "test")] +#[case::litellm_key("Litellm-Api-Key", "test")] +#[case::forwarded("Forwarded", "for=192.0.2.1;proto=https")] +#[case::forwarded_for("X-Forwarded-For", "192.0.2.1")] +#[case::forwarded_proto("X-Forwarded-Proto", "https")] +#[case::forwarded_host("X-Forwarded-Host", "public.example")] +#[case::content_type("Content-Type", "application/octet-stream")] +#[case::content_encoding("Content-Encoding", "identity")] +#[tokio::test] +async fn connection_cannot_remove_end_to_end_headers(#[case] name: &str, #[case] value: &str) { + let (upstream, mut received) = recording(ok).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + let host = if name.eq_ignore_ascii_case("host") { + "" + } else { + "host: public.example\r\n" + }; + let length = if name.eq_ignore_ascii_case("content-length") { + "" + } else { + "content-length: 7\r\n" + }; + raw.send(format!( + "POST / HTTP/1.1\r\n{host}{length}{name}: {value}\r\nconnection: x-remove\r\nconnection: {name}, keep-alive\r\nx-remove: remove\r\n\r\npayload" + )).await; + assert_eq!( + within(raw.response_head()).await.unwrap().status, + StatusCode::OK + ); + let request = within(received.recv()).await.unwrap(); + assert_eq!(request.headers[name], value); + assert_eq!(request.body, "payload"); + assert!(!request.headers.contains_key("connection")); + assert!(!request.headers.contains_key("x-remove")); +} + +#[rstest] +#[tokio::test] +async fn connection_cannot_remove_response_credentials_or_representation_metadata() { + let (upstream, _received) = recording(|| { + Response::builder() + .status(StatusCode::UNAUTHORIZED) + .header("connection", "Content-Length, Set-Cookie, WWW-Authenticate, Authentication-Info, Content-Type, Content-Encoding, x-remove") + .header("content-length", "2") + .header("set-cookie", "session=; Max-Age=0") + .header("www-authenticate", "Bearer") + .header("authentication-info", "nextnonce=test") + .header("content-type", "text/plain") + .header("content-encoding", "identity") + .header("x-remove", "remove") + .body(Body::from("ok")) + .unwrap() + }).await; + let gateway = Gateway::start(upstream).await; + let mut raw = gateway.connect().await; + raw.send("GET / HTTP/1.1\r\nhost: x\r\n\r\n").await; + let head = within(raw.response_head()).await.unwrap(); + assert_eq!(head.status, StatusCode::UNAUTHORIZED); + assert_eq!(head.content_length(), Some(2)); + assert_eq!(head.header("set-cookie"), Some(&b"session=; Max-Age=0"[..])); + assert_eq!(head.header("www-authenticate"), Some(&b"Bearer"[..])); + assert_eq!( + head.header("authentication-info"), + Some(&b"nextnonce=test"[..]) + ); + assert_eq!(head.header("content-type"), Some(&b"text/plain"[..])); + assert_eq!(head.header("content-encoding"), Some(&b"identity"[..])); + assert_eq!(head.header("x-remove"), None); + assert_eq!(raw.body(2).await, b"ok"); +} + #[rstest] #[case::get("GET")] #[case::post("POST")] diff --git a/litellm-rust/crates/gateway-forward/tests/framing.rs b/litellm-rust/crates/gateway-forward/tests/framing.rs index 7f0a44ff644..32a2c3f2b8d 100644 --- a/litellm-rust/crates/gateway-forward/tests/framing.rs +++ b/litellm-rust/crates/gateway-forward/tests/framing.rs @@ -80,7 +80,10 @@ async fn a_request_framed_two_ways_reaches_python_framed_by_transfer_encoding_an Some(&b"chunked"[..]) ); assert_eq!(upload.body, b"hello"); - assert!(received.try_recv().is_err(), "python ran the smuggled request"); + assert!( + received.try_recv().is_err(), + "python ran the smuggled request" + ); } #[rstest] diff --git a/litellm-rust/crates/gateway-forward/tests/websocket.rs b/litellm-rust/crates/gateway-forward/tests/websocket.rs index b2af47382ea..3cda99fd6a5 100644 --- a/litellm-rust/crates/gateway-forward/tests/websocket.rs +++ b/litellm-rust/crates/gateway-forward/tests/websocket.rs @@ -158,7 +158,13 @@ async fn open(gateway: &Gateway, handshake: String) -> (TcpStream, Head) { #[rstest] #[tokio::test] async fn relays_messages_both_ways( - #[values("Upgrade", "keep-alive, Upgrade", "upgrade")] connection: &'static str, + #[values( + "Upgrade", + "keep-alive, Upgrade", + "upgrade", + "Upgrade, Sec-WebSocket-Key, Sec-WebSocket-Version" + )] + connection: &'static str, #[values( Message::text("hello"), Message::binary(vec![0, 1, 2]),