mockforge_proxy/

server.rs

//! Browser/Mobile Proxy Server
//!
//! Provides an intercepting proxy for frontend/mobile clients with HTTPS support,
//! certificate injection, and comprehensive request/response logging.

use crate::{body_transform::BodyTransformationMiddleware, config::ProxyConfig};
use axum::{
    extract::Request, http::StatusCode, middleware::Next, response::Response, routing::get, Router,
};
use serde::Serialize;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

/// Proxy server state
pub struct ProxyServer {
    /// Proxy configuration
    config: Arc<RwLock<ProxyConfig>>,
    /// Request logging enabled
    log_requests: bool,
    /// Response logging enabled
    log_responses: bool,
    /// Request counter for logging
    request_counter: Arc<RwLock<u64>>,
    /// Server start time for uptime and rate calculations
    start_time: std::time::Instant,
    /// Total response time in milliseconds for average calculation
    total_response_time_ms: Arc<RwLock<u64>>,
    /// Error counter for error rate calculation
    error_counter: Arc<RwLock<u64>>,
    /// #864 — optional passive conformance tap over a loaded spec.
    /// Built from MOCKFORGE_PROXY_SPEC + MOCKFORGE_PROXY_VALIDATE_CONFORMANCE.
    conformance_tap: Option<crate::conformance::ConformanceTap>,
}

impl std::fmt::Debug for ProxyServer {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProxyServer")
            .field("log_requests", &self.log_requests)
            .field("log_responses", &self.log_responses)
            .field("conformance_tap", &self.conformance_tap.as_ref().map(|_| "configured"))
            .finish()
    }
}

impl ProxyServer {
    /// Create a new proxy server
    pub fn new(config: ProxyConfig, log_requests: bool, log_responses: bool) -> Self {
        // #864 — opt-in conformance tap. Requires a spec (env or config)
        // and the validate flag; strict mode additionally rejects.
        let want_validation = std::env::var("MOCKFORGE_PROXY_VALIDATE_CONFORMANCE")
            .map(|v| v.eq_ignore_ascii_case("true") || v == "1")
            .unwrap_or(false);
        let conformance_tap = if want_validation {
            let strict = std::env::var("MOCKFORGE_PROXY_VALIDATE_CONFORMANCE_STRICT")
                .map(|v| v.eq_ignore_ascii_case("true") || v == "1")
                .unwrap_or(false);
            match std::env::var("MOCKFORGE_PROXY_SPEC") {
                Ok(path) if !path.trim().is_empty() => {
                    match std::fs::read_to_string(&path).map_err(|e| e.to_string()).and_then(
                        |raw| {
                            serde_json::from_str::<serde_json::Value>(&raw)
                                .or_else(|_| serde_yaml::from_str::<serde_json::Value>(&raw))
                                .map_err(|e| e.to_string())
                        },
                    ) {
                        Ok(spec_value) => {
                            match crate::conformance::ConformanceTap::from_spec_value(
                                spec_value, strict,
                            ) {
                                Ok(tap) => {
                                    info!(
                                        spec = %path,
                                        strict,
                                        "proxy conformance validation enabled (#864)"
                                    );
                                    Some(tap)
                                }
                                Err(e) => {
                                    warn!(spec = %path, error = %e, "conformance tap disabled");
                                    None
                                }
                            }
                        }
                        Err(e) => {
                            warn!(spec = %path, error = %e, "cannot parse conformance spec; tap disabled");
                            None
                        }
                    }
                }
                _ => {
                    warn!("MOCKFORGE_PROXY_VALIDATE_CONFORMANCE set but MOCKFORGE_PROXY_SPEC missing; tap disabled");
                    None
                }
            }
        } else {
            None
        };
        Self {
            config: Arc::new(RwLock::new(config)),
            log_requests,
            log_responses,
            request_counter: Arc::new(RwLock::new(0)),
            start_time: std::time::Instant::now(),
            total_response_time_ms: Arc::new(RwLock::new(0)),
            error_counter: Arc::new(RwLock::new(0)),
            conformance_tap,
        }
    }

