shunt/
forwarder.rs

1use anyhow::{Context, Result};
2use axum::body::Body;
3use axum::http::{HeaderMap, HeaderName, HeaderValue, Response};
4use bytes::Bytes;
5use reqwest::Client;
6use std::str::FromStr;
7use std::time::Instant;
8use tracing::info;
9use uuid::Uuid;
10
11use crate::config::AccountConfig;
12
13/// Headers that must never be forwarded in either direction.
14const HOP_BY_HOP: &[&str] = &[
15    "connection",
16    "keep-alive",
17    "proxy-authenticate",
18    "proxy-authorization",
19    "te",
20    "trailers",
21    "transfer-encoding",
22    "upgrade",
23    "host",
24    "content-length",
25];
26
27/// Auth headers that the proxy manages — always stripped from client requests
28/// and replaced with the selected account's credential.
29const CLIENT_AUTH_HEADERS: &[&str] = &["authorization", "x-api-key"];
30
31fn is_hop_by_hop(name: &str) -> bool {
32    HOP_BY_HOP.contains(&name.to_ascii_lowercase().as_str())
33}
34
35fn is_client_auth(name: &str) -> bool {
36    CLIENT_AUTH_HEADERS.contains(&name.to_ascii_lowercase().as_str())
37}
38
39pub struct Forwarder {
40    client: Client,
41    base_url: String,
42}
43
44impl Forwarder {
45    pub fn new(base_url: impl Into<String>) -> Result<Self> {
46        let client = Client::builder()
47            .timeout(std::time::Duration::from_secs(600))
48            .redirect(reqwest::redirect::Policy::none())
49            .build()
50            .context("Failed to build HTTP client")?;
51
52        Ok(Self { client, base_url: base_url.into() })
53    }
54
55    /// Forward a request to the upstream using the given account's OAuth credential.
56    ///
57    /// - Strips `Authorization` and `x-api-key` from the client request.
58    /// - Injects `Authorization: Bearer <token>` (live token, may differ from account.credential).
59    /// - Keeps the upstream TCP connection alive for streaming responses.
60    pub async fn forward(
61        &self,
62        method: &str,
63        path: &str,
64        body: Bytes,
65        client_headers: &HeaderMap,
66        account: &AccountConfig,
67        token: &str,
68    ) -> Result<Response<Body>> {
69        let request_id = &Uuid::new_v4().to_string()[..8];
70        let url = format!("{}{}", self.base_url, path);
71
72        let mut upstream_headers = reqwest::header::HeaderMap::new();
73
74        for (name, value) in client_headers.iter() {
75            let lower = name.as_str().to_ascii_lowercase();
76            if is_hop_by_hop(&lower) || is_client_auth(&lower) {
77                continue;
78            }
79            if let (Ok(n), Ok(v)) = (
80                reqwest::header::HeaderName::from_str(name.as_str()),
81                reqwest::header::HeaderValue::from_bytes(value.as_bytes()),
82            ) {
83                upstream_headers.insert(n, v);
84            }
85        }
86
87        // Inject the live OAuth Bearer token
88        upstream_headers.insert(
89            reqwest::header::HeaderName::from_static("authorization"),
90            reqwest::header::HeaderValue::from_str(&format!("Bearer {token}"))
91                .context("invalid access token value")?,
92        );
93
94        // Required by Anthropic when authenticating with OAuth tokens
95        upstream_headers.insert(
96            reqwest::header::HeaderName::from_static("anthropic-dangerous-direct-browser-access"),
97            reqwest::header::HeaderValue::from_static("true"),
98        );
99
100        // Ensure oauth-2025-04-20 is present in anthropic-beta (required for OAuth tokens).
101        // Merge with any beta flags the client already sent.
102        let beta_key = reqwest::header::HeaderName::from_static("anthropic-beta");
103        let existing_beta = upstream_headers
104            .get(&beta_key)
105            .and_then(|v| v.to_str().ok())
106            .unwrap_or("")
107            .to_owned();
108        let merged_beta = if existing_beta.split(',').any(|s| s.trim() == "oauth-2025-04-20") {
109            existing_beta
110        } else if existing_beta.is_empty() {
111            "oauth-2025-04-20".to_owned()
112        } else {
113            format!("{existing_beta},oauth-2025-04-20")
114        };
115        upstream_headers.insert(
116            beta_key,
117            reqwest::header::HeaderValue::from_str(&merged_beta).unwrap(),
118        );
119
120        let t0 = Instant::now();
121        let upstream_resp = self
122            .client
123            .request(
124                reqwest::Method::from_str(method).context("invalid method")?,
125                &url,
126            )
127            .headers(upstream_headers)
128            .body(body)
129            .send()
130            .await
131            .context("upstream request failed")?;
132
133        let latency_ms = t0.elapsed().as_millis();
134        let status = upstream_resp.status();
135
136        info!(
137            request_id = %request_id,
138            account = %account.name,
139            status = status.as_u16(),
140            latency_ms = %latency_ms,
141            path = %path,
142            "request forwarded"
143        );
144
145        let mut builder = Response::builder().status(status.as_u16());
146
147        for (name, value) in upstream_resp.headers().iter() {
148            if !is_hop_by_hop(name.as_str()) {
149                if let (Ok(n), Ok(v)) = (
150                    HeaderName::from_str(name.as_str()),
151                    HeaderValue::from_bytes(value.as_bytes()),
152                ) {
153                    builder = builder.header(n, v);
154                }
155            }
156        }
157
158        let body = Body::from_stream(upstream_resp.bytes_stream());
159        Ok(builder.body(body).expect("response builder invariant"))
160    }
161}
shunt/forwarder.rs

shunt/
forwarder.rs
