praxis_protocol/http/pingora/handler/

mod.rs

// SPDX-License-Identifier: MIT
// Copyright (c) 2024 Praxis Contributors

//! Pingora `ProxyHttp` implementation: the main HTTP reverse-proxy
//! handler.
//!
//! Bridges Pingora's hook-based lifecycle (`request_filter`,
//! `upstream_peer`, `upstream_request_filter`, etc.) to the Praxis
//! filter pipeline via `PingoraHttpHandler` (body-capable). Body
//! hooks are always available so hot reload can add body filters and
//! Pingora compression init remains one-shot.
//!
//! Each submodule implements one Pingora hook. The pipeline is held
//! behind `Arc<ArcSwap<FilterPipeline>>` for lock-free hot reload.

use std::{collections::HashMap, sync::Arc, time::Duration};

use arc_swap::ArcSwap;
use bytes::Bytes;
use pingora_core::{
    Result, apps::HttpServerOptions, protocols::http::v2::server::H2Options, server::Server,
    services::listening::Service,
};
use pingora_proxy::{Session, http_proxy};
use praxis_core::{config::ABSOLUTE_MAX_BODY_BYTES, connectivity::Upstream};
use praxis_filter::{BodyBuffer, BodyMode, CompressionConfig, FilterPipeline, HttpFilterContext, RequestExtensions};
use tokio::sync::Semaphore;
use tracing::{debug, warn};

use super::{context::PingoraRequestCtx, metrics};

/// Structured error responses for fatal proxy errors.
mod fail_to_proxy;
/// Shared hop-by-hop header stripping logic.
mod hop_by_hop;
/// Request header normalization (duplicate headers, obs-fold).
mod normalize;
/// Request body filter hook.
mod request_body_filter;
/// Request filter hook.
mod request_filter;
/// Reserved internal header helpers.
mod reserved_headers;
/// Response body filter hook.
mod response_body_filter;
/// Response filter hook.
mod response_filter;
/// Upstream peer selection hook.
mod upstream_peer;
/// Upstream request transformation hook.
mod upstream_request;
/// Upstream response hop-by-hop stripping hook.
mod upstream_response;
/// Via header injection hook.
mod via;
/// HTTP handler with body filter hooks.
mod with_body;

pub use upstream_peer::{arm_upstream_retry_gate, lock_upstream_retry_gate_tests};
pub use with_body::PingoraHttpHandler;

// -----------------------------------------------------------------------------
// Constants
// -----------------------------------------------------------------------------

/// Maximum number of upstream connection retries for idempotent requests.
const MAX_RETRIES: usize = 3;

/// Pingora still replays retry attempts from a fixed-size internal buffer.
///
/// `StreamBuffer` initial forwarding uses Praxis-owned `pre_read_body`, but
/// retry replay cannot safely cover bodies larger than Pingora's retry
/// buffer.
const RETRY_BODY_LIMIT: u64 = 65_536; // 64 KiB

// -----------------------------------------------------------------------------
// Load Handler
// -----------------------------------------------------------------------------

/// Load an HTTP handler for a single listener.
///
/// Any TLS certificate watcher shutdown senders are appended to
/// `cert_watcher_shutdowns`. The caller must keep this `Vec` alive
/// until server shutdown; dropping the senders signals the watcher
/// tasks to stop.
///
/// ```ignore
/// use std::sync::Arc;
///
/// use pingora_core::server::Server;
/// use praxis_core::config::Listener;
/// use praxis_filter::{FilterPipeline, FilterRegistry};
/// use praxis_protocol::http::pingora::handler::load_http_handler;
///
/// let mut server = Server::new(None).unwrap();
/// server.bootstrap();
/// let registry = FilterRegistry::with_builtins();
/// let pipeline = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
/// let listener = Listener {
///     name: "http".into(),
///     address: "127.0.0.1:8080".into(),
///     cluster: None,
///     downstream_read_timeout_ms: None,
///     filter_chains: vec![],
///     max_connections: None,
///     protocol: Default::default(),
///     tcp_session_timeout_ms: None,
///     tcp_max_duration_secs: None,
///     tls: None,
///     upstream: None,
/// };
/// let mut shutdowns = Vec::new();
/// load_http_handler(&mut server, &listener, pipeline, &mut shutdowns).unwrap();
/// ```
///
/// # Errors
///
/// Returns [`ProxyError`] if the listener fails to bind.
///
/// [`ProxyError`]: praxis_core::ProxyError
pub fn load_http_handler(
    server: &mut Server,
    listener: &praxis_core::config::Listener,
    pipeline: Arc<ArcSwap<FilterPipeline>>,
    cert_watcher_shutdowns: &mut Vec<tokio::sync::watch::Sender<bool>>,
) -> Result<(), praxis_core::ProxyError> {
    let downstream_read_timeout = listener.downstream_read_timeout_ms.map(Duration::from_millis);
    let connection_semaphore = listener
        .max_connections
        .map(|max| Arc::new(Semaphore::new(max as usize)));

    // Always use the body-capable handler: a reload may add body
    // filters, and compression init is one-shot in Pingora.
    debug!(listener = %listener.name, "loading HTTP handler with body filters");
    let handler = PingoraHttpHandler::new(
        pipeline,
        downstream_read_timeout,
        connection_semaphore,
        ::metrics::SharedString::from(listener.name.clone()),
    );
    wire_service(server, listener, handler, cert_watcher_shutdowns)?;
    Ok(())
}

