zlayer_proxy/

service.rs

//! Reverse proxy service implementation
//!
//! This module provides the core proxy service that handles request forwarding.
//! It uses the `ServiceRegistry` for route resolution and backend selection.

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::{transform_path, ResolvedService, ServiceRegistry};
use bytes::Bytes;
use http::{header, Request, Response, Uri, Version};
use http_body_util::{BodyExt, Full};
use hyper::body::Incoming;
use hyper::upgrade::OnUpgrade;
use hyper_util::client::legacy::Client;
use hyper_util::rt::{TokioExecutor, TokioIo};
use std::collections::VecDeque;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::{Duration, Instant};
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tower::Service;
use tracing::{debug, error, info, warn};
use zlayer_spec::ExposeType;

/// Default ceiling for how long [`ReverseProxyService::proxy_request`] will hold
/// a request while it waits for an [`Activator`] to bring a scaled-to-zero
/// service back up before falling back to the existing `503`.
const ACTIVATE_DEADLINE: Duration = Duration::from_secs(30);

/// Polling step used while waiting for a backend to become available after an
/// [`Activator::activate`] call. Small enough to feel responsive, large enough
/// not to spin the load balancer.
const ACTIVATE_POLL_STEP: Duration = Duration::from_millis(200);

/// Width of the sliding window over which [`RpsRegistry`] computes
/// requests-per-second.
const RPS_WINDOW: Duration = Duration::from_secs(10);

/// On-demand activation hook for scale-to-zero services.
///
/// When the proxy resolves a route whose backend group currently has **no
/// healthy backends** (the scale-to-zero idle state), it calls
/// [`Activator::activate`] with the resolved load-balancer group name. The
/// implementation is expected to trigger a scale-up (e.g. scale the service to
/// its activation floor) and return once it has *initiated* the scale; the
/// proxy then re-polls backend selection on a bounded backoff loop, forwarding
/// the held request the moment a healthy backend appears.
///
/// Returning `Err` is non-fatal: the proxy logs it and falls through to the
/// existing no-healthy-backends `503` path, so a flaky activator never blocks
/// the request indefinitely beyond the proxy's own deadline.
#[async_trait::async_trait]
pub trait Activator: Send + Sync {
    /// Trigger activation (scale-up) for the load-balancer group `service`.
    ///
    /// `service` is the resolved LB group name (the same string passed to
    /// [`LoadBalancer::select`](crate::lb::LoadBalancer::select)), NOT
    /// necessarily the bare service name. Implementations that need the bare
    /// service name should derive it from this key.
    ///
    /// # Errors
    ///
    /// Returns a human-readable error string if activation could not be
    /// initiated. The proxy treats this as non-fatal and falls back to `503`.
    async fn activate(&self, service: &str) -> std::result::Result<(), String>;
}

/// Per-service sliding-window request-rate counter.
///
/// Records one timestamp per successfully-routed request and reports the
/// requests-per-second rate over a fixed [`RPS_WINDOW`]. Cheap to clone (it is
/// an `Arc`-friendly wrapper around interior-mutable state) and safe to share
/// across the proxy's per-connection service clones.
///
/// This is the real per-service RPS signal the autoscaler consumes: the proxy
/// records every routed request via [`RpsRegistry::record`], and the scheduler
/// reads [`RpsRegistry::rps`] for the same service key to drive request-rate
/// scaling.
#[derive(Debug, Default)]
pub struct RpsRegistry {
    /// Per-service ring of recent request timestamps, pruned to [`RPS_WINDOW`].
    services: Mutex<std::collections::HashMap<String, VecDeque<Instant>>>,
}

impl RpsRegistry {
    /// Create an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Record a single request against `service` at the current instant,
    /// pruning timestamps older than the window.
    pub async fn record(&self, service: &str) {
        let now = Instant::now();
        let cutoff = now.checked_sub(RPS_WINDOW).unwrap_or(now);
        let mut map = self.services.lock().await;
        let ring = map.entry(service.to_string()).or_default();
        ring.push_back(now);
        while ring.front().is_some_and(|t| *t < cutoff) {
            ring.pop_front();
        }
    }

    /// Current requests-per-second for `service`, averaged over the sliding
    /// window. Returns `0.0` for an unknown or idle service.
    pub async fn rps(&self, service: &str) -> f64 {
        let now = Instant::now();
        let cutoff = now.checked_sub(RPS_WINDOW).unwrap_or(now);
        let mut map = self.services.lock().await;
        let Some(ring) = map.get_mut(service) else {
            return 0.0;
        };
        while ring.front().is_some_and(|t| *t < cutoff) {
            ring.pop_front();
        }
        let count = ring.len();
        #[allow(clippy::cast_precision_loss)]
        {
            count as f64 / RPS_WINDOW.as_secs_f64()
        }
    }

    /// Snapshot of the current per-service RPS for every service seen within
    /// the window. Services whose window has fully drained report `0.0`.
    pub async fn snapshot(&self) -> std::collections::HashMap<String, f64> {
        let now = Instant::now();
        let cutoff = now.checked_sub(RPS_WINDOW).unwrap_or(now);
        let window_secs = RPS_WINDOW.as_secs_f64();
        let mut map = self.services.lock().await;
        let mut out = std::collections::HashMap::with_capacity(map.len());
        for (name, ring) in map.iter_mut() {
            while ring.front().is_some_and(|t| *t < cutoff) {
                ring.pop_front();
            }
            #[allow(clippy::cast_precision_loss)]
            let rps = ring.len() as f64 / window_secs;
            out.insert(name.clone(), rps);
        }
        out
    }
}

