zlayer_proxy/

server.rs

//! HTTP server implementation
//!
//! This module provides the HTTP/HTTPS server for the proxy.
//! Uses `ServiceRegistry` for route resolution instead of the legacy `Router`.

use crate::acme::CertManager;
use crate::config::ProxyConfig;
use crate::error::{ProxyError, Result};
use crate::lb::LoadBalancer;
use crate::network_policy::NetworkPolicyChecker;
use crate::routes::ServiceRegistry;
use crate::service::{Activator, ReverseProxyService, RpsRegistry};
use crate::sni_resolver::SniCertResolver;
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::Request;
use hyper_util::rt::TokioIo;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use tokio::net::{lookup_host, TcpListener, TcpStream};
use tokio::sync::watch;
use tokio_rustls::TlsAcceptor;
use tracing::{debug, error, info, warn};

/// Initial backoff between failed ingress bind attempts.
const INGRESS_BIND_BACKOFF_INITIAL: Duration = Duration::from_secs(2);

/// Maximum backoff between failed ingress bind attempts. The backoff doubles
/// from [`INGRESS_BIND_BACKOFF_INITIAL`] up to this cap, then stays there.
const INGRESS_BIND_BACKOFF_MAX: Duration = Duration::from_secs(30);

/// After the first warning, only warn again every N attempts so a long-held
/// port (e.g. another process owning :80) does not spam the log every backoff
/// tick. With the cap above this works out to roughly one warning per minute
/// in steady state.
const INGRESS_BIND_WARN_EVERY: u64 = 30;

/// Compute the next backoff for the ingress bind-retry loop: double the current
/// delay, capped at [`INGRESS_BIND_BACKOFF_MAX`].
///
/// Pulled out as a pure function so the backoff schedule can be unit-tested
/// without binding real sockets.
#[must_use]
fn next_ingress_backoff(current: Duration) -> Duration {
    (current * 2).min(INGRESS_BIND_BACKOFF_MAX)
}

/// Decide whether the ingress bind-retry loop should emit a warning on this
/// attempt. Always warns on the first attempt (`attempt == 0`), then only once
/// every [`INGRESS_BIND_WARN_EVERY`] attempts thereafter.
///
/// Pure function so the warn cadence is unit-testable.
#[must_use]
fn should_warn_on_attempt(attempt: u64) -> bool {
    attempt == 0 || attempt % INGRESS_BIND_WARN_EVERY == 0
}

/// Bind a [`TcpListener`] to `addr`, retrying FOREVER on failure.
///
/// This is the ingress bind path: a bind failure is NEVER fatal. On each
/// failure it logs (warning cadence per [`should_warn_on_attempt`]), sleeps a
/// capped exponential backoff, and tries again. It only returns once the bind
/// succeeds (or `shutdown_rx` flips to `true`, in which case it returns
/// `None`).
///
/// `EACCES` (permission denied) gets a dedicated one-time hint, since binding
/// 80/443 needs root or `CAP_NET_BIND_SERVICE`.
async fn bind_with_retry(
    addr: SocketAddr,
    label: &str,
    shutdown_rx: &mut watch::Receiver<bool>,
) -> Option<TcpListener> {
    let mut attempt: u64 = 0;
    let mut backoff = INGRESS_BIND_BACKOFF_INITIAL;
    let mut warned_eacces = false;

    loop {
        if *shutdown_rx.borrow() {
            return None;
        }

        match TcpListener::bind(addr).await {
            Ok(listener) => {
                if attempt > 0 {
                    info!(addr = %addr, label, attempt, "Ingress bound after retrying");
                }
                return Some(listener);
            }
            Err(e) => {
                let is_eacces = e.kind() == std::io::ErrorKind::PermissionDenied;
                if is_eacces && !warned_eacces {
                    warned_eacces = true;
                    warn!(
                        addr = %addr,
                        label,
                        error = %e,
                        "Ingress bind denied: binding 80/443 needs root or CAP_NET_BIND_SERVICE; \
                         will keep retrying without aborting startup"
                    );
                } else if should_warn_on_attempt(attempt) {
                    warn!(
                        addr = %addr,
                        label,
                        attempt,
                        error = %e,
                        "Ingress bind failed; will keep retrying (port may be held by another process)"
                    );
                } else {
                    debug!(addr = %addr, label, attempt, error = %e, "Ingress bind retry");
                }

                tokio::select! {
                    () = tokio::time::sleep(backoff) => {}
                    _ = shutdown_rx.changed() => {
                        if *shutdown_rx.borrow() {
                            return None;
                        }
                    }
                }

                attempt = attempt.saturating_add(1);
                backoff = next_ingress_backoff(backoff);
            }
        }
    }
}

