fix(gateway): preserve required forwarding headers and separate upgrades

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

litellm-rust/crates/gateway-forward/src/forward.rs

@@ -87,27 +87,48 @@ impl Forward {
         client: ClientSocket,
         mut request: Request,
     ) -> Result<Response, Rejection> {
-        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<Response, Rejection> {
+        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<Incoming>) -> Result<Response, Rejection> {
+    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,

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"));

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,

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<OnUpgrade>,
+    downstream: OnUpgrade,
     mut response: Response<Incoming>,
 ) -> Result<Response, Rejection> {
-    let Some(downstream) = downstream else {
-        return Err(Rejection::UnrequestedUpgrade);
-    };
     if !rewrite::is_websocket_handshake(response.headers()) {
         return Err(Rejection::UnrequestedUpgrade);
     }

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,

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")]

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]

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]),