/// The overlay network CIDR used for internal service communication.
/// Source IPs outside this range are rejected for internal-only routes.
const OVERLAY_NETWORK: (u8, u8) = (10, 200); // 10.200.0.0/16

/// Check whether an IP address belongs to the overlay network (10.200.0.0/16).
fn is_overlay_ip(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            let octets = v4.octets();
            octets[0] == OVERLAY_NETWORK.0 && octets[1] == OVERLAY_NETWORK.1
        }
        IpAddr::V6(_) => false,
    }
}

/// Body type for outgoing responses
pub type BoxBody = http_body_util::combinators::BoxBody<Bytes, hyper::Error>;

/// Empty body utility
#[must_use]
pub fn empty_body() -> BoxBody {
    http_body_util::Empty::<Bytes>::new()
        .map_err(|never| match never {})
        .boxed()
}

/// Full body utility
pub fn full_body(bytes: impl Into<Bytes>) -> BoxBody {
    Full::new(bytes.into())
        .map_err(|never| match never {})
        .boxed()
}

/// The reverse proxy service
#[derive(Clone)]
pub struct ReverseProxyService {
    /// Service registry for route resolution
    registry: Arc<ServiceRegistry>,
    /// Load balancer for backend selection
    load_balancer: Arc<LoadBalancer>,
    /// HTTP client for backend requests
    client: Client<hyper_util::client::legacy::connect::HttpConnector, BoxBody>,
    /// Proxy configuration
    config: Arc<ProxyConfig>,
    /// Client remote address (set per-request)
    remote_addr: Option<SocketAddr>,
    /// Whether the connection is over TLS
    is_tls: bool,
    /// Certificate manager for ACME challenge responses
    cert_manager: Option<Arc<CertManager>>,
    /// Optional network policy checker for access control enforcement
    network_policy_checker: Option<NetworkPolicyChecker>,
    /// Trusted upstream proxies. Requests whose TCP peer IP is in this list
    /// may set `CF-Connecting-IP` / `X-Forwarded-For` and be believed. When no
    /// explicit list is provided, defaults to `TrustedProxyList::localhost_only()`
    /// — a safe default for nodes that accidentally receive direct requests.
    trusted_proxies: Arc<crate::trust::TrustedProxyList>,
    /// Optional on-demand activator for scale-to-zero services. When set and a
    /// resolved route has no healthy backend, the request is held while the
    /// activator scales the service up (see [`Activator`]).
    activator: Option<Arc<dyn Activator>>,
    /// Optional per-service request-rate counter. When set, every
    /// successfully-routed request is recorded so the autoscaler can read a
    /// real RPS signal (see [`RpsRegistry`]).
    rps_registry: Option<Arc<RpsRegistry>>,
}

impl ReverseProxyService {
    /// Create a new reverse proxy service
    pub fn new(
        registry: Arc<ServiceRegistry>,
        load_balancer: Arc<LoadBalancer>,
        config: Arc<ProxyConfig>,
    ) -> Self {
        let client = Client::builder(TokioExecutor::new())
            .pool_max_idle_per_host(config.pool.max_idle_per_backend)
            .pool_idle_timeout(config.pool.idle_timeout)
            .pool_timer(hyper_util::rt::TokioTimer::new())
            .build_http();

        Self {
            registry,
            load_balancer,
            client,
            config,
            remote_addr: None,
            is_tls: false,
            cert_manager: None,
            network_policy_checker: None,
            trusted_proxies: Arc::new(crate::trust::TrustedProxyList::localhost_only()),
            activator: None,
            rps_registry: None,
        }
    }

    /// Set the remote client address for this request
    #[must_use]
    pub fn with_remote_addr(mut self, addr: SocketAddr) -> Self {
        self.remote_addr = Some(addr);
        self
    }

    /// Mark this connection as being over TLS
    #[must_use]
    pub fn with_tls(mut self, is_tls: bool) -> Self {
        self.is_tls = is_tls;
        self
    }

    /// Override the trusted-proxy list (default: `localhost_only`).
    ///
    /// Peers in this list are believed when they set `CF-Connecting-IP` or
    /// `X-Forwarded-For` headers identifying the real client IP.
    #[must_use]
    pub fn with_trusted_proxies(mut self, trusted: Arc<crate::trust::TrustedProxyList>) -> Self {
        self.trusted_proxies = trusted;
        self
    }

    /// 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.
    ///
    /// With an activator installed, a request to a resolved route whose backend
    /// group has no healthy backend triggers [`Activator::activate`] and is held
    /// (up to [`ACTIVATE_DEADLINE`]) until a backend appears, instead of
    /// immediately returning `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.
    ///
    /// With a registry installed, every successfully-routed request is recorded
    /// so the autoscaler can read a real per-service RPS signal.
    #[must_use]
    pub fn with_rps_registry(mut self, rps_registry: Arc<RpsRegistry>) -> Self {
        self.rps_registry = Some(rps_registry);
        self
    }

    /// Check if this connection is over TLS
    #[must_use]
    pub fn is_tls(&self) -> bool {
        self.is_tls
    }

