diff --git a/docs/web-ui/agent.md b/docs/web-ui/agent.md index 50ef0adf..63f5b708 100644 --- a/docs/web-ui/agent.md +++ b/docs/web-ui/agent.md @@ -380,7 +380,8 @@ shaped). `{"": ""}`). Each entry mounts a transparent reverse-proxy at `/extra//` that forwards every request (method, headers, body) to the configured upstream, strips hop-by-hop headers - on both sides, and buffers the full response body (MVP — SSE + on both sides (repeated response headers such as two `Set-Cookie` lines + are all forwarded), and buffers the full response body (MVP — SSE connections will appear as one large response rather than streaming). The `/extra/` namespace ensures user-declared proxies can never conflict with native agent endpoints. Upstream values are either an diff --git a/hive-agent/src/web_ui/proxy.rs b/hive-agent/src/web_ui/proxy.rs index 1bdb5e45..c44c207f 100644 --- a/hive-agent/src/web_ui/proxy.rs +++ b/hive-agent/src/web_ui/proxy.rs @@ -169,7 +169,7 @@ async fn proxy_handler( HeaderName::from_bytes(name.as_str().as_bytes()), HeaderValue::from_bytes(value.as_bytes()), ) { - resp_headers.insert(n, v); + resp_headers.append(n, v); } } match upstream_resp.bytes().await { @@ -269,8 +269,75 @@ async fn forward_over_unix_socket( if HOP_BY_HOP.contains(&name.as_str()) { continue; } - resp_headers.insert(name.clone(), value.clone()); + resp_headers.append(name.clone(), value.clone()); } let body_bytes = upstream_resp.into_body().collect().await?.to_bytes(); Ok((status, resp_headers, body_bytes).into_response()) } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + /// Raw HTTP/1.1 answer with two `Set-Cookie` headers and a body. + const TWO_COOKIES: &[u8] = b"HTTP/1.1 200 OK\r\n\ + Set-Cookie: session=abc\r\n\ + Set-Cookie: csrf=def\r\n\ + Content-Length: 2\r\n\ + Connection: close\r\n\r\nok"; + + /// Read the request head, then answer with [`TWO_COOKIES`]. + async fn answer_once(mut stream: S) { + let mut buf = [0u8; 1024]; + let _ = stream.read(&mut buf).await; + stream.write_all(TWO_COOKIES).await.unwrap(); + stream.shutdown().await.unwrap(); + } + + fn cookies(resp: &Response) -> Vec<&str> { + resp.headers() + .get_all("set-cookie") + .iter() + .map(|v| v.to_str().unwrap()) + .collect() + } + + #[tokio::test] + async fn unix_upstream_keeps_repeated_response_headers() { + let dir = tempfile::tempdir().unwrap(); + let sock = dir.path().join("up.sock"); + let listener = tokio::net::UnixListener::bind(&sock).unwrap(); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + answer_once(stream).await; + }); + + let resp = + forward_over_unix_socket(&sock, Method::GET, "/", HeaderMap::new(), Bytes::new()) + .await + .unwrap(); + assert_eq!(cookies(&resp), ["session=abc", "csrf=def"]); + } + + #[tokio::test] + async fn http_upstream_keeps_repeated_response_headers() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base = format!("http://{}", listener.local_addr().unwrap()); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + answer_once(stream).await; + }); + + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let resp = proxy_handler( + State((Arc::new(client), Arc::new(base))), + Method::GET, + Uri::from_static("/"), + HeaderMap::new(), + Bytes::new(), + ) + .await; + assert_eq!(cookies(&resp), ["session=abc", "csrf=def"]); + } +}