warpdrive_proxy/middleware/

mod.rs

//! Middleware system for WarpDrive
//!
//! This module provides a composable middleware architecture adapted to Pingora's
//! filter-based model. Each middleware operates on requests and responses as they
//! flow through the proxy.
//!
//! # Architecture
//!
//! Middlewares are applied in the following order:
//! 1. **Headers** - Inject X-Forwarded-* headers (request_filter)
//! 2. **Logging** - Log request/response metrics (request_filter + response_filter)
//! 3. **Sendfile** - X-Sendfile/X-Accel-Redirect support (response_filter)
//! 4. **Compression** - Response compression (response_filter)
//!
//! # Usage with Pingora
//!
//! ```no_run
//! use warpdrive::middleware::MiddlewareStack;
//! use pingora::proxy::Session;
//!
//! // In ProxyHttp implementation:
//! async fn request_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<()> {
//!     self.middleware.apply_request_filters(session, ctx).await
//! }
//!
//! async fn response_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<()> {
//!     self.middleware.apply_response_filters(session, ctx).await
//! }
//! ```

pub mod circuit_breaker;
pub mod compression;
pub mod concurrency;
pub mod headers;
pub mod logging;
pub mod rate_limit;
pub mod sendfile;
pub mod static_files;
pub mod trusted_ranges;

#[cfg(test)]
mod tests;

use async_trait::async_trait;
use bytes::Bytes;
use pingora::http::ResponseHeader;
use pingora::prelude::*;
use std::path::PathBuf;
use std::sync::Arc;

use crate::config::Config;

pub use circuit_breaker::CircuitBreakerMiddleware;
pub use compression::CompressionMiddleware;
pub use concurrency::ConcurrencyMiddleware;
pub use headers::HeadersMiddleware;
pub use logging::LoggingMiddleware;
pub use rate_limit::RateLimitMiddleware;
pub use sendfile::SendfileMiddleware;
pub use static_files::StaticFilesMiddleware;
pub use trusted_ranges::TrustedRangesMiddleware;

/// Context for middleware execution
///
/// Stores state needed across middleware chain execution, including
/// request start time, response metadata, and flags.
pub struct MiddlewareContext {
    /// Request start timestamp (for logging)
    pub request_start: std::time::Instant,

    /// Response status code (captured by logging middleware)
    pub status_code: u16,

    /// Response body size (captured by logging middleware)
    pub body_size: usize,

    /// Sendfile state (set by sendfile middleware)
    pub sendfile: SendfileState,

    /// Compression state (set by compression middleware)
    pub compression: CompressionState,

    /// Static response (set by static files middleware to short-circuit proxy)
    pub static_response: Option<StaticResponse>,

    /// Concurrency permit (held for request duration, auto-released on drop)
    pub concurrency_permit: Option<tokio::sync::OwnedSemaphorePermit>,

    /// Streaming flag (set when SSE or other streaming content is detected)
    ///
    /// When true, disables response buffering to allow real-time streaming.
    /// Automatically set when:
    /// - Content-Type: text/event-stream (Server-Sent Events)
    /// - X-Accel-Buffering: no (nginx compatibility)
    pub streaming: bool,

    /// Trusted source flag (set by trusted_ranges middleware)
    ///
    /// When true, indicates request comes from a trusted proxy/CDN IP range.
    /// Used to bypass rate limiting and concurrency limits.
    pub trusted_source: bool,

    /// Real client IP (normalized by trusted_ranges middleware)
    ///
    /// Contains the actual client IP extracted from proxy headers (if trusted)
    /// or the socket IP (if not from trusted source). Used for:
    /// - Rate limiting (to rate limit real clients, not proxies)
    /// - Logging (to log actual client IPs)
    /// - X-Forwarded-For headers (to maintain proxy chain)
    pub real_client_ip: std::net::IpAddr,
}