    /// Handle an incoming HTTP request
    ///
    /// # Errors
    ///
    /// Returns an error if route resolution fails, no healthy backends are
    /// available, or the backend request fails.
    ///
    /// # Panics
    ///
    /// Panics if building a well-formed HTTP response for an ACME challenge
    /// or upgrade reply fails (indicates a bug, not a runtime condition).
    #[allow(clippy::too_many_lines)]
    pub async fn proxy_request(&self, mut req: Request<Incoming>) -> Result<Response<BoxBody>> {
        let start = std::time::Instant::now();
        let method = req.method().clone();
        let uri = req.uri().clone();

        let host = req
            .headers()
            .get(header::HOST)
            .and_then(|h| h.to_str().ok())
            .or_else(|| uri.host())
            .map(std::string::ToString::to_string);

        let path = uri.path().to_string();

        // ACME HTTP-01 challenge interception. This is TERMINAL: any request
        // whose path is under /.well-known/acme-challenge/ is fully handled
        // here and never falls through to vhost routing (which would return a
        // confusing 403 Forbidden for an HTTPS-only host). A stored token
        // returns 200 with the key authorization; an unknown/expired/empty
        // token (or absent cert manager) returns a clean 404.
        if let Some(token) = path.strip_prefix("/.well-known/acme-challenge/") {
            if !token.is_empty() {
                if let Some(ref cm) = self.cert_manager {
                    if let Some(auth) = cm.get_challenge_response(token) {
                        return Ok(Response::builder()
                            .status(200)
                            .header("content-type", "text/plain")
                            .body(full_body(auth))
                            .unwrap());
                    }
                }
            }
            tracing::warn!(
                token = %token,
                cert_manager = self.cert_manager.is_some(),
                host = host.as_deref().unwrap_or("<none>"),
                "ACME HTTP-01 challenge token not found; returning 404"
            );
            return Ok(Response::builder()
                .status(404)
                .header("content-type", "text/plain")
                .body(full_body("ACME challenge token not found"))
                .unwrap());
        }

        // Check for WebSocket/HTTP upgrade
        if crate::tunnel::is_upgrade_request(&req) {
            // Resolve to get backend for upgrade
            let resolved = self
                .registry
                .resolve(host.as_deref(), &path)
                .await
                .ok_or_else(|| ProxyError::RouteNotFound {
                    host: host.as_deref().unwrap_or("<none>").to_string(),
                    path: path.clone(),
                })?;

            // Enforce internal endpoints
            if resolved.expose == ExposeType::Internal {
                if let Some(addr) = self.remote_addr {
                    if !is_overlay_ip(addr.ip()) {
                        return Err(ProxyError::Forbidden(
                            "endpoint is internal-only".to_string(),
                        ));
                    }
                }
            }

            // Enforce network policy access rules
            if let (Some(checker), Some(addr)) = (&self.network_policy_checker, self.remote_addr) {
                if !checker
                    .check_access(addr.ip(), &resolved.name, "*", resolved.target_port)
                    .await
                {
                    return Err(ProxyError::Forbidden(format!(
                        "network policy denied access to service '{}'",
                        resolved.name
                    )));
                }
            }

            let backend = self
                .select_or_activate(&resolved.name)
                .await
                .ok_or_else(|| ProxyError::NoHealthyBackends {
                    service: resolved.name.clone(),
                })?;
            let _guard = backend.track_connection();
            let backend_addr = backend.addr;

            // Record the routed request for per-service RPS metrics.
            if let Some(rps) = &self.rps_registry {
                rps.record(&resolved.name).await;
            }

            info!(
                method = %method,
                host = ?host,
                path = %path,
                backend = %backend_addr,
                service = %resolved.name,
                "Forwarding upgrade request"
            );

            // Extract the client's OnUpgrade future BEFORE consuming the request
            let client_upgrade: OnUpgrade = hyper::upgrade::on(&mut req);

            // Build the backend URI
            let original_path = req.uri().path();
            let transformed_path =
                transform_path(&resolved.path_prefix, original_path, resolved.strip_prefix);
            let new_uri = format!(
                "http://{}{}{}",
                backend_addr,
                transformed_path,
                req.uri()
                    .query()
                    .map(|q| format!("?{q}"))
                    .unwrap_or_default()
            );

            // Build backend request, preserving upgrade headers
            let (orig_parts, _body) = req.into_parts();
            let mut backend_parts = http::request::Builder::new()
                .method(orig_parts.method.clone())
                .uri(
                    new_uri
                        .parse::<Uri>()
                        .map_err(|e| ProxyError::InvalidRequest(format!("Invalid URI: {e}")))?,
                )
                .body(())
                .unwrap()
                .into_parts()
                .0;

            // Copy all original headers first (preserving Host, etc.)
            for (name, value) in &orig_parts.headers {
                backend_parts.headers.insert(name.clone(), value.clone());
            }

            // Copy upgrade-specific headers (Connection, Upgrade, Sec-WebSocket-*)
            crate::tunnel::copy_upgrade_headers(&orig_parts, &mut backend_parts);

            // Add forwarding headers
            self.add_forwarding_headers(&mut backend_parts);

            // Connect directly to backend (bypass connection pool for long-lived upgrades)
            let tcp_stream = TcpStream::connect(backend_addr).await.map_err(|e| {
                error!(error = %e, backend = %backend_addr, "Backend upgrade connect failed");
                ProxyError::BackendConnectionFailed {
                    backend: backend_addr,
                    reason: e.to_string(),
                }
            })?;
            let io = TokioIo::new(tcp_stream);

            // Perform HTTP/1.1 handshake preserving header case
            let (mut sender, conn) = hyper::client::conn::http1::Builder::new()
                .preserve_header_case(true)
                .handshake(io)
                .await
                .map_err(|e| {
                    error!(error = %e, backend = %backend_addr, "Backend upgrade handshake failed");
                    ProxyError::BackendRequestFailed(format!("Upgrade handshake failed: {e}"))
                })?;

            // Spawn the connection driver
            tokio::spawn(async move {
                if let Err(e) = conn.with_upgrades().await {
                    error!(error = %e, "Backend upgrade connection driver error");
                }
            });

            // Send the request to the backend
            let backend_req =
                Request::from_parts(backend_parts, http_body_util::Empty::<Bytes>::new());
            let backend_response = sender.send_request(backend_req).await.map_err(|e| {
                error!(error = %e, backend = %backend_addr, "Backend upgrade request failed");
                ProxyError::BackendRequestFailed(e.to_string())
            })?;

            if backend_response.status() == http::StatusCode::SWITCHING_PROTOCOLS {
                // Get the server's OnUpgrade future
                let server_upgrade: OnUpgrade = hyper::upgrade::on(backend_response);

                // Build 101 response to send back to the client
                let mut resp_builder =
                    Response::builder().status(http::StatusCode::SWITCHING_PROTOCOLS);
                // Note: we need to construct the response manually since we consumed
                // the backend response to get OnUpgrade. Copy relevant headers.
                // The hyper::upgrade::on() for the response does NOT consume it —
                // it was consumed. We need to return a 101 with appropriate headers.
                // Actually, hyper::upgrade::on() takes the response by value, so we
                // must build our own 101 response for the client.

                // For the client response, set Connection: upgrade and Upgrade headers
                if let Some(upgrade_val) = orig_parts.headers.get(header::UPGRADE) {
                    resp_builder = resp_builder.header(header::UPGRADE, upgrade_val.clone());
                }
                resp_builder = resp_builder.header(header::CONNECTION, "upgrade");

                let client_response = resp_builder.body(empty_body()).map_err(|e| {
                    ProxyError::Internal(format!("Failed to build 101 response: {e}"))
                })?;

                // Spawn background task to bridge the upgraded connections
                tokio::spawn(async move {
                    if let Err(e) =
                        crate::tunnel::proxy_upgrade(client_upgrade, server_upgrade).await
                    {
                        debug!(error = %e, "Upgrade tunnel ended");
                    }
                });

                // Add timing header to the 101 response
                let (mut parts, body) = client_response.into_parts();
                if let Ok(hv) = format!("proxy;dur={}", start.elapsed().as_millis()).parse() {
                    parts.headers.insert("server-timing", hv);
                }

                return Ok(Response::from_parts(parts, body));
            }

            // Backend didn't upgrade — stream the response as-is
            let (mut parts, body) = backend_response.into_parts();
            let streaming_body: BoxBody = body.map_err(|e: hyper::Error| e).boxed();

            // Add HSTS header for TLS connections
            if self.is_tls && self.config.headers.hsts {
                let value = if self.config.headers.hsts_subdomains {
                    format!(
                        "max-age={}; includeSubDomains",
                        self.config.headers.hsts_max_age
                    )
                } else {
                    format!("max-age={}", self.config.headers.hsts_max_age)
                };
                if let Ok(hv) = value.parse() {
                    parts.headers.insert("strict-transport-security", hv);
                }
            }

            // Add Server-Timing header
            if let Ok(hv) = format!("proxy;dur={}", start.elapsed().as_millis()).parse() {
                parts.headers.insert("server-timing", hv);
            }

            return Ok(Response::from_parts(parts, streaming_body));
        }

        debug!(method = %method, host = ?host, path = %path, "Routing request");

        // Resolve route
        let resolved = self
            .registry
            .resolve(host.as_deref(), &path)
            .await
            .ok_or_else(|| ProxyError::RouteNotFound {
                host: host.as_deref().unwrap_or("<none>").to_string(),
                path: path.clone(),
            })?;

        // Enforce internal endpoints
        if resolved.expose == ExposeType::Internal {
            match self.remote_addr {
                Some(addr) if !is_overlay_ip(addr.ip()) => {
                    warn!(
                        source = %addr.ip(),
                        service = %resolved.name,
                        "Rejected non-overlay source for internal endpoint"
                    );
                    return Err(ProxyError::Forbidden(
                        "endpoint is internal-only".to_string(),
                    ));
                }
                None => {
                    debug!(
                        service = %resolved.name,
                        "No remote_addr available; skipping overlay source check"
                    );
                }
                _ => {}
            }
        }

        // Enforce network policy access rules
        if let (Some(checker), Some(addr)) = (&self.network_policy_checker, self.remote_addr) {
            if !checker
                .check_access(addr.ip(), &resolved.name, "*", resolved.target_port)
                .await
            {
                return Err(ProxyError::Forbidden(format!(
                    "network policy denied access to service '{}'",
                    resolved.name
                )));
            }
        }

        // Select backend via load balancer, activating a scaled-to-zero
        // service on demand if needed (and an activator is installed).
        let backend = self
            .select_or_activate(&resolved.name)
            .await
            .ok_or_else(|| ProxyError::NoHealthyBackends {
                service: resolved.name.clone(),
            })?;
        let _guard = backend.track_connection();
        let backend_addr = backend.addr;

        // Record the routed request for per-service RPS metrics.
        if let Some(rps) = &self.rps_registry {
            rps.record(&resolved.name).await;
        }

        info!(
            method = %method,
            host = ?host,
            path = %path,
            backend = %backend_addr,
            service = %resolved.name,
            "Forwarding request"
        );

        // Build forwarded request
        let forwarded_req = self.build_forwarded_request(req, &backend_addr, &resolved)?;

        // Forward to backend
        let response = self.client.request(forwarded_req).await.map_err(|e| {
            error!(error = %e, backend = %backend_addr, "Backend request failed");
            ProxyError::BackendRequestFailed(e.to_string())
        })?;

        let (mut parts, body) = response.into_parts();
        let streaming_body: BoxBody = body.map_err(|e: hyper::Error| e).boxed();

        // Add HSTS header for TLS connections
        if self.is_tls && self.config.headers.hsts {
            let value = if self.config.headers.hsts_subdomains {
                format!(
                    "max-age={}; includeSubDomains",
                    self.config.headers.hsts_max_age
                )
            } else {
                format!("max-age={}", self.config.headers.hsts_max_age)
            };
            if let Ok(hv) = value.parse() {
                parts.headers.insert("strict-transport-security", hv);
            }
        }

        // Add Server-Timing header
        if let Ok(hv) = format!("proxy;dur={}", start.elapsed().as_millis()).parse() {
            parts.headers.insert("server-timing", hv);
        }

        Ok(Response::from_parts(parts, streaming_body))
    }

