warpdrive_proxy/proxy/

handler.rs

//! HTTP proxy handler implementation using Pingora
//!
//! This module provides the core ProxyHttp implementation for WarpDrive, handling:
//! - Upstream peer selection (forwarding to target server)
//! - Request/response filtering via middleware chain
//! - Error handling (502 Bad Gateway responses)

use async_trait::async_trait;
use bytes::Bytes;
use flate2::{Compression as FlateCompression, write::GzEncoder};
use pingora::http::ResponseHeader;
use pingora::prelude::*;
use pingora::proxy::FailToProxy;
use std::io::Write;
use std::sync::Arc;
use tokio::io::AsyncReadExt;
use tracing::{debug, error, info};

use crate::acme::ChallengeStore;
use crate::config::Config;
use crate::middleware::{
    CompressionEncoding, CompressionState, MiddlewareContext, MiddlewareStack, StaticResponseBody,
};
use crate::router::Router;

/// WarpDrive HTTP proxy handler
///
/// Implements Pingora's ProxyHttp trait to forward HTTP requests to an upstream server.
/// Uses a middleware stack for request/response processing.
/// Supports multi-upstream routing via optional Router.
pub struct WarpDriveProxy {
    /// Shared configuration
    config: Arc<Config>,
    /// Middleware stack for request/response processing
    middleware: MiddlewareStack,
    /// Optional router for multi-upstream routing (None = simple mode)
    router: Option<Router>,
    /// ACME challenge store for HTTP-01 validation
    challenge_store: ChallengeStore,
}

impl WarpDriveProxy {
    /// Create a new proxy handler with the given configuration
    pub fn new(
        config: Arc<Config>,
        router: Option<Router>,
        challenge_store: ChallengeStore,
    ) -> Self {
        info!("Initializing WarpDrive proxy handler");

        if router.is_some() {
            info!("  Mode: Advanced (TOML routing)");
        } else {
            info!("  Mode: Simple (env vars)");
            info!("  Target: {}:{}", config.target_host, config.target_port);
        }

        info!("  Forward headers: {}", config.forward_headers);
        info!("  X-Sendfile: {}", config.x_sendfile_enabled);
        info!("  Compression: {}", config.gzip_compression_enabled);
        info!("  HTTP/2 cleartext (h2c): {}", config.h2c_enabled);
        info!("  Logging: {}", config.log_requests);

        let middleware = MiddlewareStack::new(config.clone());

        Self {
            config,
            middleware,
            router,
            challenge_store,
        }
    }
}

#[async_trait]
impl ProxyHttp for WarpDriveProxy {
    type CTX = MiddlewareContext;

    fn new_ctx(&self) -> Self::CTX {
        MiddlewareContext::default()
    }

    /// Select the upstream peer (required method)
    ///
    /// This is called for each request to determine which upstream server to connect to.
    /// - Advanced mode (TOML): Routes based on path/host/method and selects from LoadBalancer
    /// - Simple mode (env): Always forwards to configured target_host:target_port
    async fn upstream_peer(
        &self,
        session: &mut Session,
        _ctx: &mut Self::CTX,
    ) -> Result<Box<HttpPeer>> {
        if let Some(router) = &self.router {
            // Advanced mode: Use router to select upstream
            let upstream = router
                .select_upstream(session)
                .map_err(|e| Error::because(ErrorType::HTTPStatus(502), e.to_string(), e))?;
            upstream
                .get_peer()
                .await
                .map_err(|e| Error::because(ErrorType::HTTPStatus(502), e.to_string(), e))
        } else {
            // Simple mode: Use configured target
            debug!(
                "Selecting upstream peer: {}:{}",
                self.config.target_host, self.config.target_port
            );

            Ok(Box::new(HttpPeer::new(
                (self.config.target_host.as_str(), self.config.target_port),
                false,         // HTTP (not HTTPS)
                String::new(), // No SNI for HTTP
            )))
        }
    }