impl Default for MiddlewareContext {
    fn default() -> Self {
        Self {
            request_start: std::time::Instant::now(),
            status_code: 200,
            body_size: 0,
            sendfile: SendfileState::default(),
            compression: CompressionState::default(),
            static_response: None,
            concurrency_permit: None,
            streaming: false,
            trusted_source: false,
            real_client_ip: "0.0.0.0".parse().unwrap(),
        }
    }
}

/// Body payload for a static response
#[derive(Debug)]
pub enum StaticResponseBody {
    /// In-memory buffer (small files)
    InMemory(Bytes),
    /// Streamed from disk (large files)
    Stream(PathBuf),
}

/// Pre-computed static response returned by the static files middleware
#[derive(Debug)]
pub struct StaticResponse {
    pub header: ResponseHeader,
    pub body: StaticResponseBody,
}

/// Tracks sendfile middleware state for a response
#[derive(Debug, Default)]
pub struct SendfileState {
    /// Whether sendfile is engaged for this response
    pub active: bool,
    /// Optional path for logging/diagnostics
    pub path: Option<String>,
    /// In-memory body to serve (loaded once when active)
    pub body: Option<Bytes>,
    /// Flag indicating if body has been emitted downstream
    pub served: bool,
}

impl SendfileState {
    pub fn activate(&mut self, path: String, body: Bytes) {
        self.active = true;
        self.path = Some(path);
        self.body = Some(body);
        self.served = false;
    }

    pub fn reset(&mut self) {
        self.active = false;
        self.path = None;
        self.body = None;
        self.served = false;
    }
}

/// Compression algorithm in use for the current response
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum CompressionEncoding {
    /// Brotli compression (preferred, better ratio)
    Brotli,
    /// Gzip compression (fallback, wider support)
    Gzip,
}

/// Tracks compression middleware state across body chunks
#[derive(Debug, Default)]
pub enum CompressionState {
    /// Compression disabled for this response
    #[default]
    Disabled,
    /// Compression enabled and buffering body chunks
    Pending {
        buffer: Vec<u8>,
        encoding: CompressionEncoding,
    },
    /// Compression already applied
    Complete,
}

impl CompressionState {
    pub fn enable(&mut self, encoding: CompressionEncoding) {
        *self = CompressionState::Pending {
            buffer: Vec::new(),
            encoding,
        };
    }

    pub fn is_enabled(&self) -> bool {
        matches!(self, CompressionState::Pending { .. })
    }
}

/// Middleware stack that composes all middlewares
///
/// This struct holds all middleware components and orchestrates their execution
/// in the correct order for both request and response phases.
pub struct MiddlewareStack {
    pub trusted_ranges: Option<TrustedRangesMiddleware>,
    pub static_files: Option<StaticFilesMiddleware>,
    pub concurrency: Option<ConcurrencyMiddleware>,
    pub rate_limit: Option<RateLimitMiddleware>,
    pub circuit_breaker: Option<CircuitBreakerMiddleware>,
    pub headers: HeadersMiddleware,
    pub logging: Option<LoggingMiddleware>,
    pub sendfile: Option<SendfileMiddleware>,
    pub compression: Option<CompressionMiddleware>,
}

impl MiddlewareStack {
    /// Create a new middleware stack from configuration
    pub fn new(config: Arc<Config>) -> Self {
        Self {
            trusted_ranges: if config.trusted_ranges_file.is_some()
                || config.client_ip_header.is_some()
            {
                TrustedRangesMiddleware::new(
                    config.trusted_ranges_file.clone(),
                    config.client_ip_header.clone(),
                )
                .ok()
            } else {
                None
            },
            static_files: if config.static_enabled {
                Some(StaticFilesMiddleware::from_config(&config))
            } else {
                None
            },
            concurrency: if config.max_concurrent_requests > 0 {
                Some(ConcurrencyMiddleware::new(
                    true,
                    config.max_concurrent_requests,
                ))
            } else {
                None
            },
            rate_limit: if config.rate_limit_enabled {
                Some(RateLimitMiddleware::new(
                    true,
                    config.rate_limit_requests_per_sec,
                    config.rate_limit_burst_size,
                ))
            } else {
                None
            },
            circuit_breaker: if config.circuit_breaker_enabled {
                Some(CircuitBreakerMiddleware::new(
                    true,
                    config.circuit_breaker_failure_threshold,
                    config.circuit_breaker_timeout_secs,
                ))
            } else {
                None
            },
            headers: HeadersMiddleware::new(config.clone()),
            logging: if config.log_requests {
                Some(LoggingMiddleware::new())
            } else {
                None
            },
            sendfile: if config.x_sendfile_enabled {
                Some(SendfileMiddleware::new())
            } else {
                None
            },
            compression: if config.gzip_compression_enabled {
                Some(CompressionMiddleware::new())
            } else {
                None
            },
        }
    }