    /// Select a healthy backend for `service`, activating a scaled-to-zero
    /// service on demand if no backend is available and an [`Activator`] is
    /// installed.
    ///
    /// Behavior:
    /// - If [`LoadBalancer::select`](crate::lb::LoadBalancer::select) returns a
    ///   backend, it is returned immediately (the common path; zero added cost).
    /// - Otherwise, with no activator installed, `None` is returned at once so
    ///   the caller's existing `503` path is preserved unchanged.
    /// - With an activator installed, [`Activator::activate`] is called once and
    ///   then backend selection is re-polled every [`ACTIVATE_POLL_STEP`] until a
    ///   backend appears or [`ACTIVATE_DEADLINE`] elapses, at which point `None`
    ///   is returned and the caller falls back to `503`.
    ///
    /// The caller still holds the (un-consumed) request body across this await,
    /// so a successful activation forwards the original request normally.
    async fn select_or_activate(&self, service: &str) -> Option<Arc<crate::lb::Backend>> {
        if let Some(backend) = self.load_balancer.select(service) {
            return Some(backend);
        }

        let Some(activator) = &self.activator else {
            return None;
        };

        info!(
            service = %service,
            "No healthy backend; invoking activator (scale-to-zero wake-up)"
        );
        if let Err(e) = activator.activate(service).await {
            // Non-fatal: log and fall through to the bounded re-poll. The
            // service may still be coming up from a concurrent activation.
            warn!(service = %service, error = %e, "Activator returned an error; will still poll for a backend");
        }

        let deadline = Instant::now() + ACTIVATE_DEADLINE;
        loop {
            if let Some(backend) = self.load_balancer.select(service) {
                info!(service = %service, "Backend became available after activation");
                return Some(backend);
            }
            if Instant::now() >= deadline {
                warn!(
                    service = %service,
                    "Activation deadline elapsed without a healthy backend; falling back to 503"
                );
                return None;
            }
            tokio::time::sleep(ACTIVATE_POLL_STEP).await;
        }
    }