    /// Apply request filters before sending to upstream
    ///
    /// This executes the middleware chain on the incoming request.
    /// Returns false to indicate request should not be bypassed.
    async fn request_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<bool> {
        // Handle ACME HTTP-01 challenges
        let path = session.req_header().uri.path();
        if let Some(token) = path.strip_prefix("/.well-known/acme-challenge/") {
            if let Some(key_auth) = self.challenge_store.get(token).await {
                debug!("Serving ACME challenge for token: {}", token);

                // Respond with key authorization
                let mut response = ResponseHeader::build(200, None)?;
                response.insert_header("Content-Type", "text/plain")?;
                response.insert_header("Content-Length", key_auth.len().to_string())?;

                session
                    .write_response_header(Box::new(response), false)
                    .await?;
                session
                    .write_response_body(Some(Bytes::from(key_auth)), true)
                    .await?;

                return Ok(true); // Bypass - we handled the request
            } else {
                debug!("ACME challenge token not found: {}", token);
                // Let it pass through - might be handled by upstream
            }
        }

        // Validate Host/:authority header based on HTTP version
        // HTTP/1.1: Requires "Host" header
        // HTTP/2+: Uses ":authority" pseudo-header (no "Host" header)
        let version = session.req_header().version;
        if version == http::Version::HTTP_11 || version == http::Version::HTTP_10 {
            // HTTP/1.x requires Host header
            if session.req_header().headers.get("Host").is_none() {
                if let Err(err) = session.respond_error(400).await {
                    error!(
                        "Failed to send 400 response for missing Host header: {}",
                        err
                    );
                    return Err(err);
                }
                return Ok(true); // Bypass request - we already responded
            }
        }
        // HTTP/2+ uses :authority pseudo-header - handled by Pingora internally

        // Apply middleware request filters
        self.middleware.apply_request_filters(session, ctx).await?;

        // Check if static files middleware set a response
        if let Some(static_response) = ctx.static_response.take() {
            debug!("Serving static file response");

            session
                .write_response_header(Box::new(static_response.header), false)
                .await?;

            match static_response.body {
                StaticResponseBody::InMemory(body) => {
                    session.write_response_body(Some(body), true).await?;
                }
                StaticResponseBody::Stream(path) => {
                    let mut file = match tokio::fs::File::open(&path).await {
                        Ok(file) => file,
                        Err(err) => {
                            error!(
                                "Failed to open static file for streaming {:?}: {}",
                                path, err
                            );
                            return Err(Error::explain(
                                ErrorType::HTTPStatus(500),
                                "Failed to open static file",
                            ));
                        }
                    };

                    let mut buffer = vec![0u8; 64 * 1024];
                    loop {
                        let bytes_read = file.read(&mut buffer).await.map_err(|err| {
                            error!("Failed to read static file chunk {:?}: {}", path, err);
                            Error::explain(ErrorType::HTTPStatus(500), "Failed to read static file")
                        })?;

                        if bytes_read == 0 {
                            break;
                        }

                        session
                            .write_response_body(
                                Some(Bytes::copy_from_slice(&buffer[..bytes_read])),
                                false,
                            )
                            .await?;
                    }

                    session.write_response_body(None, true).await?;
                }
            }

            return Ok(true); // Bypass - we handled the request
        }

        // Don't bypass request
        Ok(false)
    }

    /// Modify request before sending to upstream (optional)
    ///
    /// This preserves the original Host header and applies path transformation if configured.
    async fn upstream_request_filter(
        &self,
        session: &mut Session,
        upstream_request: &mut RequestHeader,
        _ctx: &mut Self::CTX,
    ) -> Result<()> {
        // Preserve the original Host header for upstream request
        // Note: Host header validation happens earlier in request_filter()
        if let Some(host) = session.req_header().headers.get("Host") {
            // Convert header value to string
            // Invalid headers are handled by returning 400 in request_filter
            let host_str = host.to_str().map_err(|_| {
                Error::explain(
                    ErrorType::HTTPStatus(400),
                    "Invalid Host header (non-ASCII characters)",
                )
            })?;
            upstream_request.insert_header("Host", host_str)?;
        }

        // Apply path transformation if using router
        if let Some(router) = &self.router {
            if let Some(route) = router.find_matching_route(session) {
                let original_path = session.req_header().uri.path();
                let transformed_path = route.transform_path(original_path);

                if transformed_path != original_path {
                    debug!(
                        "Transforming path: {} -> {}",
                        original_path, transformed_path
                    );

                    // Build new URI with transformed path
                    let mut parts = upstream_request.uri.clone().into_parts();
                    let path_and_query = if let Some(query) = session.req_header().uri.query() {
                        format!("{}?{}", transformed_path, query)
                    } else {
                        transformed_path.to_string()
                    };

                    parts.path_and_query = Some(path_and_query.parse().map_err(|_| {
                        Error::explain(
                            ErrorType::HTTPStatus(500),
                            "Failed to construct transformed URI",
                        )
                    })?);

                    upstream_request.set_uri(http::Uri::from_parts(parts).map_err(|_| {
                        Error::explain(
                            ErrorType::HTTPStatus(500),
                            "Failed to apply path transformation",
                        )
                    })?);
                }
            }
        }

        debug!(
            "Forwarding request: {} {} to upstream",
            upstream_request.method, upstream_request.uri
        );

        Ok(())
    }

    /// Apply response filters after receiving from upstream
    ///
    /// This executes the middleware chain on the upstream response.
    async fn response_filter(
        &self,
        session: &mut Session,
        upstream_response: &mut ResponseHeader,
        ctx: &mut Self::CTX,
    ) -> Result<()> {
        self.middleware
            .apply_response_filters(session, upstream_response, ctx)
            .await
    }