    /// Get the Axum router for the proxy server
    pub fn router(self) -> Router {
        let state = Arc::new(self);
        let state_for_middleware = state.clone();

        Router::new()
            // Health check endpoint
            .route("/proxy/health", get(health_check))
            // Catch-all proxy handler - use fallback for all methods
            .fallback(proxy_handler)
            .with_state(state)
            .layer(axum::middleware::from_fn_with_state(state_for_middleware, logging_middleware))
    }
}

/// Health check endpoint for the proxy
async fn health_check() -> Result<Response<String>, StatusCode> {
    // Response builder should never fail with known-good values, but handle errors gracefully
    Response::builder()
        .status(StatusCode::OK)
        .header("Content-Type", "application/json")
        .body(r#"{"status":"healthy","service":"mockforge-proxy"}"#.to_string())
        .map_err(|e| {
            tracing::error!("Failed to build health check response: {}", e);
            StatusCode::INTERNAL_SERVER_ERROR
        })
}

/// Main proxy handler that intercepts and forwards requests
async fn proxy_handler(
    axum::extract::State(state): axum::extract::State<Arc<ProxyServer>>,
    request: http::Request<axum::body::Body>,
) -> Result<Response<String>, StatusCode> {
    // Extract client address from request extensions (set by ConnectInfo middleware)
    let client_addr = request
        .extensions()
        .get::<SocketAddr>()
        .copied()
        .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));

    let method = request.method().clone();
    let uri = request.uri().clone();
    let headers = request.headers().clone();

    // Read request body early for conditional evaluation (consume the body)
    let body_bytes = match axum::body::to_bytes(request.into_body(), usize::MAX).await {
        Ok(bytes) => Some(bytes.to_vec()),
        Err(e) => {
            error!("Failed to read request body: {}", e);
            None
        }
    };

    let config = state.config.read().await;

    // Check if proxy is enabled
    if !config.enabled {
        return Err(StatusCode::SERVICE_UNAVAILABLE);
    }

    // Determine if this request should be proxied (with conditional evaluation)
    if !config.should_proxy_with_condition(&method, &uri, &headers, body_bytes.as_deref()) {
        return Err(StatusCode::NOT_FOUND);
    }

    // Get the stripped path (without proxy prefix)
    let stripped_path = config.strip_prefix(uri.path());

    // Get the base upstream URL and construct the full URL
    let base_upstream_url = config.get_upstream_url(uri.path());
    // strip_prefix re-adds a leading slash, so tolerate both `http://…`
    // and `/http://…` spellings before testing for an absolute URL.
    let candidate = stripped_path.trim_start_matches('/');
    let is_absolute_upstream =
        candidate.starts_with("http://") || candidate.starts_with("https://");

    // #1012 / MF-002: an absolute URL embedded in the request path used to be
    // forwarded verbatim — an open proxy/SSRF vector. It now requires the
    // explicit `allow_absolute_url_upstream` opt-in, and even then passes the
    // egress guard (denylisted IPs/metadata hosts, DNS re-check).
    if is_absolute_upstream && !config.allow_absolute_url_upstream {
        warn!(
            path = %uri.path(),
            "Rejected absolute-URL-in-path proxying; set allow_absolute_url_upstream=true to opt in (#1012)"
        );
        let body =
            "Forbidden: absolute-URL proxying is disabled (allow_absolute_url_upstream=false)"
                .to_string();
        return Ok(Response::builder()
            .status(StatusCode::FORBIDDEN)
            .body(body)
            .unwrap_or_else(|_| Response::new(String::new())));
    }

    let full_upstream_url = if is_absolute_upstream {
        use crate::egress::{EgressDecision, EgressGuard};
        let guard = EgressGuard::new(config.upstream_allowlist.clone());
        match guard.check(candidate).await {
            EgressDecision::Allowed => candidate.to_string(),
            EgressDecision::Blocked(reason) => {
                warn!(target = %stripped_path, %reason, "Blocked request-derived upstream by egress guard");
                let body = format!("Forbidden: upstream blocked by egress guard ({})", reason);
                return Ok(Response::builder()
                    .status(StatusCode::FORBIDDEN)
                    .body(body)
                    .unwrap_or_else(|_| Response::new(String::new())));
            }
        }
    } else {
        let base = base_upstream_url.trim_end_matches('/');
        let path = stripped_path.trim_start_matches('/');
        let query = uri.query().map(|q| format!("?{}", q)).unwrap_or_default();
        if path.is_empty() || path == "/" {
            format!("{}{}", base, query)
        } else {
            format!("{}/{}", base, path) + &query
        }
    };

    // Create a new URI with the full upstream URL for the proxy handler
    let _modified_uri = full_upstream_url.parse::<http::Uri>().unwrap_or_else(|_| uri.clone());

    // Log the request if enabled
    if state.log_requests {
        let mut counter = state.request_counter.write().await;
        *counter += 1;
        let request_id = *counter;

        info!(
            request_id = request_id,
            method = %method,
            path = %uri.path(),
            upstream = %full_upstream_url,
            client_ip = %client_addr.ip(),
            "Proxy request intercepted"
        );
    }

    // Convert headers to HashMap for the proxy handler
    let mut header_map = std::collections::HashMap::new();
    for (key, value) in &headers {
        if let Ok(value_str) = value.to_str() {
            header_map.insert(key.to_string(), value_str.to_string());
        }
    }

    // Use ProxyClient directly with the full upstream URL to bypass ProxyHandler's URL construction
    use crate::client::ProxyClient;
    let proxy_client = ProxyClient::new();

    // Convert method to reqwest method
    let reqwest_method = match method.as_str() {
        "GET" => reqwest::Method::GET,
        "POST" => reqwest::Method::POST,
        "PUT" => reqwest::Method::PUT,
        "DELETE" => reqwest::Method::DELETE,
        "HEAD" => reqwest::Method::HEAD,
        "OPTIONS" => reqwest::Method::OPTIONS,
        "PATCH" => reqwest::Method::PATCH,
        _ => {
            error!("Unsupported HTTP method: {}", method);
            return Err(StatusCode::METHOD_NOT_ALLOWED);
        }
    };

    // Add any configured headers
    for (key, value) in &config.headers {
        header_map.insert(key.clone(), value.clone());
    }

    // Apply request body transformations if configured
    let mut transformed_request_body = body_bytes.clone();
    if !config.request_replacements.is_empty() {
        let transform_middleware = BodyTransformationMiddleware::new(
            config.request_replacements.clone(),
            Vec::new(), // No response rules needed here
        );
        if let Err(e) =
            transform_middleware.transform_request_body(uri.path(), &mut transformed_request_body)
        {
            warn!("Failed to transform request body: {}", e);
            // Continue with original body if transformation fails
        }
    }

    // #864 — passive request validation against the loaded spec. The
    // findings land in the shared conformance buffer (bench/TUI/admin read
    // the same store). In STRICT mode a violation rejects with the spec
    // status instead of forwarding.
    if let Some(tap) = &state.conformance_tap {
        if let Some((status, payload)) = tap
            .validate_request(
                method.as_str(),
                uri.path(),
                uri.query(),
                &headers,
                Some(transformed_request_body.as_deref().unwrap_or(&[])),
            )
            .await
        {
            warn!(status = status, path = %uri.path(), "strict conformance rejection");
            return Response::builder()
                .status(status)
                .body(payload.to_string())
                .map_err(|_| StatusCode::BAD_GATEWAY);
        }
    }

    match proxy_client
        .send_request(
            reqwest_method,
            &full_upstream_url,
            &header_map,
            transformed_request_body.as_deref(),
        )
        .await
    {
        Ok(response) => {
            let status = StatusCode::from_u16(response.status().as_u16())
                .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);

            // Log the response if enabled
            if state.log_responses {
                info!(
                    method = %method,
                    path = %uri.path(),
                    status = status.as_u16(),
                    "Proxy response sent"
                );
            }

            // Convert response headers
            let mut response_headers = http::HeaderMap::new();
            for (name, value) in response.headers() {
                if let (Ok(header_name), Ok(header_value)) = (
                    http::HeaderName::try_from(name.as_str()),
                    http::HeaderValue::try_from(value.as_bytes()),
                ) {
                    response_headers.insert(header_name, header_value);
                }
            }

            // Read response body
            let response_body_bytes = response.bytes().await.map_err(|e| {
                error!("Failed to read proxy response body: {}", e);
                StatusCode::BAD_GATEWAY
            })?;

            // Apply response body transformations if configured
            let mut final_body_bytes = response_body_bytes.to_vec();
            {
                let config_for_response = state.config.read().await;
                if !config_for_response.response_replacements.is_empty() {
                    let transform_middleware = BodyTransformationMiddleware::new(
                        Vec::new(), // No request rules needed here
                        config_for_response.response_replacements.clone(),
                    );
                    let mut body_option = Some(final_body_bytes.clone());
                    if let Err(e) = transform_middleware.transform_response_body(
                        uri.path(),
                        status.as_u16(),
                        &mut body_option,
                    ) {
                        warn!("Failed to transform response body: {}", e);
                        // Continue with original body if transformation fails
                    } else if let Some(transformed_body) = body_option {
                        final_body_bytes = transformed_body;
                    }
                }
            }

            // #864 — observational response validation against the spec.
            if let Some(tap) = &state.conformance_tap {
                tap.validate_response(
                    method.as_str(),
                    uri.path(),
                    status.as_u16(),
                    Some(final_body_bytes.as_slice()),
                )
                .await;
            }

            let body_string = String::from_utf8_lossy(&final_body_bytes).to_string();

            // Build Axum response
            let mut response_builder = Response::builder().status(status);
            for (name, value) in response_headers.iter() {
                response_builder = response_builder.header(name, value);
            }

            response_builder
                .body(body_string)
                .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
        }
        Err(e) => {
            error!("Proxy request failed: {}", e);
            Err(StatusCode::BAD_GATEWAY)
        }
    }
}

