warpdrive_proxy/middleware/

compression.rs

//! Response compression middleware
//!
//! Implements Brotli and gzip compression for responses based on client
//! Accept-Encoding headers.
//!
//! # Compression Strategy
//!
//! - Only compress if client supports gzip (Accept-Encoding: gzip)
//! - Skip compression for already-compressed content (Content-Encoding present)
//! - Skip compression for small responses (< 1KB)
//! - Skip compression for non-compressible types (images, video, etc.)
//!
//! # Note on Pingora Integration
//!
//! Pingora doesn't have built-in response body modification in filters, so this
//! middleware sets headers to signal compression intent. Actual compression happens
//! in the response body stream using async-compression.

use async_trait::async_trait;
use pingora::http::ResponseHeader;
use pingora::prelude::*;
use tracing::debug;

use super::{CompressionEncoding, Middleware, MiddlewareContext};

/// Compression middleware
///
/// Adds compression support for responses based on client capabilities.
pub struct CompressionMiddleware {
    /// Minimum response size to compress (bytes)
    min_compress_size: usize,
}

impl CompressionMiddleware {
    /// Create new compression middleware
    pub fn new() -> Self {
        Self {
            min_compress_size: 1024, // 1KB minimum
        }
    }

    /// Check if content type should be compressed
    pub(crate) fn should_compress_content_type(content_type: &str) -> bool {
        // Compress text-based content
        let compressible_prefixes = [
            "text/",
            "application/json",
            "application/javascript",
            "application/xml",
            "application/x-javascript",
            "application/xhtml+xml",
            "application/rss+xml",
            "application/atom+xml",
        ];

        compressible_prefixes
            .iter()
            .any(|prefix| content_type.starts_with(prefix))
    }

    /// Determine the best compression encoding supported by the client
    ///
    /// Returns the preferred encoding in order: Brotli > Gzip > None
    /// Brotli provides better compression ratios (~15-25% smaller than gzip)
    fn best_supported_encoding(session: &Session) -> Option<CompressionEncoding> {
        if let Some(accept_encoding) = session.req_header().headers.get("accept-encoding") {
            if let Ok(value) = accept_encoding.to_str() {
                let encodings: Vec<&str> = value.split(',').map(|s| s.trim()).collect();

                // Prefer Brotli (better compression ratio)
                if encodings.contains(&"br") {
                    return Some(CompressionEncoding::Brotli);
                }

                // Fallback to Gzip (wider support)
                if encodings.contains(&"gzip") {
                    return Some(CompressionEncoding::Gzip);
                }
            }
        }
        None
    }
}

impl Default for CompressionMiddleware {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl Middleware for CompressionMiddleware {
    /// Process response to determine if compression should be applied
    ///
    /// This sets a custom header (X-WarpDrive-Compress) to signal that compression
    /// should be applied to the response body downstream.
    async fn response_filter(
        &self,
        session: &mut Session,
        upstream_response: &mut ResponseHeader,
        ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        // Detect streaming responses
        // Check for Content-Type: text/event-stream (Server-Sent Events)
        if let Some(content_type) = upstream_response.headers.get("content-type") {
            if let Ok(ct_str) = content_type.to_str() {
                if ct_str.contains("text/event-stream") {
                    debug!("Streaming detected: SSE (text/event-stream)");
                    ctx.streaming = true;
                    return Ok(());
                }
            }
        }

        // Check for X-Accel-Buffering: no (nginx compatibility)
        if let Some(buffering) = upstream_response.headers.get("x-accel-buffering") {
            if let Ok(value) = buffering.to_str() {
                if value.to_lowercase() == "no" {
                    debug!("Streaming detected: X-Accel-Buffering: no");
                    ctx.streaming = true;
                    return Ok(());
                }
            }
        }

        // Determine best compression encoding supported by client
        let encoding = match Self::best_supported_encoding(session) {
            Some(enc) => enc,
            None => {
                debug!("Client doesn't support compression");
                return Ok(());
            }
        };

        // Skip if sendfile already took over body streaming
        if ctx.sendfile.active {
            debug!("Skipping compression: sendfile response");
            return Ok(());
        }

        // Skip if already compressed
        if upstream_response.headers.contains_key("content-encoding") {
            debug!("Skipping compression: already encoded");
            return Ok(());
        }

        // Check content type
        if let Some(content_type) = upstream_response.headers.get("content-type") {
            if let Ok(ct_str) = content_type.to_str() {
                if !Self::should_compress_content_type(ct_str) {
                    debug!(
                        "Skipping compression: non-compressible content type {}",
                        ct_str
                    );
                    return Ok(());
                }
            }
        }

        // Check content length
        if let Some(content_length) = upstream_response.headers.get("content-length") {
            if let Ok(length_str) = content_length.to_str() {
                if let Ok(length) = length_str.parse::<usize>() {
                    if length < self.min_compress_size {
                        debug!(
                            "Skipping compression: response too small ({} bytes)",
                            length
                        );
                        return Ok(());
                    }
                }
            }
        }

        // At this point, compression is eligible
        debug!("Response eligible for {:?} compression", encoding);

        // Set Vary header to indicate response varies by Accept-Encoding
        upstream_response.insert_header("Vary", "Accept-Encoding")?;

        // Set Content-Encoding and remove Content-Length (size changes post-compression)
        let encoding_str = match encoding {
            CompressionEncoding::Brotli => "br",
            CompressionEncoding::Gzip => "gzip",
        };
        upstream_response.insert_header("Content-Encoding", encoding_str)?;
        upstream_response.remove_header("Content-Length");

        // Initialize compression state in context for body filter
        if !ctx.compression.is_enabled() {
            ctx.compression.enable(encoding);
        }

        Ok(())
    }
}

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