    fn build_forwarded_request(
        &self,
        req: Request<Incoming>,
        backend: &SocketAddr,
        resolved: &ResolvedService,
    ) -> Result<Request<BoxBody>> {
        let (mut parts, body) = req.into_parts();

        // Transform the path if needed
        let original_path = parts.uri.path();
        let transformed_path =
            transform_path(&resolved.path_prefix, original_path, resolved.strip_prefix);

        // Build new URI for backend
        let new_uri = format!(
            "http://{}{}{}",
            backend,
            transformed_path,
            parts
                .uri
                .query()
                .map(|q| format!("?{q}"))
                .unwrap_or_default()
        );

        parts.uri = new_uri
            .parse::<Uri>()
            .map_err(|e| ProxyError::InvalidRequest(format!("Invalid URI: {e}")))?;

        // Add forwarding headers
        self.add_forwarding_headers(&mut parts);

        // Remove hop-by-hop headers
        Self::remove_hop_by_hop_headers(&mut parts);

        let streaming_body: BoxBody = body.map_err(|e: hyper::Error| e).boxed();

        let req = Request::from_parts(parts, streaming_body);
        Ok(req)
    }

    fn add_forwarding_headers(&self, parts: &mut http::request::Parts) {
        let config = &self.config.headers;

        // Determine whether the immediate TCP peer is a trusted upstream proxy
        // that may dictate the real client IP via CF-Connecting-IP or XFF.
        let peer_is_trusted = self
            .remote_addr
            .is_some_and(|addr| self.trusted_proxies.is_trusted(addr.ip()));

        // Compute the effective client IP:
        //   - Trusted peer + CF-Connecting-IP (parseable) -> use CF header
        //   - Trusted peer + leftmost X-Forwarded-For (parseable) -> use XFF
        //   - Otherwise -> fall back to the TCP peer IP
        let effective_client_ip: Option<IpAddr> = if peer_is_trusted {
            let cf_ip = parts
                .headers
                .get("cf-connecting-ip")
                .and_then(|h| h.to_str().ok())
                .and_then(|s| s.trim().parse::<IpAddr>().ok());

            let xff_leftmost = parts
                .headers
                .get("x-forwarded-for")
                .and_then(|h| h.to_str().ok())
                .and_then(|s| s.split(',').next())
                .and_then(|s| s.trim().parse::<IpAddr>().ok());

            cf_ip
                .or(xff_leftmost)
                .or_else(|| self.remote_addr.map(|a| a.ip()))
        } else {
            self.remote_addr.map(|a| a.ip())
        };

        // X-Forwarded-For
        if config.x_forwarded_for {
            if let Some(addr) = self.remote_addr {
                let existing_xff = parts
                    .headers
                    .get("x-forwarded-for")
                    .and_then(|h| h.to_str().ok())
                    .map(std::string::ToString::to_string);

                let new_value = if peer_is_trusted {
                    // Trusted proxy: prepend the real client IP (from CF /
                    // leftmost XFF / peer) to any existing chain so downstream
                    // sees [real_client, ...upstream_chain].
                    let real = effective_client_ip.unwrap_or_else(|| addr.ip()).to_string();
                    match existing_xff {
                        Some(chain) if !chain.trim().is_empty() => format!("{real}, {chain}"),
                        _ => real,
                    }
                } else {
                    // Untrusted peer: preserve existing behavior — append the
                    // peer IP to any existing chain.
                    match existing_xff {
                        Some(chain) => format!("{}, {}", chain, addr.ip()),
                        None => addr.ip().to_string(),
                    }
                };

                if let Ok(value) = new_value.parse() {
                    parts.headers.insert("x-forwarded-for", value);
                }
            }
        }

        // X-Forwarded-Proto
        if config.x_forwarded_proto && parts.headers.get("x-forwarded-proto").is_none() {
            let proto = if self.is_tls { "https" } else { "http" };
            if let Ok(value) = proto.parse() {
                parts.headers.insert("x-forwarded-proto", value);
            }
        }

        // X-Forwarded-Host
        if config.x_forwarded_host {
            if let Some(host) = parts.headers.get(header::HOST).cloned() {
                if parts.headers.get("x-forwarded-host").is_none() {
                    parts.headers.insert("x-forwarded-host", host);
                }
            }
        }

        // X-Real-IP — set to the effective client IP only if the header is
        // currently absent (conservative: do not overwrite a value set by an
        // upstream component).
        if config.x_real_ip {
            if let Some(ip) = effective_client_ip {
                if parts.headers.get("x-real-ip").is_none() {
                    if let Ok(value) = ip.to_string().parse() {
                        parts.headers.insert("x-real-ip", value);
                    }
                }
            }
        }

        // Via header
        if config.via {
            let proto_version = match parts.version {
                Version::HTTP_09 => "0.9",
                Version::HTTP_10 => "1.0",
                Version::HTTP_2 => "2.0",
                Version::HTTP_3 => "3.0",
                _ => "1.1",
            };

            let via_value = format!("{} {}", proto_version, config.server_name);
            let existing = parts
                .headers
                .get(header::VIA)
                .and_then(|h| h.to_str().ok())
                .map(|s| format!("{s}, {via_value}"))
                .unwrap_or(via_value);

            if let Ok(value) = existing.parse() {
                parts.headers.insert(header::VIA, value);
            }
        }
    }