/// The proxy server
pub struct ProxyServer {
    /// Server configuration
    config: Arc<ProxyConfig>,
    /// Service registry for route resolution
    registry: Arc<ServiceRegistry>,
    /// Load balancer for backend selection
    load_balancer: Arc<LoadBalancer>,
    /// Shutdown signal sender
    shutdown_tx: watch::Sender<bool>,
    /// Shutdown signal receiver
    shutdown_rx: watch::Receiver<bool>,
    /// TLS acceptor for HTTPS connections
    tls_acceptor: Option<TlsAcceptor>,
    /// Certificate manager for ACME challenge responses
    cert_manager: Option<Arc<CertManager>>,
    /// Optional network policy checker for access control enforcement
    network_policy_checker: Option<NetworkPolicyChecker>,
    /// Optional on-demand activator for scale-to-zero services. Threaded into
    /// every per-connection [`ReverseProxyService`].
    activator: Option<Arc<dyn Activator>>,
    /// Optional per-service request-rate registry. Threaded into every
    /// per-connection [`ReverseProxyService`] so routed requests are counted.
    rps_registry: Option<Arc<RpsRegistry>>,
}

impl ProxyServer {
    /// Create a new proxy server
    pub fn new(
        config: ProxyConfig,
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
    ) -> Self {
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        Self {
            config: Arc::new(config),
            registry,
            load_balancer,
            shutdown_tx,
            shutdown_rx,
            tls_acceptor: None,
            cert_manager: None,
            network_policy_checker: None,
            activator: None,
            rps_registry: None,
        }
    }

    /// Create a proxy server with an existing registry (alias for `new`)
    pub fn with_registry(
        config: ProxyConfig,
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
    ) -> Self {
        Self::new(config, registry, load_balancer)
    }

    /// Create a proxy server with TLS via SNI resolver
    pub fn with_tls_resolver(
        config: ProxyConfig,
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
        resolver: Arc<SniCertResolver>,
    ) -> Self {
        let tls_config = rustls::ServerConfig::builder()
            .with_no_client_auth()
            .with_cert_resolver(resolver);
        let acceptor = TlsAcceptor::from(Arc::new(tls_config));
        let (shutdown_tx, shutdown_rx) = watch::channel(false);

        Self {
            config: Arc::new(config),
            registry,
            load_balancer,
            shutdown_tx,
            shutdown_rx,
            tls_acceptor: Some(acceptor),
            cert_manager: None,
            network_policy_checker: None,
            activator: None,
            rps_registry: None,
        }
    }

    /// Set the certificate manager for ACME challenge interception
    #[must_use]
    pub fn with_cert_manager(mut self, cm: Arc<CertManager>) -> Self {
        self.cert_manager = Some(cm);
        self
    }

    /// Set the network policy checker for access control enforcement
    #[must_use]
    pub fn with_network_policy_checker(mut self, checker: NetworkPolicyChecker) -> Self {
        self.network_policy_checker = Some(checker);
        self
    }

    /// Set the on-demand activator for scale-to-zero services.
    ///
    /// The activator is threaded into every per-connection
    /// [`ReverseProxyService`] this server builds, so a request to an idle
    /// (scaled-to-zero) service triggers a wake-up instead of an immediate
    /// `503`.
    #[must_use]
    pub fn with_activator(mut self, activator: Arc<dyn Activator>) -> Self {
        self.activator = Some(activator);
        self
    }

    /// Set the per-service request-rate registry.
    ///
    /// The registry is threaded into every per-connection
    /// [`ReverseProxyService`] this server builds, so each routed request is
    /// recorded for autoscaler RPS metrics.
    #[must_use]
    pub fn with_rps_registry(mut self, rps_registry: Arc<RpsRegistry>) -> Self {
        self.rps_registry = Some(rps_registry);
        self
    }

    /// Check if TLS is enabled
    #[must_use]
    pub fn has_tls(&self) -> bool {
        self.tls_acceptor.is_some()
    }

    /// Get the TLS acceptor if configured
    #[must_use]
    pub fn tls_acceptor(&self) -> Option<&TlsAcceptor> {
        self.tls_acceptor.as_ref()
    }

    /// Get the service registry
    #[must_use]
    pub fn registry(&self) -> Arc<ServiceRegistry> {
        self.registry.clone()
    }

    /// Get the certificate manager, if one is configured.
    ///
    /// Present on a server when [`Self::with_cert_manager`] was called; the
    /// ACME HTTP-01 challenge interception (and TLS resolution on `:443`)
    /// require it. The `:80` ingress carries this purely for the HTTP-01
    /// carve-out — challenges always arrive on `:80`.
    #[must_use]
    pub fn cert_manager(&self) -> Option<&Arc<CertManager>> {
        self.cert_manager.as_ref()
    }

    /// Get the configuration
    #[must_use]
    pub fn config(&self) -> Arc<ProxyConfig> {
        self.config.clone()
    }

    /// Signal the server to shut down
    pub fn shutdown(&self) {
        let _ = self.shutdown_tx.send(true);
    }

    /// Run the HTTP server
    ///
    /// # Errors
    ///
    /// Returns an error if binding to the configured HTTP address fails
    /// or if the accept loop encounters a fatal error.
    pub async fn run(&self) -> Result<()> {
        let addr = self.config.server.http_addr;
        let listener = TcpListener::bind(addr)
            .await
            .map_err(|e| ProxyError::BindFailed {
                addr,
                reason: e.to_string(),
            })?;

        info!(addr = %addr, "HTTP proxy server listening");

        self.accept_loop(listener).await
    }