    /// Apply all request filters in order
    ///
    /// This is called from ProxyHttp::request_filter() to process the incoming request
    /// before it's sent to the upstream server.
    pub async fn apply_request_filters(
        &self,
        session: &mut Session,
        ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        // 0. Trusted ranges - normalize client IP and set trusted flag (must run FIRST)
        if let Some(ref trusted_ranges) = self.trusted_ranges {
            trusted_ranges.request_filter(session, ctx).await?;
        }

        // 1. Static files - serve directly if path matches (skip all other middleware + proxy)
        if let Some(ref static_files) = self.static_files {
            static_files.request_filter(session, ctx).await?;

            // If static response was set, short-circuit the proxy
            if ctx.static_response.is_some() {
                return Ok(());
            }
        }

        // 2. Concurrency limiting - limit total concurrent requests
        if let Some(ref concurrency) = self.concurrency {
            concurrency.request_filter(session, ctx).await?;
        }

        // 3. Rate limiting - check limits (fail fast)
        if let Some(ref rate_limit) = self.rate_limit {
            rate_limit.request_filter(session, ctx).await?;
        }

        // 4. Circuit breaker - check if upstream is available
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            circuit_breaker.request_filter(session, ctx).await?;
        }

        // 5. Headers middleware - add X-Forwarded-* headers
        self.headers.request_filter(session, ctx).await?;

        // 6. Logging middleware - record request start
        if let Some(ref logging) = self.logging {
            logging.request_filter(session, ctx).await?;
        }

        Ok(())
    }

    /// Apply all response filters in order
    ///
    /// This is called from ProxyHttp::response_filter() to process the upstream response
    /// before it's sent to the client.
    pub async fn apply_response_filters(
        &self,
        session: &mut Session,
        upstream_response: &mut ResponseHeader,
        ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        // 1. Circuit breaker - record upstream response (success/failure)
        if let Some(ref circuit_breaker) = self.circuit_breaker {
            circuit_breaker
                .response_filter(session, upstream_response, ctx)
                .await?;
        }

        // 2. Sendfile middleware - handle X-Sendfile headers
        if let Some(ref sendfile) = self.sendfile {
            sendfile
                .response_filter(session, upstream_response, ctx)
                .await?;
        }

        // 2. Compression middleware - compress response if applicable
        if let Some(ref compression) = self.compression {
            compression
                .response_filter(session, upstream_response, ctx)
                .await?;
        }

        // 3. Logging middleware - log response (must be last to capture all metadata)
        if let Some(ref logging) = self.logging {
            logging
                .response_filter(session, upstream_response, ctx)
                .await?;
        }

        Ok(())
    }
}

/// Trait for middleware components
///
/// Each middleware implements this trait to define its behavior during
/// request and response phases. Both methods are optional - middlewares
/// can implement only the phase they care about.
#[async_trait]
pub trait Middleware: Send + Sync {
    /// Process request before sending to upstream
    async fn request_filter(
        &self,
        _session: &mut Session,
        _ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        Ok(())
    }

    /// Process response before sending to client
    async fn response_filter(
        &self,
        _session: &mut Session,
        _upstream_response: &mut ResponseHeader,
        _ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        Ok(())
    }
}