shift_proxy/
forward.rs

1//! Forward requests to upstream provider APIs and stream responses back.
2//!
3//! Handles header forwarding (auth passthrough), hop-by-hop header stripping
4//! (RFC 9110 §7.6.1), and transparent SSE/chunked response streaming.
5
6use axum::body::Body;
7use axum::http::{HeaderMap, HeaderValue, StatusCode};
8use axum::response::{IntoResponse, Response};
9use reqwest::Client;
10
11/// Headers stripped from upstream responses before forwarding to the client.
12///
13/// - `content-encoding` / `content-length`: reqwest auto-decompresses response
14///   bodies, so these are stale. Forwarding them causes double-decompression.
15///   NOTE: The `gzip`, `brotli`, and `deflate` features MUST be enabled on reqwest
16///   for this stripping to be correct. Without them, reqwest does NOT decompress,
17///   and stripping content-encoding causes clients to receive raw compressed bytes.
18/// - Hop-by-hop headers per RFC 9110 §7.6.1.
19const STRIP_RESPONSE_HEADERS: &[&str] = &[
20    "content-encoding",
21    "content-length",
22    "transfer-encoding",
23    "connection",
24    "keep-alive",
25    "proxy-authenticate",
26    "proxy-authorization",
27    "te",
28    "trailer",
29    "upgrade",
30];
31
32/// Headers stripped from the forwarded request.
33///
34/// - `host` / `content-length`: stale for the upstream connection.
35/// - `accept-encoding`: let reqwest negotiate compression based on its enabled
36///   decompression features (`gzip`, `brotli`, `deflate`). Forwarding the client's
37///   header could request encodings reqwest can't decompress (e.g., `zstd`), which
38///   would result in raw compressed bytes reaching the client after we strip
39///   `content-encoding`.
40const STRIP_REQUEST_HEADERS: &[&str] = &["host", "content-length", "accept-encoding"];
41
42/// Forward a request to an upstream URL, streaming the response back.
43///
44/// Auth headers (`authorization`, `x-api-key`, `anthropic-version`, `x-goog-api-key`)
45/// pass through unchanged. The response body is streamed directly — SSE and
46/// chunked responses are not buffered.
47pub async fn forward_request(
48    client: &Client,
49    method: &str,
50    target_url: &str,
51    request_headers: &HeaderMap,
52    body: Option<String>,
53) -> Response {
54    let forwarded_headers = forward_headers(request_headers);
55
56    let mut req = match method.to_uppercase().as_str() {
57        "POST" => client.post(target_url),
58        "GET" => client.get(target_url),
59        "PUT" => client.put(target_url),
60        "DELETE" => client.delete(target_url),
61        "PATCH" => client.patch(target_url),
62        _ => client.post(target_url),
63    };
64
65    req = req.headers(forwarded_headers);
66
67    if let Some(body) = body {
68        req = req.body(body);
69    }
70
71    match req.send().await {
72        Ok(upstream) => stream_response(upstream),
73        Err(err) => {
74            tracing::error!("upstream error: {}", err);
75            (
76                StatusCode::BAD_GATEWAY,
77                axum::Json(serde_json::json!({
78                    "error": "Bad Gateway",
79                    "detail": "Upstream provider unreachable"
80                })),
81            )
82                .into_response()
83        }
84    }
85}
86
87/// Convert a reqwest Response into an axum Response, streaming the body
88/// and stripping hop-by-hop headers.
89fn stream_response(upstream: reqwest::Response) -> Response {
90    let status = StatusCode::from_u16(upstream.status().as_u16()).unwrap_or(StatusCode::OK);
91
92    let mut response_headers = HeaderMap::new();
93    for (name, value) in upstream.headers() {
94        let name_str = name.as_str().to_lowercase();
95        if STRIP_RESPONSE_HEADERS
96            .iter()
97            .any(|h| h == &name_str.as_str())
98        {
99            continue;
100        }
101        if let Ok(v) = HeaderValue::from_bytes(value.as_bytes()) {
102            response_headers.insert(name.clone(), v);
103        }
104    }
105
106    // Stream the response body directly without buffering.
107    // This is critical for SSE (Anthropic/OpenAI streaming) to work correctly.
108    let body = Body::from_stream(upstream.bytes_stream());
109
110    let mut response = Response::new(body);
111    *response.status_mut() = status;
112    *response.headers_mut() = response_headers;
113    response
114}
115
116/// Forward request headers, stripping host/content-length but passing auth through.
117fn forward_headers(original: &HeaderMap) -> HeaderMap {
118    let strip: std::collections::HashSet<&str> = STRIP_REQUEST_HEADERS.iter().copied().collect();
119
120    let mut result = HeaderMap::new();
121    for (name, value) in original {
122        let name_lower = name.as_str().to_lowercase();
123        if !strip.contains(name_lower.as_str()) {
124            result.insert(name.clone(), value.clone());
125        }
126    }
127    result
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use axum::http::header;
134
135    #[test]
136    fn forward_headers_strips_host_content_length_and_accept_encoding() {
137        let mut headers = HeaderMap::new();
138        headers.insert(header::HOST, "example.com".parse().unwrap());
139        headers.insert(header::CONTENT_LENGTH, "42".parse().unwrap());
140        headers.insert(header::ACCEPT_ENCODING, "gzip, br".parse().unwrap());
141        headers.insert(header::AUTHORIZATION, "Bearer sk-test".parse().unwrap());
142        headers.insert("x-api-key", "sk-ant-test".parse().unwrap());
143        headers.insert("anthropic-version", "2023-06-01".parse().unwrap());
144
145        let result = forward_headers(&headers);
146
147        assert!(result.get(header::HOST).is_none());
148        assert!(result.get(header::CONTENT_LENGTH).is_none());
149        assert!(
150            result.get(header::ACCEPT_ENCODING).is_none(),
151            "accept-encoding should be stripped so reqwest negotiates its own"
152        );
153        assert_eq!(result.get(header::AUTHORIZATION).unwrap(), "Bearer sk-test");
154        assert_eq!(result.get("x-api-key").unwrap(), "sk-ant-test");
155        assert_eq!(result.get("anthropic-version").unwrap(), "2023-06-01");
156    }
157}
shift_proxy/forward.rs

shift_proxy/
forward.rs