/// Create a Pingora HTTP proxy service, bind the listener, and add it to the server.
fn wire_service<H>(
    server: &mut Server,
    listener: &praxis_core::config::Listener,
    handler: H,
    cert_watcher_shutdowns: &mut Vec<tokio::sync::watch::Sender<bool>>,
) -> Result<(), praxis_core::ProxyError>
where
    H: pingora_proxy::ProxyHttp + Send + Sync + 'static,
    H::CTX: Send + Sync,
{
    let service_name = format!("http-proxy:{name}", name = listener.name);
    let mut proxy = http_proxy(&server.configuration, handler);
    proxy.server_options = Some(h2c_server_options());
    proxy.h2_options = Some(h2_server_options());
    let mut service = Service::new(service_name, proxy);
    if let Some(tx) = super::listener::add_listener(&mut service, listener)? {
        cert_watcher_shutdowns.push(tx);
    }
    server.add_service(service);
    Ok(())
}

// -----------------------------------------------------------------------------
// Shared Utilities
// -----------------------------------------------------------------------------

/// Clamp a runtime-selected body mode to the byte ceiling implied by `baseline`.
///
/// `baseline` is the mode established before request/response-phase filter hooks
/// run (typically from pipeline capabilities + global body limits). Runtime
/// `set_*_body_mode` calls may widen limits; this helper preserves the original
/// ceiling while still allowing upgrades between body mode variants.
///
/// `Stream` mode passes through unconditionally because it delivers chunks
/// as they arrive without accumulating them — there is no buffer to cap.
/// A filter that downgrades from `StreamBuffer` to `Stream` at runtime is
/// opting out of buffering entirely, which is always safe from a memory
/// perspective. The pipeline-level body size limit (enforced separately
/// via `SizeLimit`) remains the backstop for oversized payloads.
fn clamp_body_mode_to_ceiling(mode: BodyMode, baseline: BodyMode) -> BodyMode {
    let ceiling = match baseline {
        BodyMode::StreamBuffer { max_bytes: Some(v) } | BodyMode::SizeLimit { max_bytes: v } => Some(v),
        _ => None,
    };

    match (mode, ceiling) {
        (BodyMode::StreamBuffer { max_bytes }, Some(limit)) => BodyMode::StreamBuffer {
            max_bytes: Some(max_bytes.map_or(limit, |v| v.min(limit))),
        },
        (BodyMode::SizeLimit { max_bytes }, Some(limit)) => BodyMode::SizeLimit {
            max_bytes: max_bytes.min(limit),
        },
        // Stream has no buffer to clamp; other modes pass through when the
        // baseline imposes no ceiling (e.g. unbounded StreamBuffer).
        (m, None | Some(_)) => m,
    }
}

/// Apply compression settings from the pipeline config to the Pingora response.
fn adjust_compression(
    session: &mut Session,
    upstream_response: &pingora_http::ResponseHeader,
    compression: Option<&CompressionConfig>,
) {
    use pingora_core::{modules::http::compression::ResponseCompression, protocols::http::compression::Algorithm};

    let Some(cfg) = compression else {
        return;
    };

    let Some(module) = session.downstream_modules_ctx.get_mut::<ResponseCompression>() else {
        return;
    };

    let headers = &upstream_response.headers;

    if !cfg.should_compress(headers) {
        debug!("disabling compression: response does not qualify");
        module.adjust_level(0);
        return;
    }

    for (enabled, level, algo) in [
        (cfg.gzip_enabled, cfg.gzip_level, Algorithm::Gzip),
        (cfg.brotli_enabled, cfg.brotli_level, Algorithm::Brotli),
        (cfg.zstd_enabled, cfg.zstd_level, Algorithm::Zstd),
    ] {
        if !enabled {
            module.adjust_algorithm_level(algo, 0);
        } else if let Some(lvl) = level {
            module.adjust_algorithm_level(algo, lvl);
        }
    }
}

/// Handle upstream connect failures with retry logic.
///
/// Retries are skipped when the effective forwarded body size exceeds
/// Pingora's retry buffer limit.
///
/// Body-mutating filters can change payload length after `request_filter`,
/// so the retry guard uses the larger of the original and mutated lengths.
fn handle_connect_failure(ctx: &mut PingoraRequestCtx, e: Box<pingora_core::Error>) -> Box<pingora_core::Error> {
    let cluster = ctx.metrics_cluster_shared.clone().unwrap_or_else(metrics::cluster_none);
    if let Some(start) = ctx.upstream_connect_start.take() {
        metrics::record_upstream_connect_duration(cluster.clone(), start.elapsed().as_secs_f64());
    }
    metrics::record_upstream_connect_failure(cluster.clone());
    if ctx.request_is_idempotent {
        maybe_retry_idempotent_connect(ctx, cluster, e)
    } else {
        e
    }
}