/// Middleware for logging requests and responses
async fn logging_middleware(
    axum::extract::State(state): axum::extract::State<Arc<ProxyServer>>,
    request: Request,
    next: Next,
) -> Response {
    let start = std::time::Instant::now();
    let method = request.method().clone();
    let uri = request.uri().clone();

    // Extract client address from request extensions
    let client_addr = request
        .extensions()
        .get::<SocketAddr>()
        .copied()
        .unwrap_or_else(|| SocketAddr::from(([0, 0, 0, 0], 0)));

    debug!(
        method = %method,
        uri = %uri,
        client_ip = %client_addr.ip(),
        "Request received"
    );

    let response = next.run(request).await;
    let duration = start.elapsed();

    // Track response time
    {
        let mut total_time = state.total_response_time_ms.write().await;
        *total_time += duration.as_millis() as u64;
    }

    // Track server errors
    if response.status().is_server_error() {
        let mut errors = state.error_counter.write().await;
        *errors += 1;
    }

    debug!(
        method = %method,
        uri = %uri,
        status = %response.status(),
        duration_ms = duration.as_millis(),
        "Response sent"
    );

    response
}

/// Proxy statistics for monitoring
#[derive(Debug, Serialize)]
pub struct ProxyStats {
    /// Total requests processed
    pub total_requests: u64,
    /// Requests per second
    pub requests_per_second: f64,
    /// Average response time in milliseconds
    pub avg_response_time_ms: f64,
    /// Error rate percentage
    pub error_rate_percent: f64,
}

