clawft_kernel/infer_proxy/
upstream.rs

1//! Forwarding to a loopback inference server.
2//!
3//! The destination is never taken from the request: it is a base URL the
4//! placement table verified as `http://<loopback ip>:<port>`, joined with
5//! an allowlisted origin-form path. The client keeps no proxy, follows no
6//! redirect and reuses no connection.
7
8use std::net::IpAddr;
9use std::time::Duration;
10
11use super::types::{Method, ProxyError, ProxyLimits, ProxyRequest, ResponseSink};
12
13/// `http://<loopback ip>:<port>` and nothing else (no path, userinfo or
14/// name: `localhost` could be remapped, an address cannot).
15pub fn parse_loopback_base(base: &str) -> Result<(IpAddr, u16), ProxyError> {
16    let bad = || ProxyError::NotLoopback(base.to_string());
17    let rest = base.strip_prefix("http://").ok_or_else(bad)?;
18    let rest = rest.trim_end_matches('/');
19    let sock: std::net::SocketAddr = rest.parse().map_err(|_| bad())?;
20    if !sock.ip().is_loopback() || sock.port() == 0 {
21        return Err(bad());
22    }
23    Ok((sock.ip(), sock.port()))
24}
25
26/// Streaming HTTP client for local servers.
27pub struct Upstream {
28    http: reqwest::Client,
29    limits: ProxyLimits,
30}
31
32impl Upstream {
33    /// Client with the limits' connect timeout.
34    pub fn new(limits: ProxyLimits) -> Result<Self, ProxyError> {
35        let http = reqwest::Client::builder()
36            .redirect(reqwest::redirect::Policy::none())
37            .no_proxy()
38            .pool_max_idle_per_host(0)
39            .connect_timeout(limits.connect_timeout)
40            .build()
41            .map_err(|e| ProxyError::Io(format!("http client: {e}")))?;
42        Ok(Self { http, limits })
43    }
44
45    /// Limits in force.
46    pub fn limits(&self) -> &ProxyLimits {
47        &self.limits
48    }
49
50    /// Forward `req` to the verified loopback `base`, relaying the response
51    /// into `sink` as it arrives. `send_auth` is true only for a request from
52    /// a local client on this node's loopback listener: the client chose to
53    /// send that credential to this address, and a server started with
54    /// `--api-key` needs it. It is forwarded to the local instance only and
55    /// never over the mesh; it is false for every request that arrived over
56    /// the mesh (a peer's credentials are not ours to pass on).
57    pub async fn forward(
58        &self,
59        base: &str,
60        req: &ProxyRequest,
61        send_auth: bool,
62        sink: &mut dyn ResponseSink,
63    ) -> Result<(), ProxyError> {
64        let (ip, port) = parse_loopback_base(base)?;
65        let authority = match ip {
66            IpAddr::V4(a) => format!("{a}:{port}"),
67            IpAddr::V6(a) => format!("[{a}]:{port}"),
68        };
69        let url = format!("http://{authority}{}", req.path);
70        let parsed = reqwest::Url::parse(&url).map_err(|e| ProxyError::BadRequest(e.to_string()))?;
71        // Defense in depth: whatever the path was, the host is unchanged.
72        if parsed.port() != Some(port) || parsed.host_str().is_none() {
73            return Err(ProxyError::Forbidden("destination changed".into()));
74        }
75        let mut b = match req.method {
76            Method::Get => self.http.get(parsed),
77            Method::Post => self.http.post(parsed).body(req.body.clone()),
78        };
79        if let Some(ct) = &req.content_type {
80            b = b.header("content-type", ct);
81        }
82        if let Some(a) = &req.accept {
83            b = b.header("accept", a);
84        }
85        if send_auth && let Some(a) = &req.authorization {
86            b = b.header("authorization", a);
87        }
88        let stall = self.limits.stall_timeout;
89        let mut resp = tokio::time::timeout(stall, b.send())
90            .await
91            .map_err(|_| ProxyError::Timeout("upstream did not answer".into()))?
92            .map_err(|e| ProxyError::Upstream(e.without_url().to_string()))?;
93        let ct = resp
94            .headers()
95            .get("content-type")
96            .and_then(|v| v.to_str().ok())
97            .map(str::to_string);
98        let code = resp.status().as_u16();
99        if !(200..=599).contains(&code) {
100            // 1xx (or 101 Upgrade) is not a final answer and must not reach
101            // the client as one.
102            return Err(ProxyError::Upstream(format!("unexpected status {code}")));
103        }
104        sink.head(code, ct.as_deref()).await?;
105        let mut total: u64 = 0;
106        loop {
107            let next = tokio::time::timeout(stall, resp.chunk())
108                .await
109                .map_err(|_| ProxyError::Timeout("upstream stalled".into()))?
110                .map_err(|e| ProxyError::Upstream(e.without_url().to_string()))?;
111            let Some(c) = next else { break };
112            total += c.len() as u64;
113            if total > self.limits.max_response_body {
114                return Err(ProxyError::TooLarge("response body".into()));
115            }
116            sink.chunk(&c).await?;
117        }
118        Ok(())
119    }
120}
121
122/// Connect timeout used when probing whether a port is held.
123pub const PROBE_TIMEOUT: Duration = Duration::from_millis(300);
clawft_kernel/infer_proxy/upstream.rs

clawft_kernel/infer_proxy/
upstream.rs