/// Decide whether an idempotent request may retry after a connect failure.
fn maybe_retry_idempotent_connect(
    ctx: &mut PingoraRequestCtx,
    cluster: ::metrics::SharedString,
    e: Box<pingora_core::Error>,
) -> Box<pingora_core::Error> {
    let mutated_len = ctx.mutated_request_body_len.unwrap_or(0) as u64;
    if std::cmp::max(ctx.request_body_bytes, mutated_len) > RETRY_BODY_LIMIT {
        warn!(
            body_bytes = ctx.request_body_bytes,
            mutated_len = ?ctx.mutated_request_body_len,
            limit = RETRY_BODY_LIMIT,
            "skipping retry: request body exceeds Pingora retry buffer limit"
        );
        record_retry_exhausted_if_attempted(ctx, cluster);
        return e;
    }
    if (ctx.retries as usize) < MAX_RETRIES {
        ctx.retries += 1;
        debug!(
            retries = ctx.retries,
            max = MAX_RETRIES,
            "retrying idempotent request after connect failure"
        );
        let mut e = e;
        e.set_retry(true);
        return e;
    }
    warn!(
        retries = ctx.retries,
        max = MAX_RETRIES,
        "retry limit reached for idempotent request"
    );
    metrics::record_upstream_retry(cluster, metrics::RETRY_RESULT_EXHAUSTED);
    e
}

/// Record `result=exhausted` only when at least one retry was already attempted.
fn record_retry_exhausted_if_attempted(ctx: &PingoraRequestCtx, cluster: ::metrics::SharedString) {
    if ctx.retries > 0 {
        metrics::record_upstream_retry(cluster, metrics::RETRY_RESULT_EXHAUSTED);
    }
}

/// Run response filters during the logging phase if the
/// response phase never executed (upstream error, filter
/// rejection, etc.).
async fn logging_cleanup(pipeline: &FilterPipeline, ctx: &mut PingoraRequestCtx) {
    if !ctx.response_phase_done
        && let Some(mut filter_ctx) = ctx.filter_context_for(pipeline, None)
    {
        let _result = pipeline.execute_http_response(&mut filter_ctx).await;
        let extensions = filter_ctx.extensions;
        let metadata = filter_ctx.filter_metadata;
        let state = filter_ctx.filter_state;
        let exec_idx = filter_ctx.executed_filter_indices;
        let body_idx = filter_ctx.body_done_indices;
        ctx.extensions = extensions;
        ctx.filter_metadata = metadata;
        ctx.filter_state = state;
        ctx.cached_executed_filter_indices = exec_idx;
        ctx.cached_body_done_indices = body_idx;
    }
}

/// Emit Prometheus metrics for a completed HTTP request.
///
/// No-op when the Prometheus recorder has not been installed.
fn emit_request_metrics(session: &Session, ctx: &PingoraRequestCtx) {
    if !metrics::is_recorder_installed() {
        return;
    }

    let status_code = session.response_written().map_or(0, |resp| resp.status.as_u16());
    let status_class = metrics::status_class(status_code);

    let request_method = session.req_header().method.as_str();
    let raw_method = if request_method.is_empty() {
        ctx.request_snapshot.as_ref().map_or("UNKNOWN", |r| r.method.as_str())
    } else {
        request_method
    };
    let method = metrics::method_label(raw_method);

    let cluster = ctx.metrics_cluster_shared.clone().unwrap_or_else(metrics::cluster_none);

    let route = ctx.metrics_route.clone().unwrap_or_else(metrics::route_unknown);

    let labels = metrics::RequestMetricLabels {
        cluster: cluster.clone(),
        method,
        route,
        status_class,
    };

    let duration_secs = ctx.request_start.elapsed().as_secs_f64();
    metrics::record_request_metrics(labels, duration_secs);
    metrics::record_body_size_metrics(
        method,
        status_class,
        cluster,
        ctx.request_body_bytes,
        ctx.response_body_bytes,
    );
}

/// Record a passive health observation for the selected upstream endpoint.
///
/// Called from the `logging` hook on every completed request. Determines
/// success/failure from the error argument and the stashed upstream
/// response status code.
///
/// No-op when no upstream was selected, no health registry is available,
/// or passive checking is not configured for the cluster.
fn record_passive_health(pipeline: &FilterPipeline, error: Option<&pingora_core::Error>, ctx: &PingoraRequestCtx) {
    let cluster_name = ctx.cluster.as_ref().or(ctx.metrics_cluster.as_ref());
    let Some(cluster_name) = cluster_name else {
        return;
    };
    let Some(idx) = ctx.selected_endpoint_index else {
        return;
    };
    let Some(registry) = pipeline.health_registry() else {
        return;
    };
    let Some(health) = registry.get(cluster_name) else {
        return;
    };

    let is_failure = error.is_some() || ctx.upstream_response_status.is_some_and(|s| s >= 500);
    apply_passive_threshold(health, idx, cluster_name, is_failure);
}

/// Apply passive health threshold for a single endpoint observation.
fn apply_passive_threshold(
    health: &praxis_core::health::ClusterHealthEntry,
    idx: usize,
    cluster_name: &Arc<str>,
    is_failure: bool,
) {
    if is_failure {
        if let Some(threshold) = health.passive_unhealthy_threshold()
            && health
                .endpoints()
                .get(idx)
                .is_some_and(|ep| ep.record_failure(threshold))
        {
            tracing::warn!(
                cluster = %cluster_name,
                endpoint_index = idx,
                threshold,
                "passive health: endpoint marked unhealthy"
            );
            emit_passive_health_transition(health, cluster_name, metrics::HEALTH_RESULT_UNHEALTHY);
        }
    } else if let Some(threshold) = health.passive_healthy_threshold()
        && health
            .endpoints()
            .get(idx)
            .is_some_and(|ep| ep.record_success(threshold))
    {
        tracing::info!(
            cluster = %cluster_name,
            endpoint_index = idx,
            threshold,
            "passive health: endpoint recovered"
        );
        emit_passive_health_transition(health, cluster_name, metrics::HEALTH_RESULT_HEALTHY);
    }
}