/// Get proxy statistics
pub async fn get_proxy_stats(state: &ProxyServer) -> ProxyStats {
    let total_requests = *state.request_counter.read().await;
    let total_response_time_ms = *state.total_response_time_ms.read().await;
    let error_count = *state.error_counter.read().await;

    let elapsed_secs = state.start_time.elapsed().as_secs_f64();
    let requests_per_second = if elapsed_secs > 0.0 {
        total_requests as f64 / elapsed_secs
    } else {
        0.0
    };

    let avg_response_time_ms = if total_requests > 0 {
        total_response_time_ms as f64 / total_requests as f64
    } else {
        0.0
    };

    let error_rate_percent = if total_requests > 0 {
        (error_count as f64 / total_requests as f64) * 100.0
    } else {
        0.0
    };

    ProxyStats {
        total_requests,
        requests_per_second,
        avg_response_time_ms,
        error_rate_percent,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::ProxyConfig;
    use axum::http::StatusCode;

    #[tokio::test]
    async fn test_proxy_server_creation() {
        let config = ProxyConfig::default();
        let server = ProxyServer::new(config, true, true);

        // Test that the server can be created
        assert!(server.log_requests);
        assert!(server.log_responses);
    }

    #[tokio::test]
    async fn test_absolute_url_in_path_rejected_by_default() {
        let config = ProxyConfig {
            enabled: true,
            prefix: Some("/proxy/".to_string()),
            ..Default::default()
        };
        // allow_absolute_url_upstream defaults to false (#1012)
        let server = Arc::new(ProxyServer::new(config, false, false));

        let request = http::Request::builder()
            .method("GET")
            .uri("/proxy/http://169.254.169.254/latest/meta-data/")
            .body(axum::body::Body::empty())
            .unwrap();
        let response = proxy_handler(axum::extract::State(server), request).await.unwrap();

        assert_eq!(response.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn test_opt_in_absolute_url_still_blocked_for_metadata_ip() {
        let config = ProxyConfig {
            enabled: true,
            prefix: Some("/proxy/".to_string()),
            allow_absolute_url_upstream: true,
            ..Default::default()
        };
        let server = Arc::new(ProxyServer::new(config, false, false));

        let request = http::Request::builder()
            .method("GET")
            .uri("/proxy/http://169.254.169.254/latest/meta-data/")
            .body(axum::body::Body::empty())
            .unwrap();
        let response = proxy_handler(axum::extract::State(server), request).await.unwrap();

        assert_eq!(response.status(), StatusCode::FORBIDDEN);
        assert!(response.into_body().contains("egress guard"));
    }

    #[tokio::test]
    async fn test_opt_in_absolute_url_blocked_for_loopback() {
        let config = ProxyConfig {
            enabled: true,
            prefix: Some("/proxy/".to_string()),
            allow_absolute_url_upstream: true,
            ..Default::default()
        };
        let server = Arc::new(ProxyServer::new(config, false, false));

        let request = http::Request::builder()
            .method("GET")
            .uri("/proxy/http://127.0.0.1:9090/admin")
            .body(axum::body::Body::empty())
            .unwrap();
        let response = proxy_handler(axum::extract::State(server), request).await.unwrap();

        assert_eq!(response.status(), StatusCode::FORBIDDEN);
    }

    #[tokio::test]
    async fn test_allowlisted_private_host_bypasses_guard() {
        use crate::egress::{EgressDecision, EgressGuard, UpstreamAllowlist};

        // Explicit operator opt-in re-opens SSRF by design; verify it wins.
        let allowlist = UpstreamAllowlist {
            url_prefixes: vec!["http://internal.local".to_string()],
            hosts: Vec::new(),
        };
        let guard = EgressGuard::new(Some(allowlist));
        match guard.check_without_dns("http://internal.local/v1") {
            EgressDecision::Allowed => {}
            EgressDecision::Blocked(e) => panic!("allowlisted host blocked: {}", e),
        }
    }

    #[tokio::test]
    async fn test_health_check() {
        let response = health_check().await.unwrap();
        assert_eq!(response.status(), StatusCode::OK);

        // Response body is already a String
        let body = response.into_body();

        assert!(body.contains("healthy"));
        assert!(body.contains("mockforge-proxy"));
    }

    #[tokio::test]
    async fn test_proxy_stats() {
        let config = ProxyConfig::default();
        let server = ProxyServer::new(config, false, false);

        // With zero requests, all derived stats should be zero
        let stats = get_proxy_stats(&server).await;
        assert_eq!(stats.total_requests, 0);
        assert_eq!(stats.avg_response_time_ms, 0.0);
        assert_eq!(stats.error_rate_percent, 0.0);
    }
}