warpdrive_proxy/middleware/

headers.rs

//! Headers middleware
//!
//! Manages X-Forwarded-* headers to preserve client information through the proxy.
//! This is critical for applications that need to know the original client IP,
//! protocol, and host when behind a reverse proxy.
//!
//! # Headers Managed
//!
//! - **X-Forwarded-For**: Client IP address chain
//! - **X-Forwarded-Proto**: Original protocol (http/https)
//! - **X-Forwarded-Host**: Original host header
//! - **X-Real-IP**: Direct client IP (non-standard but widely used)
//!
//! # Behavior
//!
//! When `forward_headers` is true:
//! - Preserves existing X-Forwarded-* headers from client
//! - Appends to X-Forwarded-For chain
//!
//! When `forward_headers` is false:
//! - Replaces X-Forwarded-* headers with current request values
//! - Starts new X-Forwarded-For chain with client IP

use async_trait::async_trait;
use pingora::prelude::*;
use std::sync::Arc;
use tracing::debug;

use super::{Middleware, MiddlewareContext};
use crate::config::Config;

/// Headers middleware
///
/// Injects and manages X-Forwarded-* headers for proxy transparency.
pub struct HeadersMiddleware {
    pub(crate) config: Arc<Config>,
}

impl HeadersMiddleware {
    /// Create new headers middleware
    pub fn new(config: Arc<Config>) -> Self {
        Self { config }
    }

    /// Get request protocol (http/https)
    ///
    /// Detects if the downstream connection is using TLS by checking the session digest.
    /// This is critical for X-Forwarded-Proto header accuracy, which apps use for
    /// building absolute URLs (redirects, asset URLs, etc.).
    fn get_protocol(session: &Session) -> &'static str {
        // Check URI scheme first (for explicit https:// requests)
        if let Some(scheme) = session.req_header().uri.scheme_str() {
            if scheme == "https" {
                return "https";
            }
        }

        // Check if downstream connection is TLS via session digest
        // See: https://github.com/cloudflare/pingora/issues/403
        if let Some(digest) = session.digest() {
            if digest.ssl_digest.is_some() {
                return "https";
            }
        }

        "http"
    }
}

#[async_trait]
impl Middleware for HeadersMiddleware {
    /// Add X-Forwarded-* headers to upstream request
    ///
    /// This preserves client information when proxying requests.
    async fn request_filter(
        &self,
        session: &mut Session,
        ctx: &mut MiddlewareContext,
    ) -> Result<()> {
        let forward_headers = self.config.forward_headers;

        // Use real client IP from context (normalized by trusted_ranges middleware)
        let client_ip = ctx.real_client_ip.to_string();

        // Extract all header values we need before mutating
        let existing_xff = session
            .req_header()
            .headers
            .get("x-forwarded-for")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string());

        let existing_proto = session
            .req_header()
            .headers
            .get("x-forwarded-proto")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string());

        let existing_host = session
            .req_header()
            .headers
            .get("x-forwarded-host")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string());

        let current_host = session
            .req_header()
            .headers
            .get("host")
            .and_then(|v| v.to_str().ok())
            .map(|s| s.to_string());

        let protocol = Self::get_protocol(session);

        // Now we can mutate the session safely

        // Handle X-Forwarded-For
        let mut xff_value = if forward_headers { existing_xff } else { None };

        // Append client IP to X-Forwarded-For chain
        xff_value = Some(match xff_value {
            Some(existing) => format!("{}, {}", existing, client_ip),
            None => client_ip.clone(),
        });

        // Set X-Forwarded-For header
        if let Some(xff) = xff_value {
            session
                .req_header_mut()
                .insert_header("X-Forwarded-For", &xff)?;
            debug!("Set X-Forwarded-For: {}", xff);
        }

        // Handle X-Forwarded-Proto
        let proto = if forward_headers {
            existing_proto.as_deref().unwrap_or(protocol)
        } else {
            protocol
        };

        session
            .req_header_mut()
            .insert_header("X-Forwarded-Proto", proto)?;
        debug!("Set X-Forwarded-Proto: {}", proto);

        // Handle X-Forwarded-Host
        let forwarded_host = if forward_headers {
            existing_host.or(current_host)
        } else {
            current_host
        };

        if let Some(host) = forwarded_host {
            session
                .req_header_mut()
                .insert_header("X-Forwarded-Host", &host)?;
            debug!("Set X-Forwarded-Host: {}", host);
        }

        // Set X-Real-IP (non-standard but widely used)
        session
            .req_header_mut()
            .insert_header("X-Real-IP", &client_ip)?;
        debug!("Set X-Real-IP: {}", client_ip);

        Ok(())
    }
}

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

    #[test]
    fn test_headers_middleware_creation() {
        let config = Arc::new(Config::default());
        let middleware = HeadersMiddleware::new(config.clone());

        assert_eq!(middleware.config.forward_headers, config.forward_headers);
    }

    #[test]
    fn test_xff_chain_building() {
        // Test X-Forwarded-For chain building logic
        let existing = Some("10.0.0.1, 10.0.0.2".to_string());
        let new_ip = "192.168.1.1";

        let result = match existing {
            Some(existing) => format!("{}, {}", existing, new_ip),
            None => new_ip.to_string(),
        };

        assert_eq!(result, "10.0.0.1, 10.0.0.2, 192.168.1.1");

        // Test with no existing
        let existing: Option<String> = None;
        let result = match existing {
            Some(existing) => format!("{}, {}", existing, new_ip),
            None => new_ip.to_string(),
        };

        assert_eq!(result, "192.168.1.1");
    }
}