/// Refresh health gauges and increment the transition counter after a passive flip.
fn emit_passive_health_transition(
    health: &praxis_core::health::ClusterHealthEntry,
    cluster_name: &Arc<str>,
    result: &'static str,
) {
    let (healthy, total) = metrics::count_healthy_endpoints(health);
    metrics::record_health_transition(
        ::metrics::SharedString::from(Arc::clone(cluster_name)),
        result,
        healthy,
        total,
    );
}

/// Build [`HttpServerOptions`] with h2c enabled.
///
/// [`HttpServerOptions`]: pingora_core::apps::HttpServerOptions
fn h2c_server_options() -> HttpServerOptions {
    let mut opts = HttpServerOptions::default();
    opts.h2c = true;
    opts
}

/// Build [`H2Options`] with limits to mitigate HPACK amplification attacks
/// (CWE-409).
///
/// Without explicit limits the `h2` crate defaults allow unbounded header
/// list sizes and concurrent streams, enabling a small compressed request
/// to allocate hundreds of megabytes on the server.
///
/// [`H2Options`]: pingora_core::protocols::http::v2::server::H2Options
fn h2_server_options() -> H2Options {
    let mut opts = H2Options::new();
    opts.max_header_list_size(65_536); // 64 KiB
    opts.max_concurrent_streams(128);
    opts
}

/// Accumulate `chunk.len()` into `accumulated_bytes` and return `true` when
/// the total exceeds `max_bytes`. Returns `false` when the body is `None`.
fn check_body_size_limit(body: &Option<Bytes>, accumulated_bytes: &mut u64, max_bytes: usize) -> bool {
    if let Some(chunk) = body {
        #[expect(clippy::allow_attributes, reason = "cast lint is platform-dependent")]
        #[allow(clippy::cast_possible_truncation, reason = "chunk length fits u64")]
        let chunk_len = chunk.len() as u64;
        *accumulated_bytes += chunk_len;

        #[expect(clippy::allow_attributes, reason = "cast lint is platform-dependent")]
        #[allow(clippy::cast_possible_truncation, reason = "max_bytes fits u64")]
        let limit = max_bytes as u64;
        return *accumulated_bytes > limit;
    }
    false
}

/// Push `chunk` into the stream buffer, creating it if absent. At end-of-stream
/// the buffer is frozen into `body`. Returns `true` when the push overflows.
fn accumulate_stream_buffer(
    body: &mut Option<Bytes>,
    body_buffer: &mut Option<BodyBuffer>,
    end_of_stream: bool,
    max_bytes: Option<usize>,
) -> bool {
    if let Some(chunk) = &*body {
        let limit = max_bytes.unwrap_or(ABSOLUTE_MAX_BODY_BYTES);
        let buf = body_buffer.get_or_insert_with(|| BodyBuffer::new(limit));

        if buf.push(chunk.clone()).is_err() {
            return true;
        }
    }

    if end_of_stream {
        tracing::trace!("stream buffer: freezing accumulated body before pipeline at EOS");
        *body = body_buffer.take().map(BodyBuffer::freeze);
    } else {
        tracing::trace!("stream buffer: filters see the original chunk");
    }
    false
}

/// Suppress the body chunk while the stream buffer is still accumulating
/// (i.e. `Continue`/`BodyDone` before release).
#[expect(
    clippy::fn_params_excessive_bools,
    reason = "mirrors the caller's existing condition flags"
)]
fn suppress_stream_buffer_chunk(body: &mut Option<Bytes>, is_stream_buffer: bool, released: bool, end_of_stream: bool) {
    if is_stream_buffer && !released && !end_of_stream {
        *body = None;
    }
}

/// Release the accumulated stream buffer on `FilterAction::Release`.
fn release_stream_buffer(
    body: &mut Option<Bytes>,
    is_stream_buffer: bool,
    released: &mut bool,
    body_buffer: &mut Option<BodyBuffer>,
    end_of_stream: bool,
) {
    if is_stream_buffer && !*released {
        *released = true;
        if !end_of_stream {
            *body = body_buffer.take().map(BodyBuffer::freeze);
        }
    }
}

/// Shared fields extracted from an `HttpFilterContext` after body filter
/// execution. Written back to `PingoraRequestCtx` via [`write_back`].
///
/// [`write_back`]: BodyFilterOutput::write_back
struct BodyFilterOutput {
    /// Cluster selected by the filter pipeline.
    cluster: Option<Arc<str>>,
    /// Upstream endpoint selected by the load balancer.
    upstream: Option<Upstream>,
    /// Type-safe request-scoped extension container.
    extensions: RequestExtensions,
    /// Durable per-request metadata that persists across phases.
    filter_metadata: HashMap<String, String>,
    /// Typed per-filter state keyed by stable filter invocation ID.
    filter_state: HashMap<usize, Box<dyn std::any::Any + Send + Sync>>,
    /// Per-filter execution tracking indices.
    executed_filter_indices: Vec<bool>,
    /// Per-filter body-done tracking indices.
    body_done_indices: Vec<bool>,
}