    /// Run the server on a specific address
    ///
    /// # Errors
    ///
    /// Returns an error if binding to the given address fails
    /// or if the accept loop encounters a fatal error.
    pub async fn run_on(&self, addr: SocketAddr) -> Result<()> {
        let listener = TcpListener::bind(addr)
            .await
            .map_err(|e| ProxyError::BindFailed {
                addr,
                reason: e.to_string(),
            })?;

        info!(addr = %addr, "HTTP proxy server listening");

        self.accept_loop(listener).await
    }

    /// Run the HTTP ingress server on a specific address, retrying the bind
    /// FOREVER on failure instead of returning a fatal error.
    ///
    /// This is the ingress entry point: if `addr` (typically `0.0.0.0:80`) is
    /// already held, it logs a warning and keeps retrying the bind with a
    /// capped exponential backoff until the port frees, then proceeds to the
    /// normal accept loop. It NEVER returns [`ProxyError::BindFailed`].
    ///
    /// # Errors
    ///
    /// Returns an error only if the accept loop itself encounters a fatal
    /// error after a successful bind. Bind failures are non-fatal.
    pub async fn run_with_retry(&self, addr: SocketAddr) -> Result<()> {
        let mut shutdown_rx = self.shutdown_rx.clone();
        let Some(listener) = bind_with_retry(addr, "http", &mut shutdown_rx).await else {
            // Shutdown requested before we ever bound.
            return Ok(());
        };

        info!(addr = %addr, "HTTP ingress server listening");

        self.accept_loop(listener).await
    }

    /// Run the HTTP accept loop on a caller-supplied, already-bound listener.
    ///
    /// Used to attach a second listener (e.g. an explicitly v6-only `[::]:80`
    /// socket) to an existing [`ProxyServer`] so it shares the same shutdown
    /// signal and routing state as the primary v4 listener.
    ///
    /// # Errors
    ///
    /// Returns an error only if the accept loop encounters a fatal error.
    pub async fn run_on_listener(&self, listener: TcpListener) -> Result<()> {
        if let Ok(addr) = listener.local_addr() {
            info!(addr = %addr, "HTTP proxy server listening (pre-bound)");
        }
        self.accept_loop(listener).await
    }

