Skip to main content

rust_web_server/proxy/
mod.rs

1//! Reverse proxy middleware with round-robin load balancing.
2//!
3//! `ReverseProxy` implements [`Middleware`] — wrap any application with it and
4//! all matching requests are forwarded to one of the configured backends over
5//! plain HTTP/1.1.  Failed backends are skipped and the next one is tried
6//! before returning `502 Bad Gateway`.
7//!
8//! # Example
9//!
10//! ```rust,no_run
11//! use rust_web_server::app::App;
12//! use rust_web_server::core::New;
13//! use rust_web_server::proxy::{LoadBalancing, ReverseProxy};
14//!
15//! // Proxy every request across two backends in round-robin order.
16//! let app = App::new()
17//!     .wrap(ReverseProxy::new(["http://backend-1:8080", "http://backend-2:8080"])
18//!         .strategy(LoadBalancing::RoundRobin));
19//!
20//! // Only proxy /api/* requests; everything else is handled locally.
21//! let app2 = App::new()
22//!     .wrap(ReverseProxy::new(["http://api-service:3000"])
23//!         .path_prefix("/api"));
24//! ```
25
26pub mod pool;
27#[cfg(feature = "http2")]
28pub(crate) mod async_upstream;
29
30#[cfg(test)]
31mod tests;
32
33use std::io::{Read, Write};
34use crate::ext::decimal::Decimal;
35use std::net::{TcpStream, ToSocketAddrs};
36use std::sync::atomic::{AtomicUsize, Ordering};
37use std::sync::Arc;
38use std::time::Duration;
39
40pub use pool::{ConnPool, Pool};
41
42use crate::application::Application;
43use crate::core::New;
44use crate::middleware::Middleware;
45use crate::mime_type::MimeType;
46use crate::range::Range;
47use crate::request::Request;
48use crate::response::{Response, STATUS_CODE_REASON_PHRASE};
49use crate::server::ConnectionInfo;
50
51// Hop-by-hop headers that must not be forwarded (RFC 9110 §7.6.1), plus
52// `proxy-connection` (legacy clients). Headers named in `Connection` are
53// dropped too (`connection_listed`).
54const HOP_BY_HOP: &[&str] = &[
55    "connection",
56    "keep-alive",
57    "proxy-authenticate",
58    "proxy-authorization",
59    "proxy-connection",
60    "te",
61    "trailer",
62    "transfer-encoding",
63    "upgrade",
64];
65
66/// Load balancing strategy used by [`ReverseProxy`].
67pub enum LoadBalancing {
68    /// Distribute requests across backends in a cyclic order.
69    RoundRobin,
70}
71
72/// Reverse proxy middleware.
73///
74/// Forwards incoming requests to one of the configured backends over HTTP/1.1.
75/// On connection failure the next backend in the list is tried; when all
76/// backends have failed the middleware returns `502 Bad Gateway`.
77///
78/// Hop-by-hop headers are stripped before forwarding.  `X-Forwarded-For` and
79/// `Via` are added to every forwarded request.
80///
81/// Idle connections are pooled and reused across requests (up to
82/// 64 idle per backend, 60-second timeout).
83/// This eliminates per-request TCP handshake overhead and ephemeral-port
84/// exhaustion.  Use [`ReverseProxy::with_pool`] to share a pool across
85/// multiple proxy instances or to tune pool parameters.
86pub struct ReverseProxy {
87    backends: Vec<Backend>,
88    path_prefix: Option<String>,
89    connect_timeout: Duration,
90    read_timeout: Duration,
91    counter: AtomicUsize,
92    pool: Arc<ConnPool>,
93    breaker: Option<Arc<dyn crate::circuit_breaker::Breaker>>,
94    trust: Option<Arc<UpstreamTls>>,
95}
96
97impl ReverseProxy {
98    /// Create a proxy that distributes requests across `backends` in
99    /// round-robin order.  Each entry must be `"http://host:port"` or
100    /// `"host:port"` (port defaults to 80).
101    pub fn new<I, S>(backends: I) -> Self
102    where
103        I: IntoIterator<Item = S>,
104        S: AsRef<str>,
105    {
106        Self {
107            backends: backends
108                .into_iter()
109                .filter_map(|u| Backend::parse(u.as_ref()))
110                .collect(),
111            path_prefix: None,
112            connect_timeout: Duration::from_secs(5),
113            read_timeout: Duration::from_secs(30),
114            counter: AtomicUsize::new(0),
115            pool: Arc::new(ConnPool::new(POOL_MAX_IDLE, Duration::from_secs(60))),
116            breaker: None,
117            trust: None,
118        }
119    }
120
121    /// Wire a [`crate::circuit_breaker::CircuitBreaker`] or
122    /// [`crate::circuit_breaker::RedisCircuitBreaker`] into this proxy:
123    /// before dialing a backend, its breaker state is checked (an `Open`
124    /// backend is skipped with no TCP attempt, exactly like a caller
125    /// manually checking `is_available` around each proxied call); after a
126    /// dial, `record_success`/`record_failure` is called automatically.
127    /// No breaker is wired by default — existing `ReverseProxy` callers see
128    /// no behavior change unless they opt in here.
129    ///
130    /// ```rust,no_run
131    /// use std::sync::Arc;
132    /// use rust_web_server::circuit_breaker::global as global_breaker;
133    /// use rust_web_server::proxy::ReverseProxy;
134    ///
135    /// let proxy = ReverseProxy::new(["http://backend-1:8080", "http://backend-2:8080"])
136    ///     .with_circuit_breaker(Arc::new(global_breaker()));
137    /// ```
138    pub fn with_circuit_breaker(mut self, breaker: Arc<dyn crate::circuit_breaker::Breaker>) -> Self {
139        self.breaker = Some(breaker);
140        self
141    }
142
143    /// Only proxy requests whose URI starts with `prefix`.
144    ///
145    /// Other requests are passed through to the next layer in the middleware
146    /// chain (or the inner application).
147    pub fn path_prefix(mut self, prefix: impl Into<String>) -> Self {
148        self.path_prefix = Some(prefix.into());
149        self
150    }
151
152    /// Override the load balancing strategy (currently only `RoundRobin`).
153    pub fn strategy(self, _strategy: LoadBalancing) -> Self {
154        self
155    }
156
157    /// Override the TCP connect timeout (default: 5 000 ms).
158    pub fn connect_timeout_ms(mut self, ms: u64) -> Self {
159        self.connect_timeout = Duration::from_millis(ms);
160        self
161    }
162
163    /// Override the response read timeout (default: 30 000 ms).
164    pub fn read_timeout_ms(mut self, ms: u64) -> Self {
165        self.read_timeout = Duration::from_millis(ms);
166        self
167    }
168
169    /// Attach a shared connection pool.
170    ///
171    /// Useful for sharing one pool across multiple `ReverseProxy` instances
172    /// or for tuning pool parameters (capacity, idle timeout).
173    pub fn with_pool(mut self, pool: Arc<ConnPool>) -> Self {
174        self.pool = pool;
175        self
176    }
177
178    /// Trust the CA certificates in the PEM file `path` for `https://`
179    /// backends, **instead of** the public WebPKI roots — for internal
180    /// services behind a private CA or a self-signed certificate.
181    pub fn tls_ca_file(mut self, path: &str) -> Result<Self, String> {
182        self.trust = Some(Arc::new(UpstreamTls::from_ca_file(path)?));
183        Ok(self)
184    }
185
186    /// Set the maximum number of idle connections per backend (default: 64).
187    pub fn max_idle_conns(mut self, n: usize) -> Self {
188        self.pool = Arc::new(ConnPool::new(n, Duration::from_secs(60)));
189        self
190    }
191
192    fn proxy(&self, request: &Request, connection: &ConnectionInfo) -> Result<Response, String> {
193        if self.backends.is_empty() {
194            return Err("no backends configured".to_string());
195        }
196        let n = self.backends.len();
197        let start = self.counter.fetch_add(1, Ordering::Relaxed);
198        let mut last_err = "all backends failed".to_string();
199        for attempt in 0..n {
200            let idx = (start + attempt) % n;
201            let backend = &self.backends[idx];
202            let key = format!("{}:{}", backend.host, backend.port);
203
204            if let Some(breaker) = &self.breaker {
205                if !breaker.is_available(&key) {
206                    last_err = format!("circuit open for {}", key);
207                    continue;
208                }
209            }
210
211            match self.try_backend(request, connection, backend) {
212                Ok(resp) => {
213                    if let Some(breaker) = &self.breaker {
214                        breaker.record_success(&key);
215                    }
216                    return Ok(resp);
217                }
218                Err(e) => {
219                    if let Some(breaker) = &self.breaker {
220                        breaker.record_failure(&key);
221                    }
222                    last_err = e;
223                }
224            }
225        }
226        Err(last_err)
227    }
228
229    fn try_backend(
230        &self,
231        request: &Request,
232        connection: &ConnectionInfo,
233        backend: &Backend,
234    ) -> Result<Response, String> {
235        forward_upstream_trusting(
236            &self.pool,
237            request,
238            &connection.client.ip,
239            &backend.host,
240            backend.port,
241            backend.tls,
242            self.trust.as_deref(),
243            self.connect_timeout,
244            self.read_timeout,
245        )
246    }
247}
248
249impl Middleware for ReverseProxy {
250    fn handle(
251        &self,
252        request: &Request,
253        connection: &ConnectionInfo,
254        next: &dyn Application,
255    ) -> Result<Response, String> {
256        if let Some(prefix) = &self.path_prefix {
257            if !request.request_uri.starts_with(prefix.as_str()) {
258                return next.execute(request, connection);
259            }
260        }
261        match self.proxy(request, connection) {
262            Ok(resp) => Ok(resp),
263            Err(_) => Ok(bad_gateway()),
264        }
265    }
266}
267
268// ── helpers ───────────────────────────────────────────────────────────────────
269
270pub(crate) fn build_request(
271    request: &Request,
272    backend_host: &str,
273    client_ip: &str,
274    keep_alive: bool,
275) -> Vec<u8> {
276    // One buffer, sized up front and filled with extend_from_slice — no
277    // write!/format machinery, no lower-cased copy of each header name.
278    // `Content-Length` is regenerated from the (possibly de-chunked) body and
279    // `X-Forwarded-For` is extended below, so the client's copies are dropped:
280    // two `Content-Length` lines are a protocol error and a desync risk.
281    let connection = request.header(crate::header::Header::_CONNECTION).map(|h| h.value.as_ref());
282    let forwarded = |h: &&crate::header::Header| {
283        // An empty name would be sent as a malformed ": value" line.
284        !h.name.is_empty()
285            && !is_hop_by_hop_or_host(&h.name)
286            && !h.name.eq_ignore_ascii_case(crate::header::Header::_CONTENT_LENGTH)
287            && !h.name.eq_ignore_ascii_case("X-Forwarded-For")
288            && !connection.is_some_and(|c| connection_listed(c, &h.name))
289    };
290    let prior_xff = request.header("X-Forwarded-For").map(|h| h.value.as_ref()).filter(|v| !v.trim().is_empty());
291    let body_len = Decimal::new(request.body.len() as u128);
292    let capacity = request.method.len() + request.request_uri.len() + backend_host.len() + 25
293        + request.headers.iter().filter(forwarded).map(|h| h.name.len() + h.value.len() + 4).sum::<usize>()
294        + client_ip.len() + 19 + prior_xff.map_or(0, |v| v.len() + 2)
295        + 13
296        + 24
297        + if request.body.is_empty() { 0 } else { body_len.as_str().len() + 18 }
298        + 2
299        + request.body.len();
300    let mut out: Vec<u8> = Vec::with_capacity(capacity);
301    for part in [request.method.as_str(), " ", request.request_uri.as_str(), " HTTP/1.1\r\nHost: ", backend_host, "\r\n"] {
302        out.extend_from_slice(part.as_bytes());
303    }
304    for h in request.headers.iter().filter(forwarded) {
305        for part in [h.name.as_ref(), ": ", h.value.as_ref(), "\r\n"] {
306            out.extend_from_slice(part.as_bytes());
307        }
308    }
309    out.extend_from_slice(b"X-Forwarded-For: ");
310    if let Some(prior) = prior_xff {
311        out.extend_from_slice(prior.as_bytes());
312        out.extend_from_slice(b", ");
313    }
314    for part in [client_ip, "\r\n", "Via: 1.1 rws\r\n"] {
315        out.extend_from_slice(part.as_bytes());
316    }
317    if keep_alive {
318        out.extend_from_slice(b"Connection: keep-alive\r\n");
319    } else {
320        out.extend_from_slice(b"Connection: close\r\n");
321    }
322    if !request.body.is_empty() {
323        out.extend_from_slice(b"Content-Length: ");
324        out.extend_from_slice(body_len.as_str().as_bytes());
325        out.extend_from_slice(b"\r\n");
326    }
327    out.extend_from_slice(b"\r\n");
328    out.extend_from_slice(&request.body);
329    out
330}
331
332/// Whether a request header is dropped on the way upstream: hop-by-hop
333/// headers and `Host` (rws sends its own). Matches `to_lowercase()` against
334/// the lower-case list — Unicode lower-casing, so a non-ASCII name (where
335/// e.g. the Kelvin sign lower-cases to `k`) still takes that path.
336fn is_hop_by_hop_or_host(name: &str) -> bool {
337    if name.is_ascii() {
338        name.eq_ignore_ascii_case("host") || HOP_BY_HOP.iter().any(|h| name.eq_ignore_ascii_case(h))
339    } else {
340        let lower = name.to_lowercase();
341        lower == "host" || HOP_BY_HOP.contains(&lower.as_str())
342    }
343}
344
345/// Whether `name` is one of the comma-separated tokens of a `Connection`
346/// header value — those headers are hop-by-hop for this connection only.
347fn connection_listed(connection: &str, name: &str) -> bool {
348    connection.split(',').any(|token| token.trim().eq_ignore_ascii_case(name))
349}
350
351/// Decode HTTP/1.1 chunked transfer-encoding from `stream`.
352///
353/// `buf[header_end..]` may already contain some body bytes that arrived in
354/// the same read as the headers.  Returns the fully decoded body.
355fn decode_chunked<S: Read>(
356    stream: &mut S,
357    buf: &[u8],
358    header_end: usize,
359    tmp: &mut [u8],
360) -> Result<Vec<u8>, String> {
361    // Seed `raw` with any body bytes already buffered alongside the headers.
362    let mut raw: Vec<u8> = buf[header_end..].to_vec();
363    let mut decoded: Vec<u8> = Vec::new();
364
365    loop {
366        // Wait until we have at least one complete chunk-size line (ends with \r\n).
367        let crlf = loop {
368            if let Some(p) = raw.windows(2).position(|w| w == b"\r\n") {
369                break p;
370            }
371            let n = stream.read(tmp).map_err(|e| e.to_string())?;
372            if n == 0 {
373                return Err("chunked: premature EOF reading chunk size".to_string());
374            }
375            raw.extend_from_slice(&tmp[..n]);
376        };
377
378        // Chunk size is hex, optionally followed by chunk-extensions (";…").
379        let size_line = std::str::from_utf8(&raw[..crlf])
380            .map_err(|_| "chunked: non-UTF-8 chunk size line".to_string())?;
381        let size_str = size_line.split(';').next().unwrap_or("").trim();
382        let chunk_size = usize::from_str_radix(size_str, 16)
383            .map_err(|_| format!("chunked: invalid chunk size '{}'", size_str))?;
384        raw.drain(..crlf + 2); // consume "<size>\r\n"
385
386        if chunk_size == 0 {
387            // Last chunk: skip trailer fields up to the empty line that ends
388            // the message. Leaving any of it unread would corrupt the next
389            // response on a pooled connection.
390            loop {
391                let line_end = loop {
392                    if let Some(p) = raw.windows(2).position(|w| w == b"\r\n") {
393                        break p;
394                    }
395                    let n = stream.read(tmp).map_err(|e| e.to_string())?;
396                    if n == 0 {
397                        return Err("chunked: premature EOF in trailer".to_string());
398                    }
399                    raw.extend_from_slice(&tmp[..n]);
400                };
401                raw.drain(..line_end + 2);
402                if line_end == 0 {
403                    break;
404                }
405            }
406            break;
407        }
408
409        // Read chunk data + trailing CRLF.
410        while raw.len() < chunk_size + 2 {
411            let n = stream.read(tmp).map_err(|e| e.to_string())?;
412            if n == 0 {
413                return Err("chunked: premature EOF reading chunk body".to_string());
414            }
415            raw.extend_from_slice(&tmp[..n]);
416        }
417        decoded.extend_from_slice(&raw[..chunk_size]);
418        raw.drain(..chunk_size + 2); // consume "<data>\r\n"
419    }
420
421    Ok(decoded)
422}
423
424/// Rewrite `buf` in-place: strip `Transfer-Encoding`, add `Content-Length`,
425/// replace the old (undecoded) body with `decoded`.
426fn rewrite_as_content_length(buf: &mut Vec<u8>, header_end: usize, decoded: &[u8]) {
427    let header_str = std::str::from_utf8(&buf[..header_end]).unwrap_or("").to_string();
428    buf.clear();
429    for line in header_str.lines() {
430        if line.to_ascii_lowercase().starts_with("transfer-encoding:") || line.is_empty() {
431            continue;
432        }
433        buf.extend_from_slice(line.as_bytes());
434        buf.extend_from_slice(b"\r\n");
435    }
436    let _ = write!(buf, "Content-Length: {}\r\n\r\n", decoded.len());
437    buf.extend_from_slice(decoded);
438}
439
440/// Headers rws regenerates for a buffered body, plus hop-by-hop headers
441/// (RFC 9110 §7.6.1) that must not be forwarded.
442pub(crate) const REGENERATED_OR_HOP_BY_HOP: &[&str] = &[
443    "content-length", "transfer-encoding", "connection", "keep-alive",
444    "proxy-connection", "te", "trailer", "upgrade",
445];
446
447/// Make a buffered upstream response safe to re-serialize. The body is
448/// already held in `content_range_list`, from which `Response::generate`
449/// (and the HTTP/2 / HTTP/3 writers) emit `Content-Type` and
450/// `Content-Length` themselves, so the upstream's own copies would be
451/// duplicated — and a duplicate `content-length` is a protocol error that
452/// HTTP/2 clients reject outright. The upstream's `Content-Type` is moved
453/// onto the body instead; the status line is re-emitted as HTTP/1.1
454/// regardless of the backend's version.
455pub(crate) fn normalize_upstream_response(mut response: Response) -> Response {
456    response.http_version.clear();
457    response.http_version.push_str(crate::http::VERSION.http_1_1);
458    let has_body = !response.content_range_list.is_empty();
459    if let [range] = response.content_range_list.as_mut_slice() {
460        if let Some(h) = response
461            .headers
462            .iter()
463            .find(|h| h.name.eq_ignore_ascii_case(crate::header::Header::_CONTENT_TYPE))
464        {
465            range.content_type.clear();
466            range.content_type.push_str(&h.value);
467        }
468    }
469    // Compared in place (the list is lower-case ASCII), not lower-cased
470    // into a new String per header.
471    response.headers.retain(|h| {
472        if h.name.eq_ignore_ascii_case("content-type") {
473            // Keep it only when there's no body to carry it (e.g. HEAD, 304).
474            return !has_body;
475        }
476        !REGENERATED_OR_HOP_BY_HOP.iter().any(|name| h.name.eq_ignore_ascii_case(name))
477    });
478    response
479}
480
481fn bad_gateway() -> Response {
482    let cr = Range::get_content_range(
483        b"502 Bad Gateway".to_vec(),
484        MimeType::TEXT_PLAIN.to_string(),
485    );
486    let mut r = Response::new();
487    r.status_code = *STATUS_CODE_REASON_PHRASE.n502_bad_gateway.status_code;
488    r.reason_phrase = STATUS_CODE_REASON_PHRASE
489        .n502_bad_gateway
490        .reason_phrase
491        .to_string();
492    r.content_range_list = vec![cr];
493    r
494}
495
496// ── Streaming proxy helpers ───────────────────────────────────────────────────
497
498/// Responses larger than this threshold are streamed instead of buffered.
499const STREAM_THRESHOLD: usize = 1024 * 1024; // 1 MB
500
501/// A `Read` implementation that drains a prefix buffer before reading from the
502/// inner stream. Used to replay body bytes that arrived with the HTTP headers.
503pub(crate) struct ConcatReader<R: Read + Send> {
504    prefix: Vec<u8>,
505    prefix_pos: usize,
506    inner: R,
507}
508
509impl<R: Read + Send> ConcatReader<R> {
510    fn new(prefix: Vec<u8>, inner: R) -> Self {
511        ConcatReader { prefix, prefix_pos: 0, inner }
512    }
513}
514
515impl<R: Read + Send> Read for ConcatReader<R> {
516    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
517        if self.prefix_pos < self.prefix.len() {
518            let avail = &self.prefix[self.prefix_pos..];
519            let n = buf.len().min(avail.len());
520            buf[..n].copy_from_slice(&avail[..n]);
521            self.prefix_pos += n;
522            return Ok(n);
523        }
524        self.inner.read(buf)
525    }
526}
527
528
529/// Idle backend connections kept per `host:port`. 64 matches common nginx
530/// `keepalive` settings and the benchmark's concurrency; fewer meant most
531/// requests at 64-way concurrency still paid a TCP connect.
532pub(crate) const POOL_MAX_IDLE: usize = 64;
533
534/// The process-wide backend connection pool used by config-driven proxy
535/// routes (`proxy_config::DynamicProxy`), keyed by `host:port`.
536pub(crate) fn shared_pool() -> &'static ConnPool {
537    static POOL: std::sync::OnceLock<ConnPool> = std::sync::OnceLock::new();
538    POOL.get_or_init(|| ConnPool::new(POOL_MAX_IDLE, Duration::from_secs(60)))
539}
540
541/// A backend connection the pooled forwarder drives: plain TCP, or TLS
542/// over TCP for `https://` backends.
543pub(crate) trait UpstreamIo: Read + Write + Send + 'static {
544    fn set_io_timeouts(&self, read: Duration, write: Duration) -> std::io::Result<()>;
545}
546
547impl UpstreamIo for TcpStream {
548    fn set_io_timeouts(&self, read: Duration, write: Duration) -> std::io::Result<()> {
549        self.set_read_timeout(Some(read))?;
550        self.set_write_timeout(Some(write))
551    }
552}
553
554#[cfg(any(feature = "http-client", feature = "http2"))]
555pub(crate) type TlsStream = rustls::StreamOwned<rustls::ClientConnection, TcpStream>;
556
557#[cfg(any(feature = "http-client", feature = "http2"))]
558impl UpstreamIo for TlsStream {
559    fn set_io_timeouts(&self, read: Duration, write: Duration) -> std::io::Result<()> {
560        self.sock.set_io_timeouts(read, write)
561    }
562}
563
564/// Forward `request` to an HTTP or HTTPS backend at `host:port`, over a
565/// pooled keep-alive connection either way: plain backends use `plain_pool`,
566/// TLS backends the process-wide TLS pool (one handshake per connection, not
567/// per request). The single entry point for `ReverseProxy`, config-driven
568/// routes, `canary` and `ingress`.
569pub(crate) fn forward_upstream(
570    plain_pool: &ConnPool,
571    request: &Request,
572    client_ip: &str,
573    host: &str,
574    port: u16,
575    tls: bool,
576    connect_timeout: Duration,
577    read_timeout: Duration,
578) -> Result<Response, String> {
579    forward_upstream_trusting(plain_pool, request, client_ip, host, port, tls, None, connect_timeout, read_timeout)
580}
581
582/// [`forward_upstream`] with the trust settings of one upstream: `trust`
583/// replaces the WebPKI roots for TLS backends (a private CA).
584#[allow(clippy::too_many_arguments)]
585pub(crate) fn forward_upstream_trusting(
586    plain_pool: &ConnPool,
587    request: &Request,
588    client_ip: &str,
589    host: &str,
590    port: u16,
591    tls: bool,
592    trust: Option<&UpstreamTls>,
593    connect_timeout: Duration,
594    read_timeout: Duration,
595) -> Result<Response, String> {
596    if tls {
597        #[cfg(any(feature = "http-client", feature = "http2"))]
598        {
599            let (pool, config) = match trust {
600                Some(t) => (t.pool, Arc::clone(&t.config)),
601                None => (shared_tls_pool(), tls_client_config()),
602            };
603            return forward_via(pool, request, client_ip, host, port, read_timeout, || {
604                connect_tls_with(Arc::clone(&config), host, port, connect_timeout)
605            });
606        }
607        #[cfg(not(any(feature = "http-client", feature = "http2")))]
608        {
609            let _ = trust;
610            return Err("HTTPS backends need the http2 or http-client feature, which this build doesn't include".to_string());
611        }
612    }
613    forward_pooled(plain_pool, request, client_ip, host, port, connect_timeout, read_timeout)
614}
615
616/// Trust settings for one upstream's TLS backends: the certificates in a CA
617/// file **instead of** the WebPKI roots — for internal services behind a
618/// private CA or a self-signed certificate (nginx's
619/// `proxy_ssl_trusted_certificate`). Connections are pooled per CA file, so
620/// two upstreams at the same address with different CAs never share one.
621pub struct UpstreamTls {
622    #[cfg(any(feature = "http-client", feature = "http2"))]
623    config: Arc<rustls::ClientConfig>,
624    #[cfg(any(feature = "http-client", feature = "http2"))]
625    pool: &'static Pool<TlsStream>,
626    /// Connections for the non-blocking forwarder, made with `config`.
627    #[cfg(feature = "http2")]
628    async_pool: &'static Pool<async_upstream::AsyncTlsStream>,
629}
630
631impl UpstreamTls {
632    /// Load the PEM certificates in `path` as the only trusted roots.
633    pub fn from_ca_file(path: &str) -> Result<UpstreamTls, String> {
634        #[cfg(any(feature = "http-client", feature = "http2"))]
635        {
636            use rustls::pki_types::{pem::PemObject, CertificateDer};
637            let bytes = std::fs::read(path).map_err(|e| format!("cannot read CA file '{path}': {e}"))?;
638            let certs = CertificateDer::pem_slice_iter(&bytes)
639                .collect::<Result<Vec<_>, _>>()
640                .map_err(|e| format!("cannot parse CA file '{path}': {e}"))?;
641            if certs.is_empty() {
642                return Err(format!("no certificates in CA file '{path}'"));
643            }
644            let mut roots = rustls::RootCertStore::empty();
645            for cert in certs {
646                roots.add(cert).map_err(|e| format!("invalid CA certificate in '{path}': {e}"))?;
647            }
648            let config = Arc::new(crate::tls_client_builder().with_root_certificates(roots).with_no_client_auth());
649            Ok(UpstreamTls {
650                config,
651                pool: tls_pool_for_ca(path),
652                #[cfg(feature = "http2")]
653                async_pool: async_tls_pool_for_ca(path),
654            })
655        }
656        #[cfg(not(any(feature = "http-client", feature = "http2")))]
657        Err(format!("tls_ca_file '{path}': HTTPS backends need the http2 or http-client feature"))
658    }
659
660    /// Trust nothing: every TLS handshake fails, so the upstream answers 502.
661    /// Used when a configured CA file can't be loaded — failing closed rather
662    /// than falling back to the public roots.
663    pub(crate) fn reject_all(label: &str) -> UpstreamTls {
664        #[cfg(any(feature = "http-client", feature = "http2"))]
665        {
666            let config = Arc::new(
667                crate::tls_client_builder().with_root_certificates(rustls::RootCertStore::empty()).with_no_client_auth(),
668            );
669            UpstreamTls {
670                config,
671                pool: tls_pool_for_ca(&format!("reject-all:{label}")),
672                #[cfg(feature = "http2")]
673                async_pool: async_tls_pool_for_ca(&format!("reject-all:{label}")),
674            }
675        }
676        #[cfg(not(any(feature = "http-client", feature = "http2")))]
677        {
678            let _ = label;
679            UpstreamTls {}
680        }
681    }
682
683    /// How the non-blocking forwarder reaches these backends.
684    #[cfg(feature = "http2")]
685    pub(crate) fn async_tls(&self) -> async_upstream::AsyncTls {
686        async_upstream::AsyncTls { config: Arc::clone(&self.config), pool: self.async_pool }
687    }
688
689    /// The client configuration (for health checks).
690    #[cfg(any(feature = "http-client", feature = "http2"))]
691    pub(crate) fn client_config(&self) -> Arc<rustls::ClientConfig> {
692        Arc::clone(&self.config)
693    }
694}
695
696/// One non-blocking connection pool per CA file, for the life of the process
697/// (a reload that keeps the same file keeps its warm connections).
698#[cfg(feature = "http2")]
699fn async_tls_pool_for_ca(path: &str) -> &'static Pool<async_upstream::AsyncTlsStream> {
700    static POOLS: std::sync::OnceLock<
701        std::sync::Mutex<std::collections::HashMap<String, &'static Pool<async_upstream::AsyncTlsStream>>>,
702    > = std::sync::OnceLock::new();
703    let mut pools = POOLS.get_or_init(Default::default).lock().unwrap_or_else(|e| e.into_inner());
704    pools
705        .entry(path.to_string())
706        .or_insert_with(|| Box::leak(Box::new(Pool::new(POOL_MAX_IDLE, Duration::from_secs(60)))))
707}
708
709/// One pool per CA file, for the life of the process.
710#[cfg(any(feature = "http-client", feature = "http2"))]
711fn tls_pool_for_ca(path: &str) -> &'static Pool<TlsStream> {
712    static POOLS: std::sync::OnceLock<std::sync::Mutex<std::collections::HashMap<String, &'static Pool<TlsStream>>>> =
713        std::sync::OnceLock::new();
714    let mut pools = POOLS.get_or_init(Default::default).lock().unwrap_or_else(|e| e.into_inner());
715    pools
716        .entry(path.to_string())
717        .or_insert_with(|| Box::leak(Box::new(Pool::new(POOL_MAX_IDLE, Duration::from_secs(60)))))
718}
719
720/// Forward `request` to `host:port` over plain HTTP/1.1, reusing an idle
721/// connection from `pool` when there is one and returning the connection to
722/// `pool` after a complete, reusable response.
723pub(crate) fn forward_pooled(
724    pool: &ConnPool,
725    request: &Request,
726    client_ip: &str,
727    host: &str,
728    port: u16,
729    connect_timeout: Duration,
730    read_timeout: Duration,
731) -> Result<Response, String> {
732    forward_via(pool, request, client_ip, host, port, read_timeout, || connect_tcp(host, port, connect_timeout))
733}
734
735/// The pooled request/response cycle, for any connection type.
736///
737/// A backend may close an idle keep-alive connection at any moment, so a
738/// pooled connection that fails before *any* response byte arrives is
739/// retried once on a fresh connection — but only for idempotent methods
740/// (RFC 9110 §9.2.2): the request may already have reached the backend.
741fn forward_via<S: UpstreamIo>(
742    pool: &Pool<S>,
743    request: &Request,
744    client_ip: &str,
745    host: &str,
746    port: u16,
747    read_timeout: Duration,
748    connect: impl Fn() -> Result<S, String>,
749) -> Result<Response, String> {
750    let key = format!("{}:{}", host, port);
751    if let Some(stream) = pool.acquire(&key) {
752        match exchange(pool, &key, stream, request, host, client_ip, read_timeout) {
753            Exchange::Done(response) => return Ok(response),
754            Exchange::NoResponse(_) if is_idempotent(&request.method) => {} // stale: retry fresh
755            Exchange::NoResponse(e) | Exchange::Failed(e) => return Err(e),
756        }
757    }
758    let stream = connect()?;
759    match exchange(pool, &key, stream, request, host, client_ip, read_timeout) {
760        Exchange::Done(response) => Ok(response),
761        Exchange::NoResponse(e) | Exchange::Failed(e) => Err(e),
762    }
763}
764
765/// Connect to the first address of `host:port` that accepts. `localhost`
766/// resolves to `::1` before `127.0.0.1` on many systems; a backend listening
767/// on IPv4 only must still be reachable by name.
768fn connect_tcp(host: &str, port: u16, connect_timeout: Duration) -> Result<TcpStream, String> {
769    let key = format!("{}:{}", host, port);
770    let addrs: Vec<_> = key
771        .to_socket_addrs()
772        .map_err(|e| format!("DNS lookup for {} failed: {}", key, e))?
773        .collect();
774    let mut last_error = format!("no address resolved for {}", key);
775    for addr in addrs {
776        match TcpStream::connect_timeout(&addr, connect_timeout) {
777            Ok(stream) => {
778                let _ = stream.set_nodelay(true);
779                return Ok(stream);
780            }
781            Err(e) => last_error = format!("connect to {} ({}) failed: {}", key, addr, e),
782        }
783    }
784    Err(last_error)
785}
786
787/// The process-wide pool of TLS connections to `https://` backends.
788#[cfg(any(feature = "http-client", feature = "http2"))]
789fn shared_tls_pool() -> &'static Pool<TlsStream> {
790    static POOL: std::sync::OnceLock<Pool<TlsStream>> = std::sync::OnceLock::new();
791    POOL.get_or_init(|| Pool::new(POOL_MAX_IDLE, Duration::from_secs(60)))
792}
793
794/// TLS client settings for backends: the Mozilla root store, no client
795/// certificate. Built once — parsing the root store per request was a large
796/// part of the per-request cost of HTTPS upstreams.
797#[cfg(any(feature = "http-client", feature = "http2"))]
798pub(crate) fn tls_client_config() -> Arc<rustls::ClientConfig> {
799    static CONFIG: std::sync::OnceLock<Arc<rustls::ClientConfig>> = std::sync::OnceLock::new();
800    Arc::clone(CONFIG.get_or_init(|| {
801        let roots = rustls::RootCertStore::from_iter(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
802        Arc::new(crate::tls_client_builder().with_root_certificates(roots).with_no_client_auth())
803    }))
804}
805
806#[cfg(any(feature = "http-client", feature = "http2"))]
807fn connect_tls_with(config: Arc<rustls::ClientConfig>, host: &str, port: u16, connect_timeout: Duration) -> Result<TlsStream, String> {
808    let tcp = connect_tcp(host, port, connect_timeout)?;
809    let server_name = rustls::pki_types::ServerName::try_from(host.to_string())
810        .map_err(|e| format!("invalid upstream hostname '{}': {}", host, e))?;
811    let conn = rustls::ClientConnection::new(config, server_name).map_err(|e| e.to_string())?;
812    Ok(rustls::StreamOwned::new(conn, tcp))
813}
814
815fn is_idempotent(method: &str) -> bool {
816    matches!(method, "GET" | "HEAD" | "OPTIONS" | "TRACE" | "PUT" | "DELETE")
817}
818
819enum Exchange {
820    Done(Response),
821    /// Failed before any response byte (write error, or closed/errored on the first read).
822    NoResponse(String),
823    Failed(String),
824}
825
826/// One request/response on `stream`. A fully read, keep-alive response puts
827/// the connection back in `pool`; a streamed response keeps it until done.
828fn exchange<S: UpstreamIo>(
829    pool: &Pool<S>,
830    key: &str,
831    mut stream: S,
832    request: &Request,
833    host: &str,
834    client_ip: &str,
835    read_timeout: Duration,
836) -> Exchange {
837    if let Err(e) = stream.set_io_timeouts(read_timeout, Duration::from_secs(10)) {
838        return Exchange::Failed(e.to_string());
839    }
840    // keep_alive = true: ask the backend to hold the connection open.
841    let req_bytes = build_request(request, host, client_ip, true);
842    if let Err(e) = stream.write_all(&req_bytes) {
843        return Exchange::NoResponse(format!("write to backend failed: {}", e));
844    }
845
846    let mut tmp = [0u8; 4096];
847    let (header_bytes, body_prefix) = match read_final_response_head(&mut stream, &mut tmp) {
848        Ok(head) => head,
849        Err(HeadError::NothingReceived(e)) => return Exchange::NoResponse(e),
850        Err(HeadError::Partial(e)) => return Exchange::Failed(e),
851    };
852    let header_lower = std::str::from_utf8(&header_bytes).unwrap_or("").to_ascii_lowercase();
853
854    if response_has_no_body(&request.method, &header_lower) {
855        if !header_lower.lines().any(|l| l.starts_with("connection:") && l.contains("close")) {
856            pool.release(key, stream);
857        }
858        return match Response::parse_upstream(&header_bytes, REGENERATED_OR_HOP_BY_HOP) {
859            Ok(response) => Exchange::Done(response),
860            Err(e) => Exchange::Failed(e),
861        };
862    }
863
864    if should_stream_response(&header_lower) {
865        // Streaming path — pipe bytes straight to the client; the connection
866        // can't be reused while the body is in flight.
867        return match parse_status_and_headers(&header_bytes) {
868            Ok(mut response) => {
869                response.stream_pipe = Some(Box::new(ConcatReader::new(body_prefix, stream)));
870                Exchange::Done(response)
871            }
872            Err(e) => Exchange::Failed(e),
873        };
874    }
875    // Buffered path — read the full body, then return the connection to the pool.
876    match read_response_from_partial(&mut stream, header_bytes, body_prefix, &mut tmp) {
877        Ok((resp_bytes, reusable)) => {
878            if reusable {
879                pool.release(key, stream);
880            }
881            match Response::parse_upstream(&resp_bytes, REGENERATED_OR_HOP_BY_HOP) {
882                Ok(response) => Exchange::Done(response),
883                Err(e) => Exchange::Failed(e),
884            }
885        }
886        Err(e) => Exchange::Failed(e),
887    }
888}
889
890/// Read exactly the HTTP response headers (up to and including `\r\n\r\n`).
891///
892/// Returns `(header_bytes, body_prefix)` where `body_prefix` contains any
893/// body bytes that arrived in the same TCP segment as the headers.
894#[cfg(test)]
895fn read_headers_only(stream: &mut TcpStream, tmp: &mut [u8]) -> Result<(Vec<u8>, Vec<u8>), String> {
896    read_response_head(stream, tmp).map_err(|e| match e {
897        HeadError::NothingReceived(msg) | HeadError::Partial(msg) => msg,
898    })
899}
900
901/// Why reading a response head failed: before any byte arrived (the request
902/// may never have been processed — e.g. a pooled connection the backend had
903/// already closed) or part-way through.
904enum HeadError {
905    NothingReceived(String),
906    Partial(String),
907}
908
909fn read_response_head<S: Read>(stream: &mut S, tmp: &mut [u8]) -> Result<(Vec<u8>, Vec<u8>), HeadError> {
910    read_response_head_from(stream, tmp, Vec::with_capacity(4096), false)
911}
912
913/// The final response head: interim `1xx` heads (`100 Continue`,
914/// `103 Early Hints`, …) are read past (RFC 9110 §15.2). Returning one as
915/// the response left the real response in the connection, and once pooled
916/// the next request — possibly another client's — received it.
917fn read_final_response_head<S: Read>(stream: &mut S, tmp: &mut [u8]) -> Result<(Vec<u8>, Vec<u8>), HeadError> {
918    let (mut head, mut rest) = read_response_head(stream, tmp)?;
919    let mut interim = 0;
920    while is_interim_response(&head) {
921        interim += 1;
922        if interim > MAX_INTERIM_RESPONSES {
923            return Err(HeadError::Partial("too many interim (1xx) responses from the backend".to_string()));
924        }
925        (head, rest) = read_response_head_from(stream, tmp, rest, true)?;
926    }
927    Ok((head, rest))
928}
929
930/// Read until a complete head is in `buf` (which may already hold bytes
931/// that followed a previous head). `received` — bytes already came from this
932/// response, so a failure is never "nothing received".
933fn read_response_head_from<S: Read>(
934    stream: &mut S,
935    tmp: &mut [u8],
936    mut buf: Vec<u8>,
937    received: bool,
938) -> Result<(Vec<u8>, Vec<u8>), HeadError> {
939    loop {
940        if let Some(pos) = buf.windows(4).position(|w| w == b"\r\n\r\n") {
941            let body_prefix = buf[pos + 4..].to_vec();
942            buf.truncate(pos + 4);
943            return Ok((buf, body_prefix));
944        }
945        let nothing = buf.is_empty() && !received;
946        let n = match stream.read(tmp) {
947            Ok(n) => n,
948            Err(e) if nothing => return Err(HeadError::NothingReceived(e.to_string())),
949            Err(e) => return Err(HeadError::Partial(e.to_string())),
950        };
951        if n == 0 {
952            let msg = "backend closed connection before headers were complete".to_string();
953            return Err(if nothing { HeadError::NothingReceived(msg) } else { HeadError::Partial(msg) });
954        }
955        buf.extend_from_slice(&tmp[..n]);
956    }
957}
958
959/// Interim responses tolerated before the final one.
960pub(crate) const MAX_INTERIM_RESPONSES: usize = 16;
961
962/// An interim (1xx) response head other than `101 Switching Protocols`
963/// (which ends the HTTP exchange): the final response follows it.
964pub(crate) fn is_interim_response(head: &[u8]) -> bool {
965    let line_end = head.iter().position(|b| *b == b'\r' || *b == b'\n').unwrap_or(head.len());
966    let status = head[..line_end].split(|b| *b == b' ').nth(1).unwrap_or(&[]);
967    status.len() == 3 && status.iter().all(u8::is_ascii_digit) && status[0] == b'1' && status != b"101"
968}
969
970/// RFC 9112 §6.3: a response to HEAD, and any 1xx, 204 or 304 response, ends
971/// with its head — whatever `Content-Length` says (for HEAD it describes the
972/// GET body). Waiting for that body hung until the read timeout (a 502).
973pub(crate) fn response_has_no_body(method: &str, header_lower: &str) -> bool {
974    let status = header_lower.split_whitespace().nth(1).unwrap_or("");
975    method.eq_ignore_ascii_case("HEAD") || status.starts_with('1') || status == "204" || status == "304"
976}
977
978/// Returns `true` when the response should be streamed rather than buffered.
979///
980/// Streams when any of the following hold:
981/// - `Content-Type: text/event-stream` — SSE
982/// - `Transfer-Encoding: chunked` — AI token streams, etc.
983/// - `Content-Length` exceeds 1 MB — large file downloads
984pub(crate) fn should_stream_response(header_lower: &str) -> bool {
985    let is_sse = header_lower.lines().any(|l| {
986        l.starts_with("content-type:") && l.contains("text/event-stream")
987    });
988    let is_chunked = header_lower.lines().any(|l| {
989        l.starts_with("transfer-encoding:") && l.contains("chunked")
990    });
991    let content_length: Option<usize> = header_lower.lines().find_map(|l| {
992        l.strip_prefix("content-length:")?.trim().parse().ok()
993    });
994    let is_large = content_length.map_or(false, |n| n > STREAM_THRESHOLD);
995    is_sse || is_chunked || is_large
996}
997
998/// What the forwarders need from an upstream response head, found in one
999/// case-insensitive pass over the raw bytes — the same answers as
1000/// [`response_has_no_body`], [`should_stream_response`] and the
1001/// `connection:` / `content-length:` scans over a lower-cased copy, without
1002/// making that copy.
1003#[cfg_attr(not(feature = "http2"), allow(dead_code))]
1004pub(crate) struct HeadFacts {
1005    /// The second whitespace-separated token of the head (the status code).
1006    pub(crate) status_no_body: bool,
1007    pub(crate) sse: bool,
1008    pub(crate) chunked: bool,
1009    /// The first `content-length:` value that parses.
1010    pub(crate) content_length: Option<usize>,
1011    pub(crate) connection_close: bool,
1012}
1013
1014#[cfg_attr(not(feature = "http2"), allow(dead_code))]
1015impl HeadFacts {
1016    pub(crate) fn scan(head: &[u8]) -> HeadFacts {
1017        let text = std::str::from_utf8(head).unwrap_or("");
1018        let status = text.split_whitespace().nth(1).unwrap_or("");
1019        let mut facts = HeadFacts {
1020            status_no_body: status.starts_with('1') || status == "204" || status == "304",
1021            sse: false,
1022            chunked: false,
1023            content_length: None,
1024            connection_close: false,
1025        };
1026        for line in text.lines() {
1027            if let Some(value) = strip_prefix_ignore_ascii_case(line, "content-type:") {
1028                facts.sse |= contains_ignore_ascii_case(value, "text/event-stream");
1029            } else if let Some(value) = strip_prefix_ignore_ascii_case(line, "transfer-encoding:") {
1030                facts.chunked |= contains_ignore_ascii_case(value, "chunked");
1031            } else if let Some(value) = strip_prefix_ignore_ascii_case(line, "content-length:") {
1032                if facts.content_length.is_none() {
1033                    facts.content_length = value.trim().parse().ok();
1034                }
1035            } else if let Some(value) = strip_prefix_ignore_ascii_case(line, "connection:") {
1036                facts.connection_close |= contains_ignore_ascii_case(value, "close");
1037            }
1038        }
1039        facts
1040    }
1041
1042    /// [`response_has_no_body`].
1043    pub(crate) fn no_body(&self, method: &str) -> bool {
1044        method.eq_ignore_ascii_case("HEAD") || self.status_no_body
1045    }
1046
1047    /// [`should_stream_response`].
1048    pub(crate) fn should_stream(&self) -> bool {
1049        self.sse || self.chunked || self.content_length.is_some_and(|n| n > STREAM_THRESHOLD)
1050    }
1051}
1052
1053fn strip_prefix_ignore_ascii_case<'a>(line: &'a str, prefix: &str) -> Option<&'a str> {
1054    let bytes = line.as_bytes();
1055    // `prefix` is ASCII, so a byte match ends on a char boundary.
1056    (bytes.len() >= prefix.len() && bytes[..prefix.len()].eq_ignore_ascii_case(prefix.as_bytes()))
1057        .then(|| &line[prefix.len()..])
1058}
1059
1060fn contains_ignore_ascii_case(haystack: &str, needle: &str) -> bool {
1061    haystack.as_bytes().windows(needle.len()).any(|w| w.eq_ignore_ascii_case(needle.as_bytes()))
1062}
1063
1064/// Parse the status line and headers from raw header bytes (ending at `\r\n\r\n`).
1065fn parse_status_and_headers(header_bytes: &[u8]) -> Result<Response, String> {
1066    let s = std::str::from_utf8(header_bytes)
1067        .map_err(|e| format!("non-UTF-8 response headers: {}", e))?;
1068    let mut lines = s.lines();
1069    let status_line = lines.next().ok_or("empty backend response")?;
1070    let mut parts = status_line.splitn(3, ' ');
1071    let http_version = parts.next().unwrap_or("HTTP/1.1").to_string();
1072    let status_code: i16 = parts
1073        .next()
1074        .unwrap_or("502")
1075        .parse()
1076        .map_err(|_| format!("invalid status code in '{}'", status_line))?;
1077    let reason_phrase = parts.next().unwrap_or("").trim_end_matches('\r').to_string();
1078    let mut headers = Vec::new();
1079    for line in lines {
1080        let line = line.trim_end_matches('\r');
1081        if line.is_empty() { break; }
1082        if let Some(colon) = line.find(':') {
1083            headers.push(crate::header::Header {
1084                name: line[..colon].trim().to_string().into(),
1085                value: (line[colon + 1..].trim().to_string()).into(),
1086            });
1087        }
1088    }
1089    Ok(Response {
1090        http_version,
1091        status_code,
1092        reason_phrase,
1093        headers,
1094        content_range_list: vec![],
1095        stream_file: None,
1096        stream_pipe: None,
1097    })
1098}
1099
1100/// Read the remaining body after headers have already been read.
1101///
1102/// `header_bytes` ends with `\r\n\r\n`; `body_prefix` holds any body bytes
1103/// that arrived in the same TCP read.  Handles all three body mechanisms
1104/// (chunked, content-length, read-to-EOF).  Returns `(full_response_bytes, can_reuse)`.
1105fn read_response_from_partial<S: Read>(
1106    stream: &mut S,
1107    header_bytes: Vec<u8>,
1108    body_prefix: Vec<u8>,
1109    tmp: &mut [u8],
1110) -> Result<(Vec<u8>, bool), String> {
1111    let header_end = header_bytes.len();
1112    let mut buf = header_bytes;
1113    buf.extend_from_slice(&body_prefix);
1114
1115    let header_lower =
1116        std::str::from_utf8(&buf[..header_end]).unwrap_or("").to_ascii_lowercase();
1117    let connection_close =
1118        header_lower.lines().any(|l| l.starts_with("connection:") && l.contains("close"));
1119    let is_chunked = header_lower
1120        .lines()
1121        .any(|l| l.starts_with("transfer-encoding:") && l.contains("chunked"));
1122    let content_length: Option<usize> = header_lower.lines().find_map(|l| {
1123        l.strip_prefix("content-length:")?.trim().parse().ok()
1124    });
1125
1126    if is_chunked {
1127        let decoded = decode_chunked(stream, &buf, header_end, tmp)?;
1128        rewrite_as_content_length(&mut buf, header_end, &decoded);
1129        Ok((buf, !connection_close))
1130    } else if let Some(len) = content_length {
1131        while buf.len() < header_end + len {
1132            let n = stream.read(tmp).map_err(|e| e.to_string())?;
1133            if n == 0 { break; }
1134            buf.extend_from_slice(&tmp[..n]);
1135        }
1136        Ok((buf, !connection_close))
1137    } else {
1138        loop {
1139            let n = stream.read(tmp).map_err(|e| e.to_string())?;
1140            if n == 0 { break; }
1141            buf.extend_from_slice(&tmp[..n]);
1142        }
1143        Ok((buf, false))
1144    }
1145}
1146
1147// ── Backend URL parsing ───────────────────────────────────────────────────────
1148
1149struct Backend {
1150    host: String,
1151    port: u16,
1152    /// Whether the upstream connection should use TLS.
1153    /// Set when the URL scheme is `https://`, `h2s://`, or `grpcs://`.
1154    tls: bool,
1155}
1156
1157impl Backend {
1158    fn parse(url: &str) -> Option<Self> {
1159        let (rest, tls, default_port) = if let Some(r) = url.strip_prefix("https://") {
1160            (r, true, 443u16)
1161        } else if let Some(r) = url.strip_prefix("h2s://") {
1162            (r, true, 443u16)
1163        } else if let Some(r) = url.strip_prefix("grpcs://") {
1164            (r, true, 443u16)
1165        } else if let Some(r) = url.strip_prefix("http://") {
1166            (r, false, 80u16)
1167        } else if let Some(r) = url.strip_prefix("h2://") {
1168            (r, false, 80u16)
1169        } else if let Some(r) = url.strip_prefix("grpc://") {
1170            (r, false, 80u16)
1171        } else {
1172            (url, false, 80u16)
1173        };
1174        // Drop any path component.
1175        let host_port = rest.split('/').next().unwrap_or(rest);
1176        let (host, port) = if let Some(colon) = host_port.rfind(':') {
1177            let port_str = &host_port[colon + 1..];
1178            if let Ok(p) = port_str.parse::<u16>() {
1179                (host_port[..colon].to_string(), p)
1180            } else {
1181                (host_port.to_string(), default_port)
1182            }
1183        } else {
1184            (host_port.to_string(), default_port)
1185        };
1186        if host.is_empty() {
1187            return None;
1188        }
1189        Some(Backend { host, port, tls })
1190    }
1191}
1192
1193// ── HTTP/2 reverse proxy ──────────────────────────────────────────────────────
1194
1195/// Reverse proxy that forwards requests to HTTP/2 backends.
1196///
1197/// Wraps [`ReverseProxy`] and forces HTTP/2 (`h2`) for all upstream connections.
1198/// Requires the `http2` Cargo feature.
1199///
1200/// This proxy also transparently handles gRPC traffic
1201/// (`Content-Type: application/grpc*`) — gRPC DATA frames are forwarded
1202/// as-is because gRPC is layered directly on HTTP/2.
1203#[cfg(feature = "http2")]
1204pub struct H2ReverseProxy {
1205    inner: ReverseProxy,
1206}
1207
1208#[cfg(feature = "http2")]
1209impl H2ReverseProxy {
1210    /// Create a proxy distributing requests across `backends` in round-robin order.
1211    ///
1212    /// Each backend entry can be:
1213    /// - `"host:port"` — plain TCP (HTTP/2 cleartext)
1214    /// - `"h2://host:port"` — plain TCP (explicit scheme)
1215    /// - `"h2s://host:port"` — TLS (HTTP/2 over HTTPS; port defaults to 443)
1216    /// - `"https://host:port"` — TLS (same as `h2s://`)
1217    ///
1218    /// TLS backends require the `http2` Cargo feature (includes `rustls` +
1219    /// `webpki-roots`).  Certificate verification uses the WebPKI trust store.
1220    pub fn new<I, S>(backends: I) -> Self
1221    where
1222        I: IntoIterator<Item = S>,
1223        S: AsRef<str>,
1224    {
1225        H2ReverseProxy {
1226            inner: ReverseProxy::new(backends),
1227        }
1228    }
1229
1230    /// Only proxy requests whose URI starts with `prefix`; pass others through.
1231    pub fn path_prefix(mut self, prefix: impl Into<String>) -> Self {
1232        self.inner = self.inner.path_prefix(prefix);
1233        self
1234    }
1235
1236    /// Override the TCP connect timeout (default: 5 s).
1237    pub fn connect_timeout_ms(mut self, ms: u64) -> Self {
1238        self.inner = self.inner.connect_timeout_ms(ms);
1239        self
1240    }
1241
1242    /// Override the response read timeout (default: 30 s).
1243    pub fn read_timeout_ms(mut self, ms: u64) -> Self {
1244        self.inner = self.inner.read_timeout_ms(ms);
1245        self
1246    }
1247}
1248
1249#[cfg(feature = "http2")]
1250impl crate::middleware::Middleware for H2ReverseProxy {
1251    fn handle(
1252        &self,
1253        request: &crate::request::Request,
1254        connection: &crate::server::ConnectionInfo,
1255        next: &dyn crate::application::Application,
1256    ) -> Result<crate::response::Response, String> {
1257        if let Some(prefix) = &self.inner.path_prefix {
1258            if !request.request_uri.starts_with(prefix.as_str()) {
1259                return next.execute(request, connection);
1260            }
1261        }
1262        if self.inner.backends.is_empty() {
1263            return Ok(bad_gateway());
1264        }
1265        let n = self.inner.backends.len();
1266        let start = self.inner.counter.fetch_add(1, Ordering::Relaxed);
1267        for attempt in 0..n {
1268            let idx = (start + attempt) % n;
1269            match try_backend_h2(request, &connection.client.ip, &self.inner.backends[idx],
1270                                  self.inner.connect_timeout, self.inner.read_timeout) {
1271                Ok(resp) => return Ok(resp),
1272                Err(_) if attempt + 1 < n => continue,
1273                Err(_) => break,
1274            }
1275        }
1276        Ok(bad_gateway())
1277    }
1278}
1279
1280#[cfg(feature = "http2")]
1281fn try_backend_h2(
1282    request: &Request,
1283    client_ip: &str,
1284    backend: &Backend,
1285    connect_timeout: Duration,
1286    read_timeout: Duration,
1287) -> Result<Response, String> {
1288    // `block_on_isolated` works on any tokio runtime flavor (including
1289    // `current_thread`, unlike `tokio::task::block_in_place`, which requires
1290    // `multi_thread` and panics otherwise) and doesn't pull a worker thread
1291    // out of a shared pool; it also works with no runtime at all (HTTP/1.1
1292    // thread pool), unlike the previous implementation, which just gave up.
1293    crate::async_bridge::block_on_isolated(|| forward_h2_async(request, client_ip, backend, connect_timeout, read_timeout))
1294}
1295
1296#[cfg(feature = "http2")]
1297async fn forward_h2_async(
1298    request: &Request,
1299    client_ip: &str,
1300    backend: &Backend,
1301    connect_timeout: Duration,
1302    read_timeout: Duration,
1303) -> Result<Response, String> {
1304    let addr = format!("{}:{}", backend.host, backend.port);
1305    let tcp = tokio::time::timeout(
1306        connect_timeout,
1307        tokio::net::TcpStream::connect(&addr),
1308    )
1309    .await
1310    .map_err(|_| format!("h2 proxy: connect to {} timed out", addr))?
1311    .map_err(|e| format!("h2 proxy: connect to {} failed: {}", addr, e))?;
1312
1313    if backend.tls {
1314        use rustls::pki_types::ServerName;
1315        use std::sync::Arc;
1316        use tokio_rustls::TlsConnector;
1317
1318        let root_store = rustls::RootCertStore::from_iter(
1319            webpki_roots::TLS_SERVER_ROOTS.iter().cloned(),
1320        );
1321        let mut config = crate::tls_client_builder()
1322            .with_root_certificates(root_store)
1323            .with_no_client_auth();
1324        // Advertise h2 via ALPN so the server selects HTTP/2.
1325        config.alpn_protocols = vec![b"h2".to_vec()];
1326        let connector = TlsConnector::from(Arc::new(config));
1327        let server_name = ServerName::try_from(backend.host.as_str())
1328            .map_err(|e| format!("invalid upstream hostname '{}': {}", backend.host, e))?
1329            .to_owned();
1330        let tls_stream = within(read_timeout, &addr, connector.connect(server_name, tcp))
1331            .await?
1332            .map_err(|e| format!("h2 proxy: TLS handshake with {} failed: {}", addr, e))?;
1333        send_h2_request(request, client_ip, backend, tls_stream, read_timeout).await
1334    } else {
1335        send_h2_request(request, client_ip, backend, tcp, read_timeout).await
1336    }
1337}
1338
1339/// Await `f`, failing after `limit` — the h2 counterpart of a socket read
1340/// timeout, applied to each wait for the upstream (handshake, response head,
1341/// every body frame), so a slow but progressing response isn't cut off.
1342#[cfg(feature = "http2")]
1343async fn within<F: std::future::Future>(limit: Duration, addr: &str, f: F) -> Result<F::Output, String> {
1344    tokio::time::timeout(limit, f)
1345        .await
1346        .map_err(|_| format!("h2 proxy: {} sent nothing for {} ms", addr, limit.as_millis()))
1347}
1348
1349/// Drive the h2 client handshake + request/response over any async I/O stream.
1350///
1351/// Accepts both plain `TcpStream` and `TlsStream<TcpStream>` — anything that
1352/// satisfies `AsyncRead + AsyncWrite + Unpin + Send + 'static`.
1353#[cfg(feature = "http2")]
1354async fn send_h2_request<T>(
1355    request: &Request,
1356    client_ip: &str,
1357    backend: &Backend,
1358    stream: T,
1359    read_timeout: Duration,
1360) -> Result<Response, String>
1361where
1362    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
1363{
1364    use bytes::Bytes;
1365    use http as hc;
1366
1367    let addr = format!("{}:{}", backend.host, backend.port);
1368
1369    let (send_req, conn) = within(read_timeout, &addr, h2::client::handshake(stream))
1370        .await?
1371        .map_err(|e| format!("h2 proxy: handshake with {} failed: {}", addr, e))?;
1372
1373    tokio::spawn(async move {
1374        let _ = conn.await;
1375    });
1376
1377    let scheme = if backend.tls { "https" } else { "http" };
1378    let uri_str = format!("{}://{}{}", scheme, addr, request.request_uri);
1379    let uri: hc::Uri = uri_str.parse().map_err(|e: hc::uri::InvalidUri| e.to_string())?;
1380    let method = hc::Method::from_bytes(request.method.as_bytes()).map_err(|e| e.to_string())?;
1381
1382    let mut builder = hc::Request::builder().method(method).uri(uri);
1383    builder = builder.header("host", &backend.host);
1384    for h in &request.headers {
1385        let lower = h.name.to_lowercase();
1386        // An empty name would be sent as a malformed ": value" line.
1387        if lower.is_empty() || HOP_BY_HOP.contains(&lower.as_str()) || lower == "host" {
1388            continue;
1389        }
1390        builder = builder.header(h.name.as_ref(), h.value.as_ref());
1391    }
1392    builder = builder.header("x-forwarded-for", client_ip);
1393    builder = builder.header("via", "2 rws");
1394
1395    let body_bytes = Bytes::from(request.body.to_vec());
1396    let end_of_stream = body_bytes.is_empty();
1397    let http_req = builder.body(()).map_err(|e| e.to_string())?;
1398
1399    let mut send_req = within(read_timeout, &addr, send_req.ready()).await?.map_err(|e| e.to_string())?;
1400    let (resp_future, mut req_body) = send_req
1401        .send_request(http_req, end_of_stream)
1402        .map_err(|e| e.to_string())?;
1403    if !end_of_stream {
1404        req_body.send_data(body_bytes, true).map_err(|e| e.to_string())?;
1405    }
1406
1407    let resp = within(read_timeout, &addr, resp_future).await?.map_err(|e| e.to_string())?;
1408    let (parts, mut body) = resp.into_parts();
1409
1410    let content_type = parts
1411        .headers
1412        .get("content-type")
1413        .and_then(|v| v.to_str().ok())
1414        .unwrap_or("application/octet-stream")
1415        .to_string();
1416
1417    let mut body_bytes: Vec<u8> = Vec::new();
1418    while let Some(chunk) = within(read_timeout, &addr, body.data()).await? {
1419        body_bytes.extend_from_slice(&chunk.map_err(|e| e.to_string())?);
1420    }
1421
1422    let mut response = Response::new();
1423    response.status_code = parts.status.as_u16() as i16;
1424    response.reason_phrase = parts.status.canonical_reason().unwrap_or("").to_string();
1425
1426    const H2_HOP: &[&str] = &[
1427        "connection", "keep-alive", "transfer-encoding", "upgrade", "proxy-connection", "te",
1428    ];
1429    for (name, value) in &parts.headers {
1430        let lower = name.as_str().to_lowercase();
1431        if H2_HOP.contains(&lower.as_str()) {
1432            continue;
1433        }
1434        if let Ok(v) = value.to_str() {
1435            response.headers.push(crate::header::Header {
1436                name: name.as_str().to_string().into(),
1437                value: v.to_string().into(),
1438            });
1439        }
1440    }
1441
1442    if !body_bytes.is_empty() {
1443        response.content_range_list = vec![Range::get_content_range(body_bytes, content_type)];
1444    }
1445
1446    Ok(normalize_upstream_response(response))
1447}
1448
1449// ── gRPC proxy ────────────────────────────────────────────────────────────────
1450
1451/// gRPC reverse proxy middleware.
1452///
1453/// Recognises requests with `Content-Type: application/grpc*` and forwards them
1454/// to a backend over HTTP/2, leaving all other requests to the next layer.
1455///
1456/// Requires the `http2` Cargo feature.
1457///
1458/// # Example
1459///
1460/// ```rust,no_run
1461/// use rust_web_server::app::App;
1462/// use rust_web_server::core::New;
1463/// use rust_web_server::proxy::GrpcProxy;
1464///
1465/// let app = App::new()
1466///     .wrap(GrpcProxy::new(["grpc-service:50051"]));
1467/// ```
1468#[cfg(feature = "http2")]
1469pub struct GrpcProxy {
1470    inner: H2ReverseProxy,
1471}
1472
1473#[cfg(feature = "http2")]
1474impl GrpcProxy {
1475    /// Create a proxy distributing gRPC connections across `backends` in round-robin order.
1476    ///
1477    /// Each backend entry can be:
1478    /// - `"host:port"` — plain TCP (gRPC cleartext)
1479    /// - `"grpc://host:port"` — plain TCP (explicit scheme)
1480    /// - `"grpcs://host:port"` — TLS (gRPC over TLS; port defaults to 443)
1481    /// - `"https://host:port"` — TLS (same as `grpcs://`)
1482    ///
1483    /// TLS backends require the `http2` Cargo feature.
1484    pub fn new<I, S>(backends: I) -> Self
1485    where
1486        I: IntoIterator<Item = S>,
1487        S: AsRef<str>,
1488    {
1489        GrpcProxy { inner: H2ReverseProxy::new(backends) }
1490    }
1491
1492    /// Only proxy requests whose URI starts with `prefix`; pass others through.
1493    pub fn path_prefix(mut self, prefix: impl Into<String>) -> Self {
1494        self.inner = self.inner.path_prefix(prefix);
1495        self
1496    }
1497}
1498
1499#[cfg(feature = "http2")]
1500impl crate::middleware::Middleware for GrpcProxy {
1501    fn handle(
1502        &self,
1503        request: &crate::request::Request,
1504        connection: &crate::server::ConnectionInfo,
1505        next: &dyn crate::application::Application,
1506    ) -> Result<crate::response::Response, String> {
1507        let ct = request
1508            .get_header("content-type".to_string())
1509            .map(|h| h.value.as_ref())
1510            .unwrap_or("");
1511        if ct.starts_with("application/grpc") {
1512            self.inner.handle(request, connection, next)
1513        } else {
1514            next.execute(request, connection)
1515        }
1516    }
1517}
1518
1519// ── Backend::parse unit tests ─────────────────────────────────────────────────
1520
1521#[cfg(test)]
1522mod backend_parse_tests {
1523    use super::Backend;
1524
1525    fn parse(url: &str) -> Option<(String, u16, bool)> {
1526        Backend::parse(url).map(|b| (b.host, b.port, b.tls))
1527    }
1528
1529    #[test]
1530    fn bare_host_port() {
1531        assert_eq!(Some(("api.example.com".into(), 8080, false)), parse("api.example.com:8080"));
1532    }
1533
1534    #[test]
1535    fn http_scheme() {
1536        assert_eq!(Some(("backend".into(), 3000, false)), parse("http://backend:3000"));
1537    }
1538
1539    #[test]
1540    fn h2_scheme_plain() {
1541        assert_eq!(Some(("svc".into(), 50051, false)), parse("h2://svc:50051"));
1542    }
1543
1544    #[test]
1545    fn grpc_scheme_plain() {
1546        assert_eq!(Some(("svc".into(), 50051, false)), parse("grpc://svc:50051"));
1547    }
1548
1549    #[test]
1550    fn https_scheme_sets_tls_and_default_port() {
1551        assert_eq!(Some(("api.example.com".into(), 443, true)), parse("https://api.example.com"));
1552    }
1553
1554    #[test]
1555    fn https_scheme_explicit_port() {
1556        assert_eq!(Some(("api.example.com".into(), 8443, true)), parse("https://api.example.com:8443"));
1557    }
1558
1559    #[test]
1560    fn h2s_scheme_sets_tls() {
1561        assert_eq!(Some(("svc".into(), 443, true)), parse("h2s://svc"));
1562    }
1563
1564    #[test]
1565    fn h2s_scheme_explicit_port() {
1566        assert_eq!(Some(("svc".into(), 8443, true)), parse("h2s://svc:8443"));
1567    }
1568
1569    #[test]
1570    fn grpcs_scheme_sets_tls() {
1571        assert_eq!(Some(("grpc-svc".into(), 443, true)), parse("grpcs://grpc-svc"));
1572    }
1573
1574    #[test]
1575    fn grpcs_scheme_explicit_port() {
1576        assert_eq!(Some(("grpc-svc".into(), 50052, true)), parse("grpcs://grpc-svc:50052"));
1577    }
1578
1579    #[test]
1580    fn empty_host_returns_none() {
1581        assert_eq!(None, parse("https://"));
1582    }
1583
1584    #[test]
1585    fn bare_host_no_port_defaults_to_80() {
1586        assert_eq!(Some(("myhost".into(), 80, false)), parse("myhost"));
1587    }
1588}