impl BodyFilterOutput {
    /// Move the shared fields out of the filter context, replacing each
    /// with its `Default` value (zero-allocation no-ops for the types involved).
    fn take_from(fctx: &mut HttpFilterContext<'_>) -> Self {
        Self {
            cluster: fctx.cluster.take(),
            upstream: fctx.upstream.take(),
            extensions: std::mem::take(&mut fctx.extensions),
            filter_metadata: std::mem::take(&mut fctx.filter_metadata),
            filter_state: std::mem::take(&mut fctx.filter_state),
            executed_filter_indices: std::mem::take(&mut fctx.executed_filter_indices),
            body_done_indices: std::mem::take(&mut fctx.body_done_indices),
        }
    }

    /// Write the shared fields back to the protocol context.
    fn write_back(self, ctx: &mut PingoraRequestCtx) {
        ctx.cluster = self.cluster;
        ctx.upstream = self.upstream;
        ctx.extensions = self.extensions;
        ctx.filter_metadata = self.filter_metadata;
        ctx.filter_state = self.filter_state;
        ctx.cached_executed_filter_indices = self.executed_filter_indices;
        ctx.cached_body_done_indices = self.body_done_indices;
    }
}

// -----------------------------------------------------------------------------
// Tests
// -----------------------------------------------------------------------------

#[cfg(test)]
#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::indexing_slicing,
    clippy::field_reassign_with_default,
    clippy::too_many_lines,
    clippy::cast_possible_truncation,
    clippy::significant_drop_tightening,
    reason = "tests"
)]
mod tests {
    use std::sync::Arc;

    use praxis_core::connectivity::ConnectionOptions;

    use super::*;