    fn remove_hop_by_hop_headers(parts: &mut http::request::Parts) {
        // Standard hop-by-hop headers that should not be forwarded
        const HOP_BY_HOP: &[&str] = &[
            "connection",
            "keep-alive",
            "proxy-authenticate",
            "proxy-authorization",
            "te",
            "trailer",
            "transfer-encoding",
            "upgrade",
        ];

        // First, collect headers listed in the Connection header before we remove it
        let connection_headers: Vec<String> = parts
            .headers
            .get(header::CONNECTION)
            .and_then(|h| h.to_str().ok())
            .map(|value| value.split(',').map(|s| s.trim().to_lowercase()).collect())
            .unwrap_or_default();

        for header_name in HOP_BY_HOP {
            parts.headers.remove(*header_name);
        }

        // Also remove headers that were listed in the Connection header
        for header_name in connection_headers {
            parts.headers.remove(header_name.as_str());
        }
    }

    /// Build a client-facing error response with a **generic** body.
    ///
    /// This is the default-deny safety boundary for the ingress proxy. The
    /// proxy binds `0.0.0.0:80`/`:443`, so it MUST NOT leak internal details
    /// (the requested Host/path, a backend address, or the internal
    /// load-balancer group name) to an unauthenticated caller. The full,
    /// detailed [`ProxyError`] is logged by the caller (`error!(error = %e)`
    /// in `server.rs`); the body returned here carries only a minimal,
    /// status-appropriate phrase so an unmatched / no-target request gets a
    /// clean deny rather than an internal echo.
    ///
    /// # Panics
    ///
    /// Panics if building a valid HTTP response with a plain-text body fails,
    /// which should never occur with well-formed status codes.
    pub fn error_response(error: &ProxyError) -> Response<BoxBody> {
        let status = error.status_code();
        // Generic, non-leaking body keyed purely off the status code. We
        // deliberately do NOT interpolate `error` (which can contain the Host,
        // path, backend address, or LB group name) into the client-visible
        // body.
        let body = status.canonical_reason().map_or_else(
            || status.as_str().to_string(),
            |reason| format!("{} {reason}", status.as_u16()),
        );

        Response::builder()
            .status(status)
            .header(header::CONTENT_TYPE, "text/plain; charset=utf-8")
            .body(full_body(body))
            .unwrap()
    }
}