    /// Handle upstream connection errors (optional)
    ///
    /// This is called when we fail to connect to the upstream server.
    /// We log the error and return it (Pingora will send 502 Bad Gateway).
    fn fail_to_connect(
        &self,
        _session: &mut Session,
        _peer: &HttpPeer,
        _ctx: &mut Self::CTX,
        e: Box<Error>,
    ) -> Box<Error> {
        error!("Failed to connect to upstream: {}", e);
        // TODO: In future, serve custom bad gateway page from config.bad_gateway_page
        e
    }

    /// Handle upstream errors during request/response (optional)
    ///
    /// This is called when the proxy operation fails after connection.
    /// We can customize the error response and connection reuse behavior.
    async fn fail_to_proxy(
        &self,
        _session: &mut Session,
        e: &Error,
        _ctx: &mut Self::CTX,
    ) -> FailToProxy {
        error!("Failed to proxy request: {}", e);

        // Determine error code based on error type
        let error_code = if e.etype() == &ErrorType::ReadError {
            // Could be request entity too large
            // TODO: Better error type detection
            413
        } else {
            // Return 502 Bad Gateway for other errors
            502
        };

        // Don't reuse downstream connection on proxy errors
        // to avoid sending partial/corrupted responses
        FailToProxy {
            error_code,
            can_reuse_downstream: false,
        }
    }

    /// Process response body chunks for compression/sendfile support
    fn response_body_filter(
        &self,
        _session: &mut Session,
        body: &mut Option<Bytes>,
        end_of_stream: bool,
        ctx: &mut Self::CTX,
    ) -> Result<Option<std::time::Duration>> {
        // X-Sendfile has highest priority: replace upstream body with file contents
        if ctx.sendfile.active {
            if let Some(chunk) = body {
                chunk.clear();
            }

            if !ctx.sendfile.served {
                if let Some(file_body) = ctx.sendfile.body.take() {
                    *body = Some(file_body);
                } else {
                    *body = None;
                }
                ctx.sendfile.served = true;
            } else {
                *body = None;
            }

            return Ok(None);
        }

        // Streaming responses: pass through immediately without buffering
        // This enables real-time SSE, WebSockets, and other streaming protocols
        if ctx.streaming {
            // Body passes through unmodified - no buffering
            return Ok(None);
        }

        // Compression: buffer upstream body until end-of-stream then emit compressed bytes
        if let CompressionState::Pending { buffer, encoding } = &mut ctx.compression {
            if let Some(chunk) = body {
                buffer.extend_from_slice(&chunk[..]);
                chunk.clear();
            }

            if end_of_stream {
                let compressed = match encoding {
                    CompressionEncoding::Brotli => brotli_compress(buffer)?,
                    CompressionEncoding::Gzip => gzip_compress(buffer)?,
                };
                *body = Some(Bytes::from(compressed));
                ctx.compression = CompressionState::Complete;
            } else {
                *body = None;
            }

            return Ok(None);
        }

        Ok(None)
    }
}

// Note: X-Forwarded-* header logic is now handled by HeadersMiddleware

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

    #[test]
    fn test_proxy_creation() {
        let config = Arc::new(Config::default());
        let challenge_store = ChallengeStore::default();
        let proxy = WarpDriveProxy::new(config.clone(), None, challenge_store);

        assert_eq!(proxy.config.target_port, 3000);
    }
}

fn gzip_compress(buffer: &[u8]) -> Result<Vec<u8>> {
    let mut encoder = GzEncoder::new(
        Vec::with_capacity(buffer.len() / 2 + 16),
        FlateCompression::default(),
    );
    encoder.write_all(buffer).map_err(|_| {
        Error::explain(
            ErrorType::HTTPStatus(500),
            "Failed to compress response body",
        )
    })?;
    encoder.finish().map_err(|_| {
        Error::explain(
            ErrorType::HTTPStatus(500),
            "Failed to finalize compressed body",
        )
    })
}

fn brotli_compress(buffer: &[u8]) -> Result<Vec<u8>> {
    use brotli::enc::BrotliEncoderParams;

    let mut output = Vec::with_capacity(buffer.len() / 2 + 16);
    let params = BrotliEncoderParams {
        quality: 6, // Brotli quality 6 (balance of speed vs size, similar to gzip default)
        ..Default::default()
    };

    brotli::BrotliCompress(&mut std::io::Cursor::new(buffer), &mut output, &params).map_err(
        |_| {
            Error::explain(
                ErrorType::HTTPStatus(500),
                "Failed to compress response body with Brotli",
            )
        },
    )?;

    Ok(output)
}