    #[test]
    fn first_failure_idempotent_sets_retry() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(e.retry(), "first failure should set retry flag");
        assert_eq!(ctx.retries, 1);
    }

    #[test]
    fn large_body_skips_retry() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        ctx.request_body_bytes = RETRY_BODY_LIMIT + 1;
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(!e.retry(), "should not retry when body exceeds retry buffer limit");
        assert_eq!(ctx.retries, 0, "retry counter should not increment");
    }

    #[test]
    fn mutated_body_exceeding_limit_skips_retry() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        ctx.request_body_bytes = 1024;
        ctx.mutated_request_body_len = Some((RETRY_BODY_LIMIT + 1) as usize);
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(
            !e.retry(),
            "should not retry when mutated body exceeds retry buffer limit"
        );
        assert_eq!(ctx.retries, 0);
    }

    #[test]
    fn body_at_limit_allows_retry() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        ctx.request_body_bytes = RETRY_BODY_LIMIT;
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(e.retry(), "body exactly at limit should allow retry");
        assert_eq!(ctx.retries, 1);
    }

    #[test]
    fn zero_body_allows_retry() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        ctx.request_body_bytes = 0;
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(e.retry(), "zero-length body should allow retry");
        assert_eq!(ctx.retries, 1);
    }

    #[test]
    fn max_retries_exhausted_does_not_retry() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        ctx.retries = MAX_RETRIES as u32;
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(!e.retry(), "should not retry after MAX_RETRIES");
        assert_eq!(ctx.retries as usize, MAX_RETRIES);
    }

    #[test]
    fn counter_increments_across_calls() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = true;
        for expected in 1..=MAX_RETRIES {
            let _result = handle_connect_failure(&mut ctx, make_error());
            assert_eq!(ctx.retries as usize, expected);
        }
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(!e.retry(), "should not retry after reaching MAX_RETRIES");
        assert_eq!(ctx.retries as usize, MAX_RETRIES);
    }

    #[test]
    fn non_idempotent_request_never_retries() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.request_is_idempotent = false;
        let e = handle_connect_failure(&mut ctx, make_error());
        assert!(!e.retry(), "non-idempotent request should never retry");
        assert_eq!(ctx.retries, 0);
    }

    #[test]
    fn connect_failure_clears_upstream_connect_start() {
        let mut ctx = PingoraRequestCtx::default();
        ctx.upstream_connect_start = Some(std::time::Instant::now());
        let _e = handle_connect_failure(&mut ctx, make_error());
        assert!(
            ctx.upstream_connect_start.is_none(),
            "failed connect should consume upstream_connect_start for duration recording"
        );
    }

    #[tokio::test]
    async fn logging_cleanup_noop_when_response_phase_done() {
        let registry = praxis_filter::FilterRegistry::with_builtins();
        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        let mut ctx = PingoraRequestCtx::default();
        ctx.response_phase_done = true;
        ctx.request_snapshot = Some(praxis_filter::Request {
            method: http::Method::GET,
            uri: "/".parse().unwrap(),
            headers: http::HeaderMap::new(),
        });
        logging_cleanup(&pipeline, &mut ctx).await;
    }

    #[tokio::test]
    async fn logging_cleanup_noop_when_no_snapshot() {
        let registry = praxis_filter::FilterRegistry::with_builtins();
        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        let mut ctx = PingoraRequestCtx::default();
        ctx.response_phase_done = false;
        ctx.request_snapshot = None;
        logging_cleanup(&pipeline, &mut ctx).await;
    }

    #[tokio::test]
    async fn logging_cleanup_runs_response_pipeline_when_needed() {
        let registry = praxis_filter::FilterRegistry::with_builtins();
        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        let mut ctx = PingoraRequestCtx::default();
        ctx.response_phase_done = false;
        ctx.cluster = Some(Arc::from("test-cluster"));
        ctx.request_snapshot = Some(praxis_filter::Request {
            method: http::Method::GET,
            uri: "/test".parse().unwrap(),
            headers: http::HeaderMap::new(),
        });
        logging_cleanup(&pipeline, &mut ctx).await;
        assert!(ctx.cluster.is_none(), "cluster should be taken by logging_cleanup");
        assert!(ctx.upstream.is_none(), "upstream should be taken by logging_cleanup");
    }

    #[tokio::test]
    async fn logging_cleanup_preserves_filter_metadata() {
        let registry = praxis_filter::FilterRegistry::with_builtins();
        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        let mut ctx = PingoraRequestCtx::default();
        ctx.response_phase_done = false;
        ctx.filter_metadata
            .insert("json_rpc.method".to_owned(), "service/invoke".to_owned());
        ctx.request_snapshot = Some(praxis_filter::Request {
            method: http::Method::POST,
            uri: "/api".parse().unwrap(),
            headers: http::HeaderMap::new(),
        });
        logging_cleanup(&pipeline, &mut ctx).await;
        assert_eq!(
            ctx.filter_metadata.get("json_rpc.method").map(String::as_str),
            Some("service/invoke"),
            "filter_metadata should survive logging_cleanup"
        );
    }

    #[tokio::test]
    async fn logging_cleanup_preserves_extensions() {
        let registry = praxis_filter::FilterRegistry::with_builtins();
        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        let mut ctx = PingoraRequestCtx::default();
        ctx.response_phase_done = false;
        ctx.extensions.insert(42_u32);
        ctx.request_snapshot = Some(praxis_filter::Request {
            method: http::Method::POST,
            uri: "/test".parse().unwrap(),
            headers: http::HeaderMap::new(),
        });
        logging_cleanup(&pipeline, &mut ctx).await;
        assert_eq!(
            ctx.extensions.get::<u32>(),
            Some(&42),
            "extensions should survive logging_cleanup"
        );
    }

    #[test]
    fn passive_health_error_is_failure() {
        let (pipeline, ctx) = make_passive_scenario(Some(3), Some(2));
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(
            entry.endpoints()[0].is_healthy(),
            "single failure should not yet mark unhealthy (threshold=3)"
        );
    }

    #[test]
    fn passive_health_status_500_is_failure() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(3), Some(2));
        ctx.upstream_response_status = Some(500);
        record_passive_health(&pipeline, None, &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(
            entry.endpoints()[0].is_healthy(),
            "single 500 should not yet mark unhealthy (threshold=3)"
        );
    }

    #[test]
    fn passive_health_status_below_500_is_success() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(2), Some(1));
        ctx.upstream_response_status = Some(499);
        record_passive_health(&pipeline, None, &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(entry.endpoints()[0].is_healthy(), "status 499 should count as success");
    }

    #[test]
    fn passive_unhealthy_threshold_transition() {
        let (pipeline, ctx) = make_passive_scenario(Some(2), Some(1));
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);
        record_passive_health(&pipeline, Some(&error), &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(
            !entry.endpoints()[0].is_healthy(),
            "2 consecutive failures should mark unhealthy (threshold=2)"
        );
    }

    #[test]
    fn passive_healthy_threshold_recovery() {
        let (pipeline, ctx) = make_passive_scenario(Some(1), Some(2));
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(
            !entry.endpoints()[0].is_healthy(),
            "should be unhealthy after 1 failure"
        );

        let ctx_ok = make_passive_ctx("test-cluster", 0, Some(200));
        record_passive_health(&pipeline, None, &ctx_ok);
        assert!(
            !entry.endpoints()[0].is_healthy(),
            "one success should not recover (threshold=2)"
        );

        record_passive_health(&pipeline, None, &ctx_ok);
        assert!(
            entry.endpoints()[0].is_healthy(),
            "2 consecutive successes should recover (threshold=2)"
        );
    }

    #[test]
    fn passive_health_no_thresholds_is_noop() {
        let (pipeline, ctx) = make_passive_scenario(None, None);
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(
            entry.endpoints()[0].is_healthy(),
            "no passive thresholds means failures are no-op"
        );
    }

    #[test]
    fn passive_health_endpoint_index_out_of_bounds() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(1), Some(1));
        ctx.selected_endpoint_index = Some(999);
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(entry.endpoints()[0].is_healthy(), "out-of-bounds index should be no-op");
    }

    #[test]
    fn passive_health_missing_cluster_is_noop() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(1), Some(1));
        ctx.cluster = None;
        ctx.metrics_cluster = None;
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);
    }

    #[test]
    fn passive_health_falls_back_to_metrics_cluster() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(2), Some(1));
        ctx.cluster = None;
        ctx.metrics_cluster = Some(Arc::from("test-cluster"));
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);
        record_passive_health(&pipeline, Some(&error), &ctx);

        let registry = pipeline.health_registry().unwrap();
        let entry = registry.get("test-cluster").unwrap();
        assert!(
            !entry.endpoints()[0].is_healthy(),
            "fallback to metrics_cluster should still record passive health"
        );
    }

    #[test]
    fn passive_health_missing_endpoint_index_is_noop() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(1), Some(1));
        ctx.selected_endpoint_index = None;
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);
    }

    #[test]
    fn passive_health_missing_registry_is_noop() {
        let registry = praxis_filter::FilterRegistry::with_builtins();
        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        let mut ctx = PingoraRequestCtx::default();
        ctx.cluster = Some(Arc::from("test-cluster"));
        ctx.selected_endpoint_index = Some(0);
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);
    }

    #[test]
    fn passive_health_unknown_cluster_is_noop() {
        let (pipeline, mut ctx) = make_passive_scenario(Some(1), Some(1));
        ctx.cluster = Some(Arc::from("nonexistent"));
        let error = make_error();
        record_passive_health(&pipeline, Some(&error), &ctx);
    }

    #[test]
    fn size_limit_none_body_returns_false() {
        let mut bytes = 0_u64;
        assert!(!check_body_size_limit(&None, &mut bytes, 100));
        assert_eq!(bytes, 0, "accumulated bytes unchanged for None body");
    }

    #[test]
    fn size_limit_within_limit() {
        let mut bytes = 0_u64;
        let body = Some(Bytes::from_static(b"hello"));
        assert!(!check_body_size_limit(&body, &mut bytes, 10));
        assert_eq!(bytes, 5);
    }

    #[test]
    fn size_limit_at_exact_limit() {
        let mut bytes = 0_u64;
        let body = Some(Bytes::from_static(b"exact"));
        assert!(!check_body_size_limit(&body, &mut bytes, 5));
        assert_eq!(bytes, 5);
    }

    #[test]
    fn size_limit_exceeds_limit() {
        let mut bytes = 0_u64;
        let body = Some(Bytes::from_static(b"toolong"));
        assert!(check_body_size_limit(&body, &mut bytes, 3));
    }

    #[test]
    fn size_limit_cumulative_overflow() {
        let mut bytes = 0_u64;
        let first = Some(Bytes::from_static(b"aaa"));
        assert!(!check_body_size_limit(&first, &mut bytes, 5));

        let second = Some(Bytes::from_static(b"bbb"));
        assert!(check_body_size_limit(&second, &mut bytes, 5));
        assert_eq!(bytes, 6);
    }

    #[test]
    fn stream_buffer_accumulates_chunks() {
        let mut body = Some(Bytes::from_static(b"hello "));
        let mut buf: Option<BodyBuffer> = None;
        assert!(!accumulate_stream_buffer(&mut body, &mut buf, false, Some(100)));
        assert!(buf.is_some());

        body = Some(Bytes::from_static(b"world"));
        assert!(!accumulate_stream_buffer(&mut body, &mut buf, false, Some(100)));

        let frozen = buf.take().unwrap().freeze();
        assert_eq!(frozen, Bytes::from_static(b"hello world"));
    }

    #[test]
    fn stream_buffer_freezes_at_eos() {
        let mut body = Some(Bytes::from_static(b"data"));
        let mut buf: Option<BodyBuffer> = None;
        assert!(!accumulate_stream_buffer(&mut body, &mut buf, false, Some(100)));

        body = Some(Bytes::from_static(b" end"));
        assert!(!accumulate_stream_buffer(&mut body, &mut buf, true, Some(100)));
        assert!(buf.is_none(), "buffer should be taken at EOS");
        assert_eq!(body.unwrap(), Bytes::from_static(b"data end"));
    }

    #[test]
    fn stream_buffer_overflow() {
        let mut body = Some(Bytes::from_static(b"too long"));
        let mut buf: Option<BodyBuffer> = None;
        assert!(accumulate_stream_buffer(&mut body, &mut buf, false, Some(5)));
    }

    #[test]
    fn stream_buffer_none_body() {
        let mut body: Option<Bytes> = None;
        let mut buf: Option<BodyBuffer> = None;
        assert!(!accumulate_stream_buffer(&mut body, &mut buf, false, Some(100)));
        assert!(buf.is_none());
    }

    #[test]
    fn stream_buffer_uses_absolute_max_when_none() {
        let mut body = Some(Bytes::from_static(b"data"));
        let mut buf: Option<BodyBuffer> = None;
        assert!(!accumulate_stream_buffer(&mut body, &mut buf, false, None));
        assert!(buf.is_some(), "should create buffer with absolute max");
    }

    #[test]
    fn suppress_clears_body_when_buffering() {
        let mut body = Some(Bytes::from_static(b"data"));
        suppress_stream_buffer_chunk(&mut body, true, false, false);
        assert!(body.is_none());
    }

    #[test]
    fn suppress_noop_when_not_stream_buffer() {
        let mut body = Some(Bytes::from_static(b"data"));
        suppress_stream_buffer_chunk(&mut body, false, false, false);
        assert!(body.is_some());
    }

    #[test]
    fn suppress_noop_when_released() {
        let mut body = Some(Bytes::from_static(b"data"));
        suppress_stream_buffer_chunk(&mut body, true, true, false);
        assert!(body.is_some());
    }

    #[test]
    fn suppress_noop_at_eos() {
        let mut body = Some(Bytes::from_static(b"data"));
        suppress_stream_buffer_chunk(&mut body, true, false, true);
        assert!(body.is_some());
    }

    #[test]
    fn release_sets_flag_and_flushes_buffer() {
        let mut body: Option<Bytes> = None;
        let mut released = false;
        let mut buf = Some(BodyBuffer::new(100));
        buf.as_mut().unwrap().push(Bytes::from_static(b"buffered")).unwrap();

        release_stream_buffer(&mut body, true, &mut released, &mut buf, false);
        assert!(released);
        assert_eq!(body.unwrap(), Bytes::from_static(b"buffered"));
        assert!(buf.is_none());
    }

    #[test]
    fn release_noop_when_already_released() {
        let mut body: Option<Bytes> = None;
        let mut released = true;
        let mut buf: Option<BodyBuffer> = None;

        release_stream_buffer(&mut body, true, &mut released, &mut buf, false);
        assert!(body.is_none(), "body should be unchanged when already released");
    }

    #[test]
    fn release_noop_when_not_stream_buffer() {
        let mut body: Option<Bytes> = None;
        let mut released = false;
        let mut buf: Option<BodyBuffer> = None;

        release_stream_buffer(&mut body, false, &mut released, &mut buf, false);
        assert!(!released, "released flag should be unchanged for non-stream-buffer");
    }

    #[test]
    fn release_at_eos_sets_flag_but_no_flush() {
        let mut body: Option<Bytes> = None;
        let mut released = false;
        let mut buf = Some(BodyBuffer::new(100));
        buf.as_mut().unwrap().push(Bytes::from_static(b"data")).unwrap();

        release_stream_buffer(&mut body, true, &mut released, &mut buf, true);
        assert!(released);
        assert!(body.is_none(), "body should not be overwritten at EOS");
        assert!(buf.is_some(), "buffer should not be taken at EOS");
    }

    #[test]
    fn write_back_transfers_fields() {
        let mut ctx = PingoraRequestCtx::default();

        let mut extensions = RequestExtensions::new();
        extensions.insert(42_u32);

        let state_val: Box<dyn std::any::Any + Send + Sync> = Box::new(99_i32);
        let filter_state = HashMap::from([(0_usize, state_val)]);

        let output = BodyFilterOutput {
            cluster: Some(Arc::from("test-cluster")),
            upstream: Some(Upstream {
                address: Arc::from("10.0.0.1:80"),
                connection: Arc::new(ConnectionOptions::default()),
                tls: None,
            }),
            extensions,
            filter_metadata: HashMap::from([("key".to_owned(), "val".to_owned())]),
            filter_state,
            executed_filter_indices: vec![true, false],
            body_done_indices: vec![false, true],
        };
        output.write_back(&mut ctx);

        assert_eq!(ctx.cluster.as_deref(), Some("test-cluster"));
        assert!(ctx.upstream.is_some(), "upstream should transfer");
        assert_eq!(ctx.upstream.as_ref().unwrap().address.as_ref(), "10.0.0.1:80");
        assert_eq!(ctx.extensions.get::<u32>(), Some(&42));
        assert_eq!(ctx.filter_metadata.get("key").map(String::as_str), Some("val"));
        assert_eq!(ctx.filter_state.len(), 1, "filter_state should transfer");
        assert_eq!(
            ctx.filter_state.get(&0).and_then(|v| v.downcast_ref::<i32>()),
            Some(&99)
        );
        assert_eq!(ctx.cached_executed_filter_indices, vec![true, false]);
        assert_eq!(ctx.cached_body_done_indices, vec![false, true]);
    }

    // -------------------------------------------------------------------------
    // Test Utilities
    // -------------------------------------------------------------------------

    /// Create a connect error for tests.
    fn make_error() -> Box<pingora_core::Error> {
        pingora_core::Error::explain(pingora_core::ErrorType::ConnectError, "test connect failure")
    }

    /// Build a [`PingoraRequestCtx`] for passive health testing.
    fn make_passive_ctx(cluster: &str, endpoint_idx: usize, status: Option<u16>) -> PingoraRequestCtx {
        let mut ctx = PingoraRequestCtx::default();
        ctx.cluster = Some(Arc::from(cluster));
        ctx.selected_endpoint_index = Some(endpoint_idx);
        ctx.upstream_response_status = status;
        ctx
    }

    /// Build a pipeline with a health registry and a matching context
    /// for passive health testing.
    fn make_passive_scenario(
        passive_unhealthy: Option<u32>,
        passive_healthy: Option<u32>,
    ) -> (FilterPipeline, PingoraRequestCtx) {
        use std::collections::HashMap;

        use praxis_core::health::{ClusterHealthEntry, EndpointHealth};

        let entry = ClusterHealthEntry::new(
            vec![EndpointHealth::new()],
            vec![Arc::from("10.0.0.1:80")],
            passive_unhealthy,
            passive_healthy,
        );
        let mut map = HashMap::new();
        map.insert(Arc::from("test-cluster"), Arc::new(entry));
        let health_registry = Arc::new(map);

        let registry = praxis_filter::FilterRegistry::with_builtins();
        let mut pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
        pipeline.set_health_registry(health_registry);

        let ctx = make_passive_ctx("test-cluster", 0, None);

        (pipeline, ctx)
    }
}