    #[test]
    fn test_should_compress_content_type() {
        assert!(CompressionMiddleware::should_compress_content_type(
            "text/html"
        ));
        assert!(CompressionMiddleware::should_compress_content_type(
            "text/plain"
        ));
        assert!(CompressionMiddleware::should_compress_content_type(
            "application/json"
        ));
        assert!(CompressionMiddleware::should_compress_content_type(
            "application/javascript"
        ));
        assert!(CompressionMiddleware::should_compress_content_type(
            "application/xml"
        ));

        assert!(!CompressionMiddleware::should_compress_content_type(
            "image/png"
        ));
        assert!(!CompressionMiddleware::should_compress_content_type(
            "image/jpeg"
        ));
        assert!(!CompressionMiddleware::should_compress_content_type(
            "video/mp4"
        ));
        assert!(!CompressionMiddleware::should_compress_content_type(
            "application/pdf"
        ));
    }

    #[test]
    fn test_compression_middleware_creation() {
        let middleware = CompressionMiddleware::new();
        assert_eq!(middleware.min_compress_size, 1024);

        let middleware_default = CompressionMiddleware::default();
        assert_eq!(middleware_default.min_compress_size, 1024);
    }

    #[test]
    fn test_sse_streaming_detection() {
        // Test that SSE content type is detected
        use MiddlewareContext;
        use pingora::http::ResponseHeader;

        let mut ctx = MiddlewareContext::default();
        assert!(!ctx.streaming, "Context should start with streaming=false");

        // Create response with SSE content type
        let mut response = ResponseHeader::build(200, None).unwrap();
        response
            .insert_header("content-type", "text/event-stream")
            .unwrap();

        // Content type should contain event-stream
        if let Some(ct) = response.headers.get("content-type") {
            if let Ok(ct_str) = ct.to_str() {
                assert!(ct_str.contains("text/event-stream"));
            }
        }
    }

    #[test]
    fn test_x_accel_buffering_detection() {
        use MiddlewareContext;
        use pingora::http::ResponseHeader;

        let mut ctx = MiddlewareContext::default();
        assert!(!ctx.streaming);

        // Test X-Accel-Buffering: no
        let mut response = ResponseHeader::build(200, None).unwrap();
        response.insert_header("x-accel-buffering", "no").unwrap();

        if let Some(buffering) = response.headers.get("x-accel-buffering") {
            if let Ok(value) = buffering.to_str() {
                assert_eq!(value.to_lowercase(), "no");
            }
        }
    }

    #[test]
    fn test_brotli_preferred_over_gzip() {
        use crate::middleware::CompressionEncoding;

        // Verify that Brotli enum value comes before Gzip in definition order
        // This ensures best_supported_encoding() checks Brotli first
        let br = CompressionEncoding::Brotli;
        let gz = CompressionEncoding::Gzip;

        // Just verify both variants exist and are distinct
        assert_ne!(br, gz);
    }

    #[test]
    fn test_compression_encoding_equality() {
        use crate::middleware::CompressionEncoding;

        assert_eq!(CompressionEncoding::Brotli, CompressionEncoding::Brotli);
        assert_eq!(CompressionEncoding::Gzip, CompressionEncoding::Gzip);
        assert_ne!(CompressionEncoding::Brotli, CompressionEncoding::Gzip);
    }
}