    async fn accept_loop(&self, listener: TcpListener) -> Result<()> {
        let mut shutdown_rx = self.shutdown_rx.clone();

        loop {
            tokio::select! {
                // Check for shutdown signal
                _ = shutdown_rx.changed() => {
                    if *shutdown_rx.borrow() {
                        info!("Shutting down proxy server");
                        break;
                    }
                }

                // Accept new connections
                result = listener.accept() => {
                    match result {
                        Ok((stream, remote_addr)) => {
                            let registry = self.registry.clone();
                            let load_balancer = self.load_balancer.clone();
                            let config = self.config.clone();
                            let cert_manager = self.cert_manager.clone();
                            let npc = self.network_policy_checker.clone();
                            let activator = self.activator.clone();
                            let rps_registry = self.rps_registry.clone();

                            tokio::spawn(async move {
                                if let Err(e) = Self::handle_connection(
                                    stream,
                                    remote_addr,
                                    registry,
                                    load_balancer,
                                    config,
                                    cert_manager,
                                    npc,
                                    activator,
                                    rps_registry,
                                ).await {
                                    debug!(
                                        error = %e,
                                        remote_addr = %remote_addr,
                                        "Connection error"
                                    );
                                }
                            });
                        }
                        Err(e) => {
                            warn!(error = %e, "Failed to accept connection");
                        }
                    }
                }
            }
        }

        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    async fn handle_connection(
        stream: tokio::net::TcpStream,
        remote_addr: SocketAddr,
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
        config: Arc<ProxyConfig>,
        cert_manager: Option<Arc<CertManager>>,
        network_policy_checker: Option<NetworkPolicyChecker>,
        activator: Option<Arc<dyn Activator>>,
        rps_registry: Option<Arc<RpsRegistry>>,
    ) -> Result<()> {
        let io = TokioIo::new(stream);

        let mut service =
            ReverseProxyService::new(registry, load_balancer, config).with_remote_addr(remote_addr);
        if let Some(cm) = cert_manager {
            service = service.with_cert_manager(cm);
        }
        if let Some(checker) = network_policy_checker {
            service = service.with_network_policy_checker(checker);
        }
        if let Some(activator) = activator {
            service = service.with_activator(activator);
        }
        if let Some(rps_registry) = rps_registry {
            service = service.with_rps_registry(rps_registry);
        }

        let service = service_fn(move |req: Request<Incoming>| {
            let svc = service.clone();
            async move {
                match svc.proxy_request(req).await {
                    Ok(response) => Ok::<_, hyper::Error>(response),
                    Err(e) => {
                        error!(error = %e, "Proxy error");
                        Ok(ReverseProxyService::error_response(&e))
                    }
                }
            }
        });

        http1::Builder::new()
            .preserve_header_case(true)
            .title_case_headers(false)
            .serve_connection(io, service)
            .with_upgrades()
            .await
            .map_err(ProxyError::Hyper)?;

        Ok(())
    }

    /// Run the HTTPS server
    ///
    /// This requires TLS to be configured when creating the `ProxyServer`.
    ///
    /// # Errors
    ///
    /// Returns an error if TLS is not configured, if binding to the
    /// configured HTTPS address fails, or if the accept loop encounters a fatal error.
    pub async fn run_https(&self) -> Result<()> {
        let acceptor = self
            .tls_acceptor
            .as_ref()
            .ok_or_else(|| ProxyError::Config("TLS not configured".to_string()))?;

        let addr = self.config.server.https_addr;
        let listener = TcpListener::bind(addr)
            .await
            .map_err(|e| ProxyError::BindFailed {
                addr,
                reason: e.to_string(),
            })?;

        info!(addr = %addr, "HTTPS proxy server listening");

        self.accept_loop_tls(listener, acceptor.clone()).await
    }

    /// Run the HTTPS server on a specific address
    ///
    /// # Errors
    ///
    /// Returns an error if TLS is not configured, if binding to the
    /// given address fails, or if the accept loop encounters a fatal error.
    pub async fn run_https_on(&self, addr: SocketAddr) -> Result<()> {
        let acceptor = self
            .tls_acceptor
            .as_ref()
            .ok_or_else(|| ProxyError::Config("TLS not configured".to_string()))?;

        let listener = TcpListener::bind(addr)
            .await
            .map_err(|e| ProxyError::BindFailed {
                addr,
                reason: e.to_string(),
            })?;

        info!(addr = %addr, "HTTPS proxy server listening");

        self.accept_loop_tls(listener, acceptor.clone()).await
    }

    /// Run the HTTPS ingress server on a specific address, retrying the bind
    /// FOREVER on failure instead of returning a fatal error.
    ///
    /// This is the TLS ingress entry point: if `addr` (typically
    /// `0.0.0.0:443`) is already held, it logs a warning and keeps retrying
    /// the bind with a capped exponential backoff until the port frees, then
    /// proceeds to the normal TLS accept loop. It NEVER returns
    /// [`ProxyError::BindFailed`].
    ///
    /// # Errors
    ///
    /// Returns an error if TLS is not configured, or if the accept loop itself
    /// encounters a fatal error after a successful bind. Bind failures are
    /// non-fatal.
    pub async fn run_https_with_retry(&self, addr: SocketAddr) -> Result<()> {
        let acceptor = self
            .tls_acceptor
            .as_ref()
            .ok_or_else(|| ProxyError::Config("TLS not configured".to_string()))?;

        let mut shutdown_rx = self.shutdown_rx.clone();
        let Some(listener) = bind_with_retry(addr, "https", &mut shutdown_rx).await else {
            // Shutdown requested before we ever bound.
            return Ok(());
        };

        info!(addr = %addr, "HTTPS ingress server listening");

        self.accept_loop_tls(listener, acceptor.clone()).await
    }

    /// Run the HTTPS accept loop on a caller-supplied, already-bound listener.
    ///
    /// Uses this server's configured [`TlsAcceptor`] (the same SNI resolver as
    /// the primary listener), so a second listener (e.g. an explicitly v6-only
    /// `[::]:443` socket) shares certificates, shutdown, and routing state.
    ///
    /// # Errors
    ///
    /// Returns an error if TLS is not configured, or if the accept loop
    /// encounters a fatal error.
    pub async fn run_https_on_listener(&self, listener: TcpListener) -> Result<()> {
        let acceptor = self
            .tls_acceptor
            .as_ref()
            .ok_or_else(|| ProxyError::Config("TLS not configured".to_string()))?;

        if let Ok(addr) = listener.local_addr() {
            info!(addr = %addr, "HTTPS proxy server listening (pre-bound)");
        }
        self.accept_loop_tls(listener, acceptor.clone()).await
    }

    /// Run both HTTP and HTTPS servers concurrently
    ///
    /// This requires TLS to be configured when creating the `ProxyServer`.
    ///
    /// # Errors
    ///
    /// Returns an error if TLS is not configured, if binding to either
    /// the HTTP or HTTPS address fails, or if either accept loop encounters
    /// a fatal error.
    #[allow(clippy::similar_names)]
    pub async fn run_both(&self) -> Result<()> {
        let http_addr = self.config.server.http_addr;
        let https_addr = self.config.server.https_addr;

        let acceptor = self
            .tls_acceptor
            .as_ref()
            .ok_or_else(|| ProxyError::Config("TLS not configured".to_string()))?;

        let http_listener =
            TcpListener::bind(http_addr)
                .await
                .map_err(|e| ProxyError::BindFailed {
                    addr: http_addr,
                    reason: e.to_string(),
                })?;

        let https_listener =
            TcpListener::bind(https_addr)
                .await
                .map_err(|e| ProxyError::BindFailed {
                    addr: https_addr,
                    reason: e.to_string(),
                })?;

        info!(http = %http_addr, https = %https_addr, "Proxy server listening");

        // Run both accept loops concurrently
        let http_future = self.accept_loop(http_listener);
        let https_future = self.accept_loop_tls(https_listener, acceptor.clone());

        tokio::select! {
            result = http_future => result,
            result = https_future => result,
        }
    }

    async fn accept_loop_tls(&self, listener: TcpListener, acceptor: TlsAcceptor) -> Result<()> {
        let mut shutdown_rx = self.shutdown_rx.clone();

        loop {
            tokio::select! {
                // Check for shutdown signal
                _ = shutdown_rx.changed() => {
                    if *shutdown_rx.borrow() {
                        info!("Shutting down HTTPS proxy server");
                        break;
                    }
                }

                // Accept new connections
                result = listener.accept() => {
                    match result {
                        Ok((stream, remote_addr)) => {
                            let registry = self.registry.clone();
                            let load_balancer = self.load_balancer.clone();
                            let config = self.config.clone();
                            let acceptor = acceptor.clone();
                            let cert_manager = self.cert_manager.clone();
                            let npc = self.network_policy_checker.clone();
                            let activator = self.activator.clone();
                            let rps_registry = self.rps_registry.clone();

                            tokio::spawn(async move {
                                if let Err(e) = Self::handle_tls_connection(
                                    stream,
                                    remote_addr,
                                    registry,
                                    load_balancer,
                                    config,
                                    acceptor,
                                    cert_manager,
                                    npc,
                                    activator,
                                    rps_registry,
                                ).await {
                                    debug!(
                                        error = %e,
                                        remote_addr = %remote_addr,
                                        "TLS connection error"
                                    );
                                }
                            });
                        }
                        Err(e) => {
                            warn!(error = %e, "Failed to accept TLS connection");
                        }
                    }
                }
            }
        }

        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    async fn handle_tls_connection(
        stream: tokio::net::TcpStream,
        remote_addr: SocketAddr,
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
        config: Arc<ProxyConfig>,
        acceptor: TlsAcceptor,
        cert_manager: Option<Arc<CertManager>>,
        network_policy_checker: Option<NetworkPolicyChecker>,
        activator: Option<Arc<dyn Activator>>,
        rps_registry: Option<Arc<RpsRegistry>>,
    ) -> Result<()> {
        // Intercept-but-passthrough: peek the ClientHello SNI WITHOUT consuming
        // the socket bytes, so we can decide whether to terminate TLS (managed)
        // or TCP-splice raw to the real upstream (unmanaged) before rustls ever
        // sees a cert-less SNI and aborts — which would otherwise hang the
        // client.
        let sni = Self::peek_sni(&stream).await;

        // A route existing for this host is the primary "we own this" signal.
        // When there's no SNI at all we treat it as managed and fall through to
        // the normal accept path (a default cert / catch-all route may apply,
        // and we must never passthrough to an unknown destination).
        let managed = match sni.as_deref() {
            Some(host) => registry.resolve(Some(host), "/").await.is_some(),
            None => true,
        };

        if !managed {
            // `managed` is only false when `sni` is `Some`.
            if let Some(host) = sni {
                return Self::passthrough_unmanaged(stream, remote_addr, host).await;
            }
        }

        // MANAGED path (unchanged): terminate TLS and reverse-proxy. The peeked
        // bytes are still queued in the socket, so rustls reads the full
        // ClientHello with no replay needed.
        let tls_stream = acceptor
            .accept(stream)
            .await
            .map_err(|e| ProxyError::Tls(format!("TLS handshake failed: {e}")))?;

        let io = TokioIo::new(tls_stream);

        let mut service = ReverseProxyService::new(registry, load_balancer, config)
            .with_remote_addr(remote_addr)
            .with_tls(true);
        if let Some(cm) = cert_manager {
            service = service.with_cert_manager(cm);
        }
        if let Some(checker) = network_policy_checker {
            service = service.with_network_policy_checker(checker);
        }
        if let Some(activator) = activator {
            service = service.with_activator(activator);
        }
        if let Some(rps_registry) = rps_registry {
            service = service.with_rps_registry(rps_registry);
        }

        let service = service_fn(move |req: Request<Incoming>| {
            let svc = service.clone();
            async move {
                match svc.proxy_request(req).await {
                    Ok(response) => Ok::<_, hyper::Error>(response),
                    Err(e) => {
                        error!(error = %e, "Proxy error");
                        Ok(ReverseProxyService::error_response(&e))
                    }
                }
            }
        });

        http1::Builder::new()
            .preserve_header_case(true)
            .title_case_headers(false)
            .serve_connection(io, service)
            .with_upgrades()
            .await
            .map_err(ProxyError::Hyper)?;

        Ok(())
    }

    /// Peek (without consuming) the initial bytes of `stream` and extract the
    /// SNI host from the TLS `ClientHello`.
    ///
    /// Uses [`TcpStream::peek`] so the bytes stay queued in the socket for the
    /// subsequent `acceptor.accept` (managed) or raw splice (unmanaged) — there
    /// is no ClientHello-replay problem. Bounded by a byte cap and a short
    /// timeout so a slowloris cannot stall the accept task forever. Returns
    /// `None` on close, timeout, or unparseable input.
    async fn peek_sni(stream: &TcpStream) -> Option<String> {
        /// Maximum bytes to peek for the `ClientHello`.
        const MAX_PEEK: usize = 8192;
        /// Overall budget to receive enough of the `ClientHello`.
        const PEEK_TIMEOUT: Duration = Duration::from_secs(5);

        let mut buf = vec![0u8; MAX_PEEK];
        let result = tokio::time::timeout(PEEK_TIMEOUT, async {
            loop {
                let n = match stream.peek(&mut buf).await {
                    // peer closed (Ok(0)) or errored before sending a ClientHello
                    Ok(n) if n > 0 => n,
                    _ => return None,
                };

                // Once the 5-byte record header is buffered, wait until the
                // whole record (capped at MAX_PEEK) is present before parsing.
                if n >= 5 {
                    let record_len = ((buf[3] as usize) << 8) | buf[4] as usize;
                    let needed = (5 + record_len).min(MAX_PEEK);
                    if n >= needed {
                        return crate::sni_peek::parse_sni(&buf[..n]);
                    }
                }
                if n >= MAX_PEEK {
                    return crate::sni_peek::parse_sni(&buf[..n]);
                }

                // Not enough yet — yield briefly so we don't busy-spin while the
                // rest of the ClientHello trickles in.
                tokio::time::sleep(Duration::from_millis(5)).await;
            }
        })
        .await;

        result.ok().flatten()
    }

    /// Splice an unmanaged-SNI connection straight to its real upstream on the
    /// standard HTTPS port (443), terminating no TLS.
    ///
    /// Refuses to connect to loopback / unspecified addresses to avoid a
    /// proxy→self loop. On any DNS / connect / loop-guard failure the
    /// connection is dropped promptly (returns `Ok(())` which closes it) — it
    /// must never hang.
    async fn passthrough_unmanaged(
        client: TcpStream,
        remote_addr: SocketAddr,
        host: String,
    ) -> Result<()> {
        let mut addrs = match lookup_host((host.as_str(), 443u16)).await {
            Ok(it) => it,
            Err(e) => {
                debug!(host = %host, error = %e, "TLS passthrough DNS lookup failed; dropping");
                return Ok(());
            }
        };

        // LOOP-GUARD: skip any loopback/unspecified IP so we never proxy to
        // ourselves; pick the first real upstream.
        let Some(upstream_addr) = addrs.find(|a| !Self::is_self_addr(&a.ip())) else {
            warn!(
                host = %host,
                "TLS passthrough refused: SNI resolves only to loopback/unspecified (loop guard)"
            );
            return Ok(());
        };

        let upstream = match TcpStream::connect(upstream_addr).await {
            Ok(s) => s,
            Err(e) => {
                debug!(
                    host = %host,
                    upstream = %upstream_addr,
                    error = %e,
                    "TLS passthrough connect failed; dropping"
                );
                return Ok(());
            }
        };

        debug!(
            host = %host,
            upstream = %upstream_addr,
            client = %remote_addr,
            "TLS passthrough (unmanaged SNI) -> upstream"
        );
        crate::stream::TcpStreamService::splice(client, upstream).await;
        Ok(())
    }

    /// True if `ip` would loop the proxy back to itself — loopback
    /// (`127.0.0.0/8`, `::1`) or unspecified (`0.0.0.0`, `::`).
    fn is_self_addr(ip: &IpAddr) -> bool {
        ip.is_loopback() || ip.is_unspecified()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lb::LoadBalancer;
    use crate::routes::{ResolvedService, RouteEntry};
    use zlayer_spec::{ExposeType, Protocol};

    /// Helper to build a minimal `RouteEntry` for tests.
    fn make_entry(
        service: &str,
        host: Option<&str>,
        path: &str,
        backends: Vec<SocketAddr>,
    ) -> RouteEntry {
        RouteEntry {
            service_name: service.to_string(),
            endpoint_name: "http".to_string(),
            host: host.map(std::string::ToString::to_string),
            path_prefix: path.to_string(),
            resolved: ResolvedService {
                name: service.to_string(),
                backends,
                use_tls: false,
                sni_hostname: String::new(),
                expose: ExposeType::Public,
                protocol: Protocol::Http,
                strip_prefix: false,
                path_prefix: path.to_string(),
                target_port: 8080,
            },
        }
    }

    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    /// Drive a single raw HTTP/1.1 request through `handle_connection` and
    /// return the raw response bytes (status line + headers + body).
    ///
    /// This exercises the FULL ingress path — route resolution, the
    /// default-deny on an unmatched route, the ACME carve-out, and the
    /// generic `error_response` body — exactly as a real client would hit it,
    /// without binding a privileged port.
    async fn roundtrip(
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
        cert_manager: Option<Arc<CertManager>>,
        raw_request: &str,
    ) -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();

        let server = tokio::spawn(async move {
            let (stream, remote_addr) = listener.accept().await.unwrap();
            let _ = ProxyServer::handle_connection(
                stream,
                remote_addr,
                registry,
                load_balancer,
                Arc::new(ProxyConfig::default()),
                cert_manager,
                None,
                None,
                None,
            )
            .await;
        });

        let mut client = tokio::net::TcpStream::connect(addr).await.unwrap();
        client.write_all(raw_request.as_bytes()).await.unwrap();
        client.flush().await.unwrap();

        let mut buf = Vec::new();
        // Read until the peer closes (Connection: close ensures EOF).
        let _ = client.read_to_end(&mut buf).await;
        server.abort();

        String::from_utf8_lossy(&buf).into_owned()
    }

    /// An unregistered Host (no matching route) must get a clean 404 with a
    /// generic body — NOT forwarded anywhere, and NOT echoing the requested
    /// Host/path. This is the default-deny boundary for the open ingress.
    #[tokio::test]
    async fn test_unmatched_host_denied_404_generic_body() {
        let registry = Arc::new(ServiceRegistry::new());
        // Register a route for ONE host so the registry is non-empty; the
        // request below targets a DIFFERENT host that matches nothing.
        registry
            .register(make_entry(
                "known",
                Some("known.example.com"),
                "/",
                vec!["127.0.0.1:9".parse().unwrap()],
            ))
            .await;
        let lb = Arc::new(LoadBalancer::new());

        let resp = roundtrip(
            registry,
            lb,
            None,
            "GET /secret/path HTTP/1.1\r\nHost: attacker.unregistered.test\r\nConnection: close\r\n\r\n",
        )
        .await;

        assert!(
            resp.starts_with("HTTP/1.1 404"),
            "unmatched host must be denied with 404, got: {resp}"
        );
        // The generic body must not leak the requested host or path.
        assert!(
            !resp.contains("attacker.unregistered.test"),
            "response must not echo the requested host: {resp}"
        );
        assert!(
            !resp.contains("/secret/path"),
            "response must not echo the requested path: {resp}"
        );
        assert!(
            resp.contains("404 Not Found"),
            "response should carry the generic 404 body: {resp}"
        );
    }

    /// The ACME HTTP-01 challenge path must still be served (this is how certs
    /// get issued) — the default-deny must come AFTER the ACME carve-out.
    #[tokio::test]
    async fn test_acme_challenge_served_not_denied() {
        let tmp = tempfile::tempdir().unwrap();
        let cm = Arc::new(
            CertManager::new(tmp.path().to_string_lossy().into_owned(), None)
                .await
                .unwrap(),
        );
        let token = "test-token-abc";
        cm.store_challenge(token, "example.com", "key-auth-payload-123");

        // Empty registry: nothing is registered, so if the ACME carve-out
        // were missing this would be denied by default.
        let registry = Arc::new(ServiceRegistry::new());
        let lb = Arc::new(LoadBalancer::new());

        let resp = roundtrip(
            registry,
            lb,
            Some(cm),
            &format!(
                "GET /.well-known/acme-challenge/{token} HTTP/1.1\r\nHost: example.com\r\nConnection: close\r\n\r\n"
            ),
        )
        .await;

        assert!(
            resp.starts_with("HTTP/1.1 200"),
            "ACME challenge must be served with 200, got: {resp}"
        );
        assert!(
            resp.contains("key-auth-payload-123"),
            "ACME challenge must return the key authorization: {resp}"
        );
    }

    /// An ACME challenge path with an UNKNOWN token must be handled terminally
    /// with a 404 — it must NOT fall through to vhost routing (which would
    /// return the default-deny 403/404 for the host). The acme-challenge
    /// carve-out owns the path entirely, even when no token is stored.
    #[tokio::test]
    async fn test_acme_challenge_unknown_token_404_not_forbidden() {
        let tmp = tempfile::tempdir().unwrap();
        let cm = Arc::new(
            CertManager::new(tmp.path().to_string_lossy().into_owned(), None)
                .await
                .unwrap(),
        );
        // No token is stored, so the lookup misses.

        // Register a route for ONE host so the registry is non-empty; the
        // request below targets a DIFFERENT host that matches nothing — proving
        // the acme path is handled terminally and never reaches the vhost
        // default-deny (which returns Forbidden/404 for unmatched hosts).
        let registry = Arc::new(ServiceRegistry::new());
        registry
            .register(make_entry(
                "known",
                Some("known.example.com"),
                "/",
                vec!["127.0.0.1:9".parse().unwrap()],
            ))
            .await;
        let lb = Arc::new(LoadBalancer::new());

        let resp = roundtrip(
            registry,
            lb,
            Some(cm),
            "GET /.well-known/acme-challenge/does-not-exist HTTP/1.1\r\nHost: console.zatabase.io\r\nConnection: close\r\n\r\n",
        )
        .await;

        assert!(
            resp.starts_with("HTTP/1.1 404"),
            "unknown ACME token must return 404, got: {resp}"
        );
        assert!(
            !resp.starts_with("HTTP/1.1 403"),
            "unknown ACME token must NOT fall through to vhost 403: {resp}"
        );
        assert!(
            resp.contains("ACME challenge token not found"),
            "404 should carry the ACME-specific body: {resp}"
        );
    }

    /// A matched route whose backend group has zero healthy backends must
    /// return 503 with a generic body — NOT leaking the internal LB group
    /// name.
    #[tokio::test]
    async fn test_matched_no_backends_503_generic_body() {
        let registry = Arc::new(ServiceRegistry::new());
        // Route matches, but its resolved LB group name (carried in the body
        // of the old leaky error) must never reach the client.
        let lb_group = "prod/api#http-secret-group";
        let mut entry = make_entry("api", Some("api.example.com"), "/", vec![]);
        entry.resolved.name = lb_group.to_string();
        registry.register(entry).await;

        // Load balancer has NO group registered for this name → select()
        // returns None → NoHealthyBackends.
        let lb = Arc::new(LoadBalancer::new());

        let resp = roundtrip(
            registry,
            lb,
            None,
            "GET / HTTP/1.1\r\nHost: api.example.com\r\nConnection: close\r\n\r\n",
        )
        .await;

        assert!(
            resp.starts_with("HTTP/1.1 503"),
            "matched route with no healthy backends must return 503, got: {resp}"
        );
        assert!(
            !resp.contains(lb_group),
            "503 body must not leak the internal LB group name: {resp}"
        );
        assert!(
            resp.contains("503 Service Unavailable"),
            "response should carry the generic 503 body: {resp}"
        );
    }

    #[tokio::test]
    async fn test_server_shutdown() {
        let registry = Arc::new(ServiceRegistry::new());
        let lb = Arc::new(LoadBalancer::new());
        let server = ProxyServer::new(ProxyConfig::default(), registry, lb);

        // Create a separate handle for shutdown
        let shutdown_tx = server.shutdown_tx.clone();

        // Signal shutdown immediately
        let _ = shutdown_tx.send(true);

        // Server should exit gracefully
        // (In a real test, we'd spawn the server and verify it stops)
    }

    #[tokio::test]
    async fn test_registry_integration() {
        let registry = Arc::new(ServiceRegistry::new());

        // Add a route
        registry
            .register(make_entry(
                "test-service",
                None,
                "/api",
                vec!["127.0.0.1:8081".parse().unwrap()],
            ))
            .await;

        let lb = Arc::new(LoadBalancer::new());
        let server = ProxyServer::new(ProxyConfig::default(), registry, lb);

        // Verify registry is accessible
        let reg = server.registry();
        assert_eq!(reg.route_count().await, 1);
    }

    #[test]
    fn test_next_ingress_backoff_doubles_then_caps() {
        // Doubles from the initial value.
        assert_eq!(
            next_ingress_backoff(INGRESS_BIND_BACKOFF_INITIAL),
            INGRESS_BIND_BACKOFF_INITIAL * 2
        );
        // Keeps doubling until it reaches the cap.
        let mut d = INGRESS_BIND_BACKOFF_INITIAL;
        for _ in 0..20 {
            d = next_ingress_backoff(d);
        }
        assert_eq!(d, INGRESS_BIND_BACKOFF_MAX);
        // Never exceeds the cap once there.
        assert_eq!(
            next_ingress_backoff(INGRESS_BIND_BACKOFF_MAX),
            INGRESS_BIND_BACKOFF_MAX
        );
    }

    #[test]
    fn test_should_warn_cadence() {
        // Always warns on the first attempt.
        assert!(should_warn_on_attempt(0));
        // Quiet in between.
        assert!(!should_warn_on_attempt(1));
        assert!(!should_warn_on_attempt(INGRESS_BIND_WARN_EVERY - 1));
        // Warns again every N attempts.
        assert!(should_warn_on_attempt(INGRESS_BIND_WARN_EVERY));
        assert!(should_warn_on_attempt(INGRESS_BIND_WARN_EVERY * 2));
    }

    #[tokio::test]
    async fn test_bind_with_retry_succeeds_after_initial_conflict() {
        // Hold an ephemeral port, then spawn a retrying bind against it; once
        // we drop our listener the retry loop must grab the port and return.
        let held = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = held.local_addr().unwrap();

        let (_tx, mut rx) = watch::channel(false);
        let handle = tokio::spawn(async move { bind_with_retry(addr, "test", &mut rx).await });

        // Give the retry loop a chance to observe the conflict at least once.
        tokio::time::sleep(Duration::from_millis(50)).await;
        // Free the port; the retry loop should now bind it.
        drop(held);

        let bound = tokio::time::timeout(Duration::from_secs(10), handle)
            .await
            .expect("bind_with_retry did not finish")
            .expect("task panicked");
        let listener = bound.expect("expected a bound listener, got None");
        assert_eq!(listener.local_addr().unwrap().port(), addr.port());
    }

    #[tokio::test]
    async fn test_bind_with_retry_returns_none_on_shutdown() {
        // A port held for the whole test: the retry loop can never bind, so a
        // shutdown signal must unblock it and yield None (never errors).
        let held = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = held.local_addr().unwrap();

        let (tx, mut rx) = watch::channel(false);
        let handle = tokio::spawn(async move { bind_with_retry(addr, "test", &mut rx).await });

        tokio::time::sleep(Duration::from_millis(50)).await;
        tx.send(true).unwrap();

        let result = tokio::time::timeout(Duration::from_secs(10), handle)
            .await
            .expect("bind_with_retry did not respond to shutdown")
            .expect("task panicked");
        assert!(result.is_none(), "shutdown should yield None");
    }
}