impl Service<Request<Incoming>> for ReverseProxyService {
    type Response = Response<BoxBody>;
    type Error = ProxyError;
    type Future = std::pin::Pin<
        Box<
            dyn std::future::Future<Output = std::result::Result<Self::Response, Self::Error>>
                + Send,
        >,
    >;

    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<std::result::Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: Request<Incoming>) -> Self::Future {
        let this = self.clone();
        Box::pin(async move { this.proxy_request(req).await })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_error_response() {
        let error = ProxyError::RouteNotFound {
            host: "example.com".to_string(),
            path: "/api".to_string(),
        };

        let response = ReverseProxyService::error_response(&error);
        assert_eq!(response.status(), http::StatusCode::NOT_FOUND);
    }

    #[test]
    fn test_hop_by_hop_headers() {
        let mut parts = http::request::Builder::new()
            .method("GET")
            .uri("/test")
            .header("connection", "keep-alive, x-custom")
            .header("keep-alive", "timeout=5")
            .header("x-custom", "value")
            .header("x-other", "value")
            .body(())
            .unwrap()
            .into_parts()
            .0;

        ReverseProxyService::remove_hop_by_hop_headers(&mut parts);

        assert!(parts.headers.get("connection").is_none());
        assert!(parts.headers.get("keep-alive").is_none());
        assert!(parts.headers.get("x-custom").is_none());
        // x-other should remain
        assert!(parts.headers.get("x-other").is_some());
    }

    #[test]
    fn test_is_overlay_ip_accepts_overlay_range() {
        // 10.200.x.x should be recognized as overlay
        assert!(is_overlay_ip("10.200.0.1".parse().unwrap()));
        assert!(is_overlay_ip("10.200.255.254".parse().unwrap()));
        assert!(is_overlay_ip("10.200.1.100".parse().unwrap()));
    }

    #[test]
    fn test_is_overlay_ip_rejects_non_overlay() {
        // Non-overlay addresses
        assert!(!is_overlay_ip("192.168.1.1".parse().unwrap()));
        assert!(!is_overlay_ip("10.0.0.1".parse().unwrap()));
        assert!(!is_overlay_ip("10.201.0.1".parse().unwrap()));
        assert!(!is_overlay_ip("172.16.0.1".parse().unwrap()));
        assert!(!is_overlay_ip("8.8.8.8".parse().unwrap()));
    }

    #[test]
    fn test_is_overlay_ip_rejects_ipv6() {
        assert!(!is_overlay_ip("::1".parse().unwrap()));
        assert!(!is_overlay_ip("fe80::1".parse().unwrap()));
    }

    #[test]
    fn test_forbidden_error_response() {
        let error = ProxyError::Forbidden("endpoint 'ws' is internal-only".to_string());
        let response = ReverseProxyService::error_response(&error);
        assert_eq!(response.status(), http::StatusCode::FORBIDDEN);
    }

    // --- Tests for CF-Connecting-IP / X-Forwarded-For trust handling ------

    use crate::trust::TrustedProxyList;

    fn build_svc(peer: SocketAddr, trusted: TrustedProxyList) -> ReverseProxyService {
        let registry = Arc::new(ServiceRegistry::new());
        let load_balancer = Arc::new(LoadBalancer::new());
        let config = Arc::new(ProxyConfig::default());
        ReverseProxyService::new(registry, load_balancer, config)
            .with_remote_addr(peer)
            .with_trusted_proxies(Arc::new(trusted))
    }

    fn parts_with_headers(headers: &[(&str, &str)]) -> http::request::Parts {
        let mut builder = http::request::Builder::new().method("GET").uri("/");
        for (k, v) in headers {
            builder = builder.header(*k, *v);
        }
        builder.body(()).unwrap().into_parts().0
    }

    #[test]
    fn trusted_peer_cf_connecting_ip_is_honored() {
        // Peer 203.0.113.50 is inside the trusted /24. Its CF-Connecting-IP
        // should become X-Real-IP and be prepended to X-Forwarded-For.
        let peer: SocketAddr = "203.0.113.50:443".parse().unwrap();
        let trusted = TrustedProxyList::new(vec!["203.0.113.0/24".parse().unwrap()], None);
        let svc = build_svc(peer, trusted);

        let mut parts = parts_with_headers(&[("cf-connecting-ip", "198.51.100.7")]);
        svc.add_forwarding_headers(&mut parts);

        assert_eq!(parts.headers.get("x-real-ip").unwrap(), "198.51.100.7");
        let xff = parts
            .headers
            .get("x-forwarded-for")
            .unwrap()
            .to_str()
            .unwrap();
        assert!(
            xff.starts_with("198.51.100.7"),
            "XFF should start with real client IP, got {xff}"
        );
    }

    #[test]
    fn trusted_peer_xff_leftmost_is_honored_when_no_cf_header() {
        // Peer is trusted; no CF header but XFF chain is present. The leftmost
        // XFF entry is treated as the real client IP.
        let peer: SocketAddr = "203.0.113.50:443".parse().unwrap();
        let trusted = TrustedProxyList::new(vec!["203.0.113.0/24".parse().unwrap()], None);
        let svc = build_svc(peer, trusted);

        let mut parts = parts_with_headers(&[("x-forwarded-for", "198.51.100.9, 10.0.0.1")]);
        svc.add_forwarding_headers(&mut parts);

        assert_eq!(parts.headers.get("x-real-ip").unwrap(), "198.51.100.9");
        let xff = parts
            .headers
            .get("x-forwarded-for")
            .unwrap()
            .to_str()
            .unwrap();
        // Real client prepended, original chain preserved after.
        assert!(
            xff.starts_with("198.51.100.9"),
            "XFF should start with leftmost real client, got {xff}"
        );
        assert!(
            xff.contains("10.0.0.1"),
            "original chain should survive: {xff}"
        );
    }

    #[test]
    fn untrusted_peer_cf_connecting_ip_is_ignored() {
        // Peer 8.8.8.8 is NOT in the trusted list. The CF header must be
        // ignored and X-Real-IP must reflect the TCP peer.
        let peer: SocketAddr = "8.8.8.8:443".parse().unwrap();
        let trusted = TrustedProxyList::new(vec!["203.0.113.0/24".parse().unwrap()], None);
        let svc = build_svc(peer, trusted);

        let mut parts = parts_with_headers(&[("cf-connecting-ip", "198.51.100.7")]);
        svc.add_forwarding_headers(&mut parts);

        assert_eq!(parts.headers.get("x-real-ip").unwrap(), "8.8.8.8");
        let xff = parts
            .headers
            .get("x-forwarded-for")
            .unwrap()
            .to_str()
            .unwrap();
        // Untrusted peer: XFF should end with the peer IP (append behavior).
        assert!(
            xff.ends_with("8.8.8.8"),
            "XFF for untrusted peer should end with peer IP, got {xff}"
        );
    }

    #[test]
    fn no_headers_uses_peer_ip() {
        // No CF, no XFF. Any peer (trusted or not) should yield X-Real-IP ==
        // peer IP.
        let peer: SocketAddr = "198.51.100.250:443".parse().unwrap();
        let trusted = TrustedProxyList::localhost_only();
        let svc = build_svc(peer, trusted);

        let mut parts = parts_with_headers(&[]);
        svc.add_forwarding_headers(&mut parts);

        assert_eq!(parts.headers.get("x-real-ip").unwrap(), "198.51.100.250");
        assert_eq!(
            parts.headers.get("x-forwarded-for").unwrap(),
            "198.51.100.250"
        );
    }

    // --- RpsRegistry --------------------------------------------------------

    #[tokio::test]
    async fn rps_registry_counts_recorded_requests() {
        let reg = RpsRegistry::new();
        // Unknown service is 0.
        assert!((reg.rps("svc").await - 0.0).abs() < f64::EPSILON);

        // Record N requests; rps == N / window_secs.
        let n = 30;
        for _ in 0..n {
            reg.record("svc").await;
        }
        let expected = f64::from(n) / RPS_WINDOW.as_secs_f64();
        let got = reg.rps("svc").await;
        assert!(
            (got - expected).abs() < 1e-9,
            "expected {expected}, got {got}"
        );
    }

    #[tokio::test]
    async fn rps_registry_isolates_services() {
        let reg = RpsRegistry::new();
        reg.record("a").await;
        reg.record("a").await;
        reg.record("b").await;

        let snap = reg.snapshot().await;
        let a = snap.get("a").copied().unwrap_or_default();
        let b = snap.get("b").copied().unwrap_or_default();
        assert!(
            a > b,
            "service a ({a}) should have a higher rate than b ({b})"
        );
        // b recorded exactly once.
        assert!((b - 1.0 / RPS_WINDOW.as_secs_f64()).abs() < 1e-9);
    }

    #[tokio::test]
    async fn rps_registry_prunes_old_timestamps() {
        let reg = RpsRegistry::new();
        // Inject a timestamp well outside the window directly so the test does
        // not have to sleep for the full window.
        {
            let mut map = reg.services.lock().await;
            let ring = map.entry("svc".to_string()).or_default();
            let stale = Instant::now()
                .checked_sub(RPS_WINDOW + Duration::from_secs(5))
                .expect("instant underflow in test");
            ring.push_back(stale);
        }
        // The stale entry must be pruned on read.
        assert!((reg.rps("svc").await - 0.0).abs() < f64::EPSILON);
    }

    // --- Activator ----------------------------------------------------------

    use crate::lb::LbStrategy;

    /// Activator that flips a flag and registers a healthy backend on the
    /// shared load balancer, simulating a scale-from-zero wake-up.
    struct TestActivator {
        lb: Arc<LoadBalancer>,
        called: std::sync::atomic::AtomicBool,
    }

    #[async_trait::async_trait]
    impl Activator for TestActivator {
        async fn activate(&self, service: &str) -> std::result::Result<(), String> {
            self.called.store(true, std::sync::atomic::Ordering::SeqCst);
            // "Scale up": give the group a backend so the re-poll succeeds.
            self.lb.register(
                service,
                vec!["127.0.0.1:9".parse().unwrap()],
                LbStrategy::RoundRobin,
            );
            Ok(())
        }
    }

    fn build_svc_with_lb(lb: Arc<LoadBalancer>) -> ReverseProxyService {
        let registry = Arc::new(ServiceRegistry::new());
        let config = Arc::new(ProxyConfig::default());
        ReverseProxyService::new(registry, lb, config)
    }

    #[tokio::test]
    async fn select_or_activate_returns_none_without_activator() {
        let lb = Arc::new(LoadBalancer::new());
        let svc = build_svc_with_lb(Arc::clone(&lb));
        // No backends, no activator: immediate None (preserves 503 path).
        assert!(svc.select_or_activate("svc").await.is_none());
    }

    #[tokio::test]
    async fn select_or_activate_wakes_scaled_to_zero_service() {
        let lb = Arc::new(LoadBalancer::new());
        // Group exists but has no healthy backend (scale-to-zero idle).
        lb.register("svc", vec![], LbStrategy::RoundRobin);

        let activator = Arc::new(TestActivator {
            lb: Arc::clone(&lb),
            called: std::sync::atomic::AtomicBool::new(false),
        });
        let svc = build_svc_with_lb(Arc::clone(&lb)).with_activator(activator.clone());

        let backend = svc.select_or_activate("svc").await;
        assert!(
            backend.is_some(),
            "activator should have produced a backend"
        );
        assert!(
            activator.called.load(std::sync::atomic::Ordering::SeqCst),
            "activator must have been invoked"
        );
    }

    #[tokio::test]
    async fn select_or_activate_returns_existing_backend_without_calling_activator() {
        let lb = Arc::new(LoadBalancer::new());
        lb.register(
            "svc",
            vec!["127.0.0.1:9".parse().unwrap()],
            LbStrategy::RoundRobin,
        );
        let activator = Arc::new(TestActivator {
            lb: Arc::clone(&lb),
            called: std::sync::atomic::AtomicBool::new(false),
        });
        let svc = build_svc_with_lb(Arc::clone(&lb)).with_activator(activator.clone());

        assert!(svc.select_or_activate("svc").await.is_some());
        assert!(
            !activator.called.load(std::sync::atomic::Ordering::SeqCst),
            "activator must NOT be called when a backend already exists"
        );
    }
}