turbomcp-proxy 3.5.0

//! Runtime proxy layer (dynamic, no code generation)
//!
//! This module provides dynamic proxying capabilities without code generation.
//! Ideal for development, testing, and prototyping.

//! # Security Features
//!
//! - Command injection protection via allowlist
//! - SSRF protection for HTTP backends
//! - Path traversal protection
//! - Request size limits
//! - Timeout enforcement
//!
//! # Example
//!
//! ```no_run
//! # use turbomcp_proxy::runtime::{RuntimeProxyBuilder, RuntimeProxy};
//! # use turbomcp_proxy::config::{BackendConfig, FrontendType};
//! # async fn example() -> turbomcp_proxy::ProxyResult<()> {
//! let proxy = RuntimeProxyBuilder::new()
//!     .with_stdio_backend("python", vec!["server.py".to_string()])
//!     .with_http_frontend("127.0.0.1:3000")
//!     .build()
//!     .await?;
//!
//! // proxy.run().await?;
//! # Ok(())
//! # }
//! ```

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use tracing::{debug, warn};

use crate::config::{BackendConfig, BackendValidationConfig, FrontendType, SsrfProtection};
use crate::error::{ProxyError, ProxyResult};
use crate::proxy::{AtomicMetrics, BackendConnector, BackendTransport, ProxyService};
use ipnetwork::IpNetwork;

mod security;
pub use security::{OriginAllowlist, build_cors_layer, origin_guard};

/// Maximum request size in bytes (10 MB)
pub const MAX_REQUEST_SIZE: usize = 10 * 1024 * 1024;

/// Default timeout in milliseconds (30 seconds)
pub const DEFAULT_TIMEOUT_MS: u64 = 30_000;

/// Maximum timeout in milliseconds (5 minutes)
pub const MAX_TIMEOUT_MS: u64 = 300_000;

/// Allowed commands for STDIO backends (security allowlist)
///
/// Only these commands are permitted to prevent command injection attacks.
/// Add new commands here with careful security review.
pub const ALLOWED_COMMANDS: &[&str] = &["python", "python3", "node", "deno", "uv", "npx", "bun"];

/// Secure default bind address (localhost only)
pub const DEFAULT_BIND_ADDRESS: &str = "127.0.0.1:3000";

/// Runtime proxy builder following `TurboMCP` builder pattern
///
/// Provides a fluent API for constructing runtime proxies with:
/// - Comprehensive security validation
/// - Sensible defaults
/// - Type-safe configuration
#[derive(Debug)]
pub struct RuntimeProxyBuilder {
    backend_config: Option<BackendConfig>,
    frontend_type: Option<FrontendType>,
    bind_address: Option<String>,
    request_size_limit: usize,
    timeout_ms: u64,
    enable_metrics: bool,
    validation_config: BackendValidationConfig,
    /// Browser origins allowed to reach the HTTP/WebSocket frontend. Empty by
    /// default → any request carrying an `Origin` header is rejected with 403.
    allowed_origins: Vec<String>,
}

impl RuntimeProxyBuilder {
    /// Create a new runtime proxy builder with secure defaults
    #[must_use]
    pub fn new() -> Self {
        Self {
            backend_config: None,
            frontend_type: None,
            bind_address: Some(DEFAULT_BIND_ADDRESS.to_string()),
            request_size_limit: MAX_REQUEST_SIZE,
            timeout_ms: DEFAULT_TIMEOUT_MS,
            enable_metrics: true,
            validation_config: BackendValidationConfig::default(),
            allowed_origins: Vec::new(),
        }
    }

    /// Configure the list of browser origins permitted to reach the HTTP /
    /// WebSocket frontend.
    ///
    /// Each entry is matched verbatim against the request's `Origin` header
    /// (e.g. `"https://app.example.com"`). When empty (the default), any
    /// request carrying an `Origin` header is rejected with 403 — server-to-
    /// server clients without `Origin` are unaffected. Non-`null` matches in
    /// the allowlist also enable a `CorsLayer` that emits the corresponding
    /// preflight responses; an empty list installs no CORS layer.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_allowed_origins(vec!["https://app.example.com".to_string()]);
    /// ```
    #[must_use]
    pub fn with_allowed_origins<I, S>(mut self, origins: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.allowed_origins = origins.into_iter().map(Into::into).collect();
        self
    }

    /// Configure a STDIO backend (subprocess)
    ///
    /// # Arguments
    ///
    /// * `command` - Command to execute (must be in allowlist)
    /// * `args` - Command arguments
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_stdio_backend("python", vec!["server.py".to_string()]);
    /// ```
    #[must_use]
    pub fn with_stdio_backend(mut self, command: impl Into<String>, args: Vec<String>) -> Self {
        self.backend_config = Some(BackendConfig::Stdio {
            command: command.into(),
            args,
            working_dir: None,
        });
        self
    }

    /// Configure a STDIO backend with working directory
    ///
    /// # Arguments
    ///
    /// * `command` - Command to execute (must be in allowlist)
    /// * `args` - Command arguments
    /// * `working_dir` - Working directory for the subprocess
    #[must_use]
    pub fn with_stdio_backend_and_dir(
        mut self,
        command: impl Into<String>,
        args: Vec<String>,
        working_dir: impl Into<String>,
    ) -> Self {
        self.backend_config = Some(BackendConfig::Stdio {
            command: command.into(),
            args,
            working_dir: Some(working_dir.into()),
        });
        self
    }

    /// Configure an HTTP backend
    ///
    /// # Arguments
    ///
    /// * `url` - Base URL of the HTTP server (HTTPS required for non-localhost)
    /// * `auth_token` - Optional authentication token
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_http_backend("https://api.example.com", Some("token123".to_string()));
    /// ```
    #[must_use]
    pub fn with_http_backend(mut self, url: impl Into<String>, auth_token: Option<String>) -> Self {
        self.backend_config = Some(BackendConfig::Http {
            url: url.into(),
            endpoint_path: None,
            auth_token,
        });
        self
    }

    /// Configure an HTTP backend with an explicit MCP endpoint path.
    ///
    /// Use this when the upstream MCP server mounts at a path other than
    /// the default `/mcp` (e.g. `"/api/mcp"`).
    #[must_use]
    pub fn with_http_backend_path(
        mut self,
        url: impl Into<String>,
        endpoint_path: impl Into<String>,
        auth_token: Option<String>,
    ) -> Self {
        self.backend_config = Some(BackendConfig::Http {
            url: url.into(),
            endpoint_path: Some(endpoint_path.into()),
            auth_token,
        });
        self
    }

    /// Configure a WebSocket backend
    ///
    /// # Arguments
    ///
    /// * `url` - WebSocket URL (e.g., "<ws://localhost:8080>" or "<wss://server.example.com>")
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_websocket_backend("wss://mcp.example.com");
    /// ```
    #[must_use]
    pub fn with_websocket_backend(mut self, url: impl Into<String>) -> Self {
        self.backend_config = Some(BackendConfig::WebSocket { url: url.into() });
        self
    }

    /// Configure a TCP backend
    ///
    /// # Arguments
    ///
    /// * `host` - Host or IP address to connect to
    /// * `port` - Port number
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_tcp_backend("localhost", 5000);
    /// ```
    #[must_use]
    pub fn with_tcp_backend(mut self, host: impl Into<String>, port: u16) -> Self {
        self.backend_config = Some(BackendConfig::Tcp {
            host: host.into(),
            port,
        });
        self
    }

    /// Configure a Unix domain socket backend
    ///
    /// # Arguments
    ///
    /// * `path` - Path to Unix socket file
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_unix_backend("/tmp/mcp.sock");
    /// ```
    #[cfg(unix)]
    #[must_use]
    pub fn with_unix_backend(mut self, path: impl Into<String>) -> Self {
        self.backend_config = Some(BackendConfig::Unix { path: path.into() });
        self
    }

    /// Configure an HTTP frontend
    ///
    /// # Arguments
    ///
    /// * `bind` - Address to bind to (e.g., "127.0.0.1:3000")
    ///
    /// # Security Note
    ///
    /// Default is localhost-only. Only bind to 0.0.0.0 if you have proper
    /// authentication and network security in place.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_http_frontend("127.0.0.1:8080");
    /// ```
    #[must_use]
    pub fn with_http_frontend(mut self, bind: impl Into<String>) -> Self {
        self.frontend_type = Some(FrontendType::Http);
        self.bind_address = Some(bind.into());
        self
    }

    /// Configure a STDIO frontend
    ///
    /// Reads JSON-RPC from stdin, writes to stdout. Ideal for CLI tools.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_stdio_frontend();
    /// ```
    #[must_use]
    pub fn with_stdio_frontend(mut self) -> Self {
        self.frontend_type = Some(FrontendType::Stdio);
        self
    }

    /// Configure a WebSocket frontend
    ///
    /// Bidirectional WebSocket server for real-time communication.
    /// Ideal for browser clients and bidirectional elicitation.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// let builder = RuntimeProxyBuilder::new()
    ///     .with_websocket_frontend("127.0.0.1:8080");
    /// ```
    #[must_use]
    pub fn with_websocket_frontend(mut self, bind: impl Into<String>) -> Self {
        self.frontend_type = Some(FrontendType::WebSocket);
        self.bind_address = Some(bind.into());
        self
    }

    /// Set maximum request size limit
    ///
    /// # Arguments
    ///
    /// * `limit` - Maximum size in bytes (default: 10 MB)
    ///
    /// # Security Note
    ///
    /// Prevents memory exhaustion from large requests.
    #[must_use]
    pub fn with_request_size_limit(mut self, limit: usize) -> Self {
        self.request_size_limit = limit;
        self
    }

    /// Set request timeout
    ///
    /// # Arguments
    ///
    /// * `timeout_ms` - Timeout in milliseconds (max: 5 minutes)
    ///
    /// # Errors
    ///
    /// Returns an error if timeout exceeds maximum.
    pub fn with_timeout(mut self, timeout_ms: u64) -> ProxyResult<Self> {
        if timeout_ms > MAX_TIMEOUT_MS {
            return Err(ProxyError::configuration_with_key(
                format!("Timeout {timeout_ms}ms exceeds maximum {MAX_TIMEOUT_MS}ms"),
                "timeout_ms",
            ));
        }
        self.timeout_ms = timeout_ms;
        Ok(self)
    }

    /// Enable or disable metrics collection
    ///
    /// # Arguments
    ///
    /// * `enable` - Whether to collect metrics (default: true)
    #[must_use]
    pub fn with_metrics(mut self, enable: bool) -> Self {
        self.enable_metrics = enable;
        self
    }

    /// Configure backend URL validation and SSRF protection
    ///
    /// # Arguments
    ///
    /// * `config` - Backend validation configuration
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// # use turbomcp_proxy::config::{BackendValidationConfig, SsrfProtection};
    /// # use ipnetwork::IpNetwork;
    /// # use std::str::FromStr;
    /// # async fn example() -> turbomcp_proxy::ProxyResult<()> {
    /// // Allow connections to specific private network
    /// let validation = BackendValidationConfig {
    ///     ssrf_protection: SsrfProtection::Balanced {
    ///         allowed_private_networks: vec![
    ///             IpNetwork::from_str("10.0.0.0/8").unwrap(),
    ///         ],
    ///     },
    ///     ..Default::default()
    /// };
    ///
    /// let proxy = RuntimeProxyBuilder::new()
    ///     .with_websocket_backend("ws://10.0.1.5:8080")
    ///     .with_http_frontend("127.0.0.1:3000")
    ///     .with_backend_validation(validation)
    ///     .build()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    #[must_use]
    pub fn with_backend_validation(mut self, config: BackendValidationConfig) -> Self {
        self.validation_config = config;
        self
    }

    /// Build and validate the runtime proxy
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Backend configuration is missing
    /// - Frontend type is missing
    /// - Security validation fails (command not in allowlist, invalid URL, etc.)
    /// - Backend connection fails
    ///
    /// # Panics
    ///
    /// Panics if `backend_config` is None after successful validation (should never happen as validation ensures it's Some).
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use turbomcp_proxy::runtime::RuntimeProxyBuilder;
    /// # async fn example() -> turbomcp_proxy::ProxyResult<()> {
    /// let proxy = RuntimeProxyBuilder::new()
    ///     .with_stdio_backend("python", vec!["server.py".to_string()])
    ///     .with_http_frontend("127.0.0.1:3000")
    ///     .build()
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn build(self) -> ProxyResult<RuntimeProxy> {
        // Ensure required fields are set
        let backend_config = self
            .backend_config
            .as_ref()
            .ok_or_else(|| ProxyError::configuration("Backend configuration is required"))?;

        let frontend_type = self
            .frontend_type
            .ok_or_else(|| ProxyError::configuration("Frontend type is required"))?;

        // Validate security constraints
        Self::validate_command(backend_config)?;
        Self::validate_url(backend_config, &self.validation_config).await?;
        Self::validate_working_dir(backend_config)?;

        // Take ownership after validation
        let backend_config = self.backend_config.unwrap();

        // Convert BackendConfig to BackendTransport for BackendConnector
        let transport = match &backend_config {
            BackendConfig::Stdio {
                command,
                args,
                working_dir,
            } => BackendTransport::Stdio {
                command: command.clone(),
                args: args.clone(),
                working_dir: working_dir.clone(),
            },
            BackendConfig::Http {
                url,
                endpoint_path,
                auth_token,
            } => BackendTransport::Http {
                url: url.clone(),
                endpoint_path: endpoint_path.clone(),
                // Wrap in SecretString as the value crosses into the internal
                // backend layer; from here on it stays redacted in Debug output.
                auth_token: auth_token.clone().map(secrecy::SecretString::from),
            },
            BackendConfig::Tcp { host, port } => BackendTransport::Tcp {
                host: host.clone(),
                port: *port,
            },
            #[cfg(unix)]
            BackendConfig::Unix { path } => BackendTransport::Unix { path: path.clone() },
            BackendConfig::WebSocket { url } => BackendTransport::WebSocket { url: url.clone() },
        };

        // Create BackendConnector configuration
        let connector_config = crate::proxy::backend::BackendConfig {
            transport,
            client_name: "turbomcp-proxy".to_string(),
            client_version: crate::VERSION.to_string(),
        };

        // Create backend connector
        let backend = BackendConnector::new(connector_config).await?;

        // Create metrics if enabled
        let metrics = if self.enable_metrics {
            Some(Arc::new(AtomicMetrics::new()))
        } else {
            None
        };

        Ok(RuntimeProxy {
            backend,
            frontend_type,
            bind_address: self.bind_address,
            request_size_limit: self.request_size_limit,
            timeout_ms: self.timeout_ms,
            metrics,
            origin_allowlist: OriginAllowlist::new(self.allowed_origins),
        })
    }

    /// Validate command is in allowlist (SECURITY CRITICAL)
    fn validate_command(config: &BackendConfig) -> ProxyResult<()> {
        if let BackendConfig::Stdio { command, .. } = config
            && !ALLOWED_COMMANDS.contains(&command.as_str())
        {
            return Err(ProxyError::configuration_with_key(
                format!("Command '{command}' not in allowlist. Allowed: {ALLOWED_COMMANDS:#?}"),
                "command",
            ));
        }
        Ok(())
    }

    /// Validate URL for SSRF protection (SECURITY CRITICAL)
    ///
    /// Validates both HTTP and WebSocket URLs against SSRF protection rules.
    /// Blocks private IPs, cloud metadata endpoints, and validates schemes.
    async fn validate_url(
        config: &BackendConfig,
        validation_config: &BackendValidationConfig,
    ) -> ProxyResult<()> {
        // Extract URL based on backend type
        let (BackendConfig::Http { url: url_str, .. } | BackendConfig::WebSocket { url: url_str }) =
            config
        else {
            return Ok(()); // Other backend types don't use URLs
        };

        let parsed = url::Url::parse(url_str)
            .map_err(|e| ProxyError::configuration_with_key(format!("Invalid URL: {e}"), "url"))?;

        // Validate scheme is allowed
        if !validation_config
            .allowed_schemes
            .contains(&parsed.scheme().to_string())
        {
            return Err(ProxyError::configuration_with_key(
                format!(
                    "Scheme '{}' not allowed. Allowed schemes: {}",
                    parsed.scheme(),
                    validation_config.allowed_schemes.join(", ")
                ),
                "url",
            ));
        }

        // Require secure protocols (HTTPS/WSS) except for localhost
        if parsed.scheme() == "http" || parsed.scheme() == "ws" {
            let host = parsed.host_str().unwrap_or("");
            if !is_localhost(host) {
                let secure_scheme = if parsed.scheme() == "http" {
                    "https"
                } else {
                    "wss"
                };
                return Err(ProxyError::configuration_with_key(
                    format!(
                        "Secure protocol required for non-localhost URLs. Use {} instead of {}",
                        secure_scheme,
                        parsed.scheme()
                    ),
                    "url",
                ));
            }
        }

        // Apply SSRF protection based on configuration
        if let Some(host) = parsed.host_str() {
            let port = parsed.port_or_known_default().ok_or_else(|| {
                ProxyError::configuration_with_key(
                    format!(
                        "URL is missing a usable port for scheme '{}'",
                        parsed.scheme()
                    ),
                    "url",
                )
            })?;
            Self::validate_host(host, port, validation_config).await?;
        }

        Ok(())
    }

    /// Validate host is not private/metadata based on SSRF protection level (SECURITY CRITICAL)
    async fn validate_host(
        host: &str,
        port: u16,
        validation_config: &BackendValidationConfig,
    ) -> ProxyResult<()> {
        // Check custom blocklist first
        if validation_config.blocked_hosts.contains(&host.to_string()) {
            return Err(ProxyError::configuration_with_key(
                format!("Host '{host}' is blocked by custom blocklist"),
                "url",
            ));
        }

        // Apply SSRF protection based on configuration
        match &validation_config.ssrf_protection {
            SsrfProtection::Disabled => {
                warn!("SSRF protection disabled for host: {}", host);
                Ok(())
            }
            SsrfProtection::Strict => Self::validate_host_strict(host, port).await,
            SsrfProtection::Balanced {
                allowed_private_networks,
            } => Self::validate_host_balanced(host, port, allowed_private_networks).await,
        }
    }

    /// Strict SSRF validation - blocks all private networks
    async fn validate_host_strict(host: &str, port: u16) -> ProxyResult<()> {
        Self::validate_host_addresses(host, port, |ip| match classify_ip(ip) {
            IpClass::Loopback | IpClass::Public => Ok(()),
            IpClass::Metadata => Err(metadata_blocked(ip)),
            IpClass::Private => Err(ProxyError::configuration_with_key(
                format!(
                    "Private {} address blocked: {ip}. \
                    For internal proxies, configure:\n  \
                    SsrfProtection::Balanced {{ \
                    allowed_private_networks: vec![IpNetwork::from_str(\"10.0.0.0/8\")?] }}",
                    match canonical_ip(ip) {
                        IpAddr::V4(_) => "IPv4",
                        IpAddr::V6(_) => "IPv6",
                    }
                ),
                "url",
            )),
        })
        .await
    }

    /// Balanced SSRF validation - allows specific private networks
    async fn validate_host_balanced(
        host: &str,
        port: u16,
        allowed_networks: &[IpNetwork],
    ) -> ProxyResult<()> {
        Self::validate_host_addresses(host, port, |ip| {
            Self::validate_ip_balanced(ip, allowed_networks)
        })
        .await
    }

    fn validate_ip_balanced(ip: IpAddr, allowed_networks: &[IpNetwork]) -> ProxyResult<()> {
        match classify_ip(ip) {
            IpClass::Loopback | IpClass::Public => Ok(()),
            // Always blocked, even in balanced mode: no allowlist entry
            // should ever be read as permission to reach instance metadata.
            IpClass::Metadata => Err(metadata_blocked(ip)),
            IpClass::Private => {
                let canonical = canonical_ip(ip);
                if allowed_networks
                    .iter()
                    .any(|net| net.contains(ip) || net.contains(canonical))
                {
                    debug!("Private IP {} allowed by configured network", ip);
                    Ok(())
                } else {
                    Err(ProxyError::configuration_with_key(
                        format!(
                            "Private IP {ip} not in allowed networks. Allowed networks: {allowed_networks:?}"
                        ),
                        "url",
                    ))
                }
            }
        }
    }

    /// Apply `validate_ip` to the host's address: the literal if it is one,
    /// otherwise every address it resolves to.
    ///
    /// Metadata hostnames are refused by name first, since they resolve only
    /// inside the cloud network they belong to.
    async fn validate_host_addresses<F>(
        host: &str,
        port: u16,
        mut validate_ip: F,
    ) -> ProxyResult<()>
    where
        F: FnMut(IpAddr) -> ProxyResult<()>,
    {
        if matches!(host, "metadata.google.internal" | "metadata") {
            return Err(ProxyError::configuration_with_key(
                format!("Cloud metadata endpoint blocked: {host}"),
                "url",
            ));
        }

        // `Url::host_str` keeps the brackets on an IPv6 literal.
        let literal = host.trim_start_matches('[').trim_end_matches(']');
        if let Ok(ip) = literal.parse::<IpAddr>() {
            return validate_ip(ip);
        }

        let resolved = tokio::net::lookup_host((host, port)).await.map_err(|e| {
            ProxyError::configuration_with_key(
                format!("Failed to resolve host '{host}': {e}"),
                "url",
            )
        })?;

        let mut saw_ip = false;
        for addr in resolved {
            saw_ip = true;
            validate_ip(addr.ip())?;
        }

        if !saw_ip {
            return Err(ProxyError::configuration_with_key(
                format!("Host '{host}' resolved to no addresses"),
                "url",
            ));
        }

        Ok(())
    }

    /// Validate working directory (path traversal protection)
    fn validate_working_dir(config: &BackendConfig) -> ProxyResult<()> {
        if let BackendConfig::Stdio {
            working_dir: Some(wd),
            ..
        } = config
        {
            let path = PathBuf::from(wd);

            // Ensure path exists
            if !path.exists() {
                return Err(ProxyError::configuration_with_key(
                    format!("Working directory does not exist: {wd}"),
                    "working_dir",
                ));
            }

            // Canonicalize to resolve symlinks and relative paths
            let canonical = path.canonicalize().map_err(|e| {
                ProxyError::configuration_with_key(
                    format!("Failed to canonicalize working directory: {e}"),
                    "working_dir",
                )
            })?;

            // Additional validation: ensure it's a directory
            if !canonical.is_dir() {
                return Err(ProxyError::configuration_with_key(
                    format!("Working directory is not a directory: {wd}"),
                    "working_dir",
                ));
            }
        }
        Ok(())
    }
}

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

/// Check if a host string refers to a loopback address.
///
/// Accepts both bracketed (`"[::1]"`) and unbracketed (`"::1"`) IPv6 forms;
/// callers obtained their host string from various sources (raw config,
/// `Url::host_str()` which strips brackets, etc.) and asking each to
/// normalize first was inviting drift. Internally we compare against the
/// bracket-stripped form so `"[::1]"` and `"::1"` route identically.
fn is_localhost(host: &str) -> bool {
    let normalized = host
        .strip_prefix('[')
        .and_then(|s| s.strip_suffix(']'))
        .unwrap_or(host);
    matches!(normalized, "localhost" | "127.0.0.1" | "::1")
}

/// What an address means for SSRF purposes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum IpClass {
    /// This host. Always allowed: the proxy fronting a local server is the
    /// primary use case.
    Loopback,
    /// A cloud instance-metadata endpoint. Never allowed.
    Metadata,
    /// Not publicly routable: private, link-local, shared (CGNAT),
    /// unspecified or "this network".
    Private,
    /// Anything else.
    Public,
}

/// The IPv4 address an IPv6 address stands for, if it stands for one.
///
/// IPv4-mapped (`::ffff:a.b.c.d`) addresses are how dual-stack sockets reach
/// IPv4 hosts, and NAT64 (`64:ff9b::/96`) addresses are translated to their
/// embedded IPv4 address by the network. Either one let a private or metadata
/// address through as an innocuous-looking IPv6 literal, because only the
/// IPv6 rules were ever applied to it.
fn canonical_ip(ip: IpAddr) -> IpAddr {
    match ip {
        IpAddr::V6(v6) => {
            if let Some(v4) = v6.to_ipv4_mapped() {
                return IpAddr::V4(v4);
            }
            let segments = v6.segments();
            if segments[..6] == [0x64, 0xff9b, 0, 0, 0, 0] {
                let [a, b] = segments[6].to_be_bytes();
                let [c, d] = segments[7].to_be_bytes();
                return IpAddr::V4(Ipv4Addr::new(a, b, c, d));
            }
            ip
        }
        IpAddr::V4(_) => ip,
    }
}

/// Classify an address, comparing addresses rather than strings so that
/// every spelling of one (`169.254.169.254`, `::ffff:a9fe:a9fe`,
/// `64:ff9b::a9fe:a9fe`) gets the same answer.
fn classify_ip(ip: IpAddr) -> IpClass {
    match canonical_ip(ip) {
        IpAddr::V4(v4) => {
            let octets = v4.octets();
            if v4.is_loopback() {
                IpClass::Loopback
            } else if v4 == Ipv4Addr::new(169, 254, 169, 254)
                || v4 == Ipv4Addr::new(168, 63, 129, 16)
            {
                // AWS/GCP/Azure IMDS and Azure's wireserver.
                IpClass::Metadata
            } else if v4.is_private()
                || v4.is_link_local()
                // 100.64.0.0/10, shared address space (RFC 6598): carrier-grade
                // NAT, and inside some clouds, internal service addresses.
                || (octets[0] == 100 && octets[1] & 0xc0 == 64)
                // 0.0.0.0/8, "this network": 0.0.0.0 reaches the local host.
                || octets[0] == 0
                || v4.is_broadcast()
            {
                IpClass::Private
            } else {
                IpClass::Public
            }
        }
        IpAddr::V6(v6) => {
            if v6.is_loopback() {
                IpClass::Loopback
            } else if v6 == Ipv6Addr::new(0xfd00, 0x0ec2, 0, 0, 0, 0, 0, 0x254) {
                // AWS IMDS over IPv6.
                IpClass::Metadata
            } else if v6.is_unspecified()
                // fc00::/7 unique local, fe80::/10 link-local
                || v6.segments()[0] & 0xfe00 == 0xfc00
                || v6.segments()[0] & 0xffc0 == 0xfe80
            {
                IpClass::Private
            } else {
                IpClass::Public
            }
        }
    }
}

fn metadata_blocked(ip: IpAddr) -> ProxyError {
    ProxyError::configuration_with_key(
        format!(
            "Cloud metadata endpoint blocked: {ip}. \
            For internal proxies, use SsrfProtection::Balanced with allowed networks."
        ),
        "url",
    )
}

/// Runtime proxy instance
///
/// Manages the proxy lifecycle, routing requests between frontend and backend.
#[derive(Debug)]
pub struct RuntimeProxy {
    /// Backend connector
    backend: BackendConnector,

    /// Frontend type
    frontend_type: FrontendType,

    /// Bind address (for HTTP frontend)
    bind_address: Option<String>,

    /// Request size limit
    request_size_limit: usize,

    /// Request timeout
    timeout_ms: u64,

    /// Metrics collector
    metrics: Option<Arc<AtomicMetrics>>,

    /// Browser origins permitted to reach the HTTP/WebSocket frontend.
    origin_allowlist: OriginAllowlist,
}

impl RuntimeProxy {
    /// Run the proxy
    ///
    /// Starts the appropriate frontend based on configuration and runs
    /// until stopped or an error occurs.
    ///
    /// # Errors
    ///
    /// Returns an error if the frontend fails to start or encounters
    /// a fatal error during operation.
    pub async fn run(&mut self) -> ProxyResult<()> {
        match self.frontend_type {
            FrontendType::Http => {
                let bind = self
                    .bind_address
                    .as_ref()
                    .ok_or_else(|| {
                        ProxyError::configuration("Bind address required for HTTP frontend")
                    })?
                    .clone();
                self.run_http(&bind).await
            }
            FrontendType::Stdio => self.run_stdio().await,
            FrontendType::WebSocket => {
                let bind = self
                    .bind_address
                    .as_ref()
                    .ok_or_else(|| {
                        ProxyError::configuration("Bind address required for WebSocket frontend")
                    })?
                    .clone();
                self.run_websocket(&bind).await
            }
        }
    }

    /// Get reference to backend connector
    #[must_use]
    pub fn backend(&self) -> &BackendConnector {
        &self.backend
    }

    /// Get metrics snapshot
    #[must_use]
    pub fn metrics(&self) -> Option<crate::proxy::metrics::ProxyMetrics> {
        self.metrics.as_ref().map(|m| m.snapshot())
    }

    /// Run HTTP frontend using Axum and `ProxyService`
    async fn run_http(&mut self, bind: &str) -> ProxyResult<()> {
        use axum::{http::StatusCode, middleware};
        use std::time::Duration;
        use tower_http::limit::RequestBodyLimitLayer;
        use tower_http::timeout::TimeoutLayer;
        use turbomcp_server::McpServerExt;

        debug!("Starting HTTP frontend on {}", bind);

        // 1. Introspect backend to get ServerSpec
        let spec = self.backend.introspect().await?;

        debug!(
            "Backend introspection complete: {} tools, {} resources, {} prompts",
            spec.tools.len(),
            spec.resources.len(),
            spec.prompts.len()
        );

        // 2. Create ProxyService (takes ownership, so clone backend)
        let service = ProxyService::new(self.backend.clone(), spec);

        // 3. Create Axum router with MCP routes and security layers
        // Note: Security layers applied in both STDIO and HTTP frontends:
        //   - origin_guard / cors: Reject browser-origin requests not on allowlist
        //   - request_size_limit: Prevents memory exhaustion DoS
        //   - timeout_ms: Prevents hanging requests (STDIO uses tokio::time::timeout, HTTP uses Tower layer)
        let allowlist = self.origin_allowlist.clone();
        let server_config = turbomcp_server::ServerConfig::builder()
            .max_message_size(self.request_size_limit)
            // The proxy owns its stricter browser-origin policy in origin_guard:
            // no Origin is allowed for server-to-server clients, any Origin must
            // match the explicit allowlist.
            .allow_any_origin(true)
            .build();
        let mut app = service
            .builder()
            .with_config(server_config)
            .into_axum_router()
            .layer(middleware::from_fn_with_state(
                allowlist.clone(),
                origin_guard,
            ))
            .layer(RequestBodyLimitLayer::new(self.request_size_limit))
            .layer(TimeoutLayer::with_status_code(
                StatusCode::REQUEST_TIMEOUT,
                Duration::from_millis(self.timeout_ms),
            ));
        if let Some(cors) = build_cors_layer(&allowlist) {
            app = app.layer(cors);
        }

        // 4. Parse bind address
        let listener = tokio::net::TcpListener::bind(bind).await.map_err(|e| {
            ProxyError::backend_connection(format!("Failed to bind to {bind}: {e}"))
        })?;

        debug!("HTTP frontend listening on {}", bind);

        // 5. Start Axum server
        axum::serve(
            listener,
            app.into_make_service_with_connect_info::<std::net::SocketAddr>(),
        )
        .await
        .map_err(|e| ProxyError::backend(format!("Axum serve error: {e}")))?;

        Ok(())
    }

    /// Run WebSocket frontend using Axum and `ProxyService`
    async fn run_websocket(&mut self, bind: &str) -> ProxyResult<()> {
        debug!("Starting WebSocket frontend on {}", bind);

        // 1. Introspect backend to get ServerSpec
        let spec = self.backend.introspect().await?;

        debug!(
            "Backend introspection complete: {} tools, {} resources, {} prompts",
            spec.tools.len(),
            spec.resources.len(),
            spec.prompts.len()
        );

        // 2. Create ProxyService (takes ownership, so clone backend)
        let service = ProxyService::new(self.backend.clone(), spec);

        // 3. Run the supported WebSocket transport. It validates browser
        // Origin headers during upgrade using the same explicit allowlist.
        let mut config_builder = turbomcp_server::ServerConfig::builder()
            .max_message_size(self.request_size_limit)
            .allow_localhost_origins(false);
        for origin in self
            .origin_allowlist
            .header_values()
            .filter_map(|origin| origin.to_str().ok())
        {
            config_builder = config_builder.allow_origin(origin.to_owned());
        }
        let server_config = config_builder.build();

        debug!("WebSocket frontend listening on {}", bind);

        turbomcp_server::transport::websocket::run_with_config(&service, bind, &server_config)
            .await
            .map_err(|e| ProxyError::backend(format!("WebSocket server error: {e}")))?;

        Ok(())
    }

    /// Run STDIO frontend using `ProxyService` and the server's stdio transport
    ///
    /// Same stack as the HTTP and WebSocket frontends: the handshake, `ping`,
    /// notifications, and framing are `turbomcp-server`'s. This used to be a
    /// hand-written read-dispatch-write loop with no `initialize` arm at all,
    /// so no MCP client could ever complete a handshake against it.
    async fn run_stdio(&mut self) -> ProxyResult<()> {
        debug!("Starting STDIO frontend");

        let spec = self.backend.introspect().await?;
        let service = ProxyService::new(self.backend.clone(), spec)
            .with_request_timeout(Duration::from_millis(self.timeout_ms));
        let server_config = turbomcp_server::ServerConfig::builder()
            .max_message_size(self.request_size_limit)
            .build();

        crate::proxy::StdioFrontend::new(service, server_config)
            .run()
            .await
    }
}

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

    #[test]
    fn test_builder_creation() {
        let builder = RuntimeProxyBuilder::new();
        assert_eq!(builder.request_size_limit, MAX_REQUEST_SIZE);
        assert_eq!(builder.timeout_ms, DEFAULT_TIMEOUT_MS);
        assert!(builder.enable_metrics);
    }

    #[test]
    fn test_builder_with_stdio_backend() {
        let builder =
            RuntimeProxyBuilder::new().with_stdio_backend("python", vec!["server.py".to_string()]);

        assert!(matches!(
            builder.backend_config,
            Some(BackendConfig::Stdio { .. })
        ));
    }

    #[test]
    fn test_builder_with_http_backend() {
        let builder = RuntimeProxyBuilder::new().with_http_backend("https://api.example.com", None);

        assert!(matches!(
            builder.backend_config,
            Some(BackendConfig::Http { .. })
        ));
    }

    #[test]
    fn test_builder_with_tcp_backend() {
        let builder = RuntimeProxyBuilder::new().with_tcp_backend("localhost", 5000);

        assert!(matches!(
            builder.backend_config,
            Some(BackendConfig::Tcp {
                host: _,
                port: 5000
            })
        ));
    }

    #[cfg(unix)]
    #[test]
    fn test_builder_with_unix_backend() {
        let builder = RuntimeProxyBuilder::new().with_unix_backend("/tmp/mcp.sock");

        assert!(matches!(
            builder.backend_config,
            Some(BackendConfig::Unix { path: _ })
        ));
    }

    #[test]
    fn test_builder_with_frontends() {
        let http_builder = RuntimeProxyBuilder::new().with_http_frontend("0.0.0.0:3000");
        assert_eq!(http_builder.frontend_type, Some(FrontendType::Http));

        let stdio_builder = RuntimeProxyBuilder::new().with_stdio_frontend();
        assert_eq!(stdio_builder.frontend_type, Some(FrontendType::Stdio));
    }

    #[test]
    fn test_builder_with_timeout() {
        let result = RuntimeProxyBuilder::new().with_timeout(60_000);
        assert!(result.is_ok());
        assert_eq!(result.unwrap().timeout_ms, 60_000);
    }

    #[test]
    fn test_builder_timeout_exceeds_max() {
        let result = RuntimeProxyBuilder::new().with_timeout(MAX_TIMEOUT_MS + 1);
        assert!(result.is_err());
        match result {
            Err(ProxyError::Configuration { key, .. }) => {
                assert_eq!(key, Some("timeout_ms".to_string()));
            }
            _ => panic!("Expected Configuration error"),
        }
    }

    #[test]
    fn test_validate_command_allowed() {
        let config = BackendConfig::Stdio {
            command: "python".to_string(),
            args: vec![],
            working_dir: None,
        };

        assert!(RuntimeProxyBuilder::validate_command(&config).is_ok());
    }

    #[test]
    fn test_validate_command_not_allowed() {
        let config = BackendConfig::Stdio {
            command: "malicious_command".to_string(),
            args: vec![],
            working_dir: None,
        };

        let result = RuntimeProxyBuilder::validate_command(&config);
        assert!(result.is_err());
        match result {
            Err(ProxyError::Configuration { message, key }) => {
                assert!(message.contains("not in allowlist"));
                assert_eq!(key, Some("command".to_string()));
            }
            _ => panic!("Expected Configuration error"),
        }
    }

    #[tokio::test]
    async fn test_validate_url_https_required() {
        let config = BackendConfig::Http {
            url: "http://api.example.com".to_string(),
            endpoint_path: None,
            auth_token: None,
        };
        let validation_config = BackendValidationConfig::default();

        let result = RuntimeProxyBuilder::validate_url(&config, &validation_config).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_validate_url_localhost_http_allowed() {
        let config = BackendConfig::Http {
            url: "http://localhost:3000".to_string(),
            endpoint_path: None,
            auth_token: None,
        };
        let validation_config = BackendValidationConfig::default();

        assert!(
            RuntimeProxyBuilder::validate_url(&config, &validation_config)
                .await
                .is_ok()
        );
    }

    #[tokio::test]
    async fn test_validate_url_https_allowed() {
        let config = BackendConfig::Http {
            url: "https://8.8.8.8".to_string(),
            endpoint_path: None,
            auth_token: None,
        };
        let validation_config = BackendValidationConfig::default();

        assert!(
            RuntimeProxyBuilder::validate_url(&config, &validation_config)
                .await
                .is_ok()
        );
    }

    #[tokio::test]
    async fn test_validate_host_blocks_metadata() {
        let validation_config = BackendValidationConfig::default();

        // AWS metadata endpoint
        assert!(
            RuntimeProxyBuilder::validate_host("169.254.169.254", 443, &validation_config)
                .await
                .is_err()
        );

        // GCP metadata endpoint
        assert!(
            RuntimeProxyBuilder::validate_host("metadata.google.internal", 443, &validation_config)
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_validate_host_blocks_private_ips() {
        let validation_config = BackendValidationConfig::default();

        // Private IP ranges
        assert!(
            RuntimeProxyBuilder::validate_host("192.168.1.1", 443, &validation_config)
                .await
                .is_err()
        );
        assert!(
            RuntimeProxyBuilder::validate_host("10.0.0.1", 443, &validation_config)
                .await
                .is_err()
        );
        assert!(
            RuntimeProxyBuilder::validate_host("172.16.0.1", 443, &validation_config)
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn test_validate_host_allows_loopback() {
        let validation_config = BackendValidationConfig::default();

        assert!(
            RuntimeProxyBuilder::validate_host("127.0.0.1", 443, &validation_config)
                .await
                .is_ok()
        );
    }

    /// Every spelling of a private or metadata address is refused, not just
    /// the dotted-quad one. IPv4-mapped and NAT64 IPv6 literals used to get
    /// only the IPv6 rules, CGNAT and `0.0.0.0/8` weren't private, and the
    /// metadata check compared strings.
    #[tokio::test]
    async fn strict_mode_refuses_every_spelling_of_an_internal_address() {
        let validation_config = BackendValidationConfig::default();
        for host in [
            "[::ffff:10.0.0.1]",
            "[::ffff:169.254.169.254]",
            "[64:ff9b::a9fe:a9fe]",
            "[64:ff9b::a00:1]",
            "100.64.0.1",
            "100.127.255.254",
            "0.0.0.0",
            "0.1.2.3",
            "[::]",
            "[fd00:ec2::254]",
        ] {
            assert!(
                RuntimeProxyBuilder::validate_host(host, 443, &validation_config)
                    .await
                    .is_err(),
                "{host} must be refused"
            );
        }

        for host in ["8.8.8.8", "100.128.0.1", "[2606:4700::1111]", "[::1]"] {
            assert!(
                RuntimeProxyBuilder::validate_host(host, 443, &validation_config)
                    .await
                    .is_ok(),
                "{host} must be allowed"
            );
        }
    }

    /// Balanced mode honours an allowlisted network whichever way the
    /// address is spelled, and never lets an allowlist reach metadata.
    #[tokio::test]
    async fn balanced_mode_compares_addresses_not_spellings() {
        let validation_config = BackendValidationConfig {
            ssrf_protection: SsrfProtection::Balanced {
                allowed_private_networks: vec![
                    "10.0.0.0/8".parse().unwrap(),
                    "169.254.0.0/16".parse().unwrap(),
                ],
            },
            ..Default::default()
        };

        assert!(
            RuntimeProxyBuilder::validate_host("[::ffff:10.1.2.3]", 443, &validation_config)
                .await
                .is_ok()
        );
        for metadata in [
            "169.254.169.254",
            "[::ffff:a9fe:a9fe]",
            "[64:ff9b::a9fe:a9fe]",
        ] {
            assert!(
                RuntimeProxyBuilder::validate_host(metadata, 443, &validation_config)
                    .await
                    .is_err(),
                "{metadata} is metadata and must stay blocked"
            );
        }
        assert!(
            RuntimeProxyBuilder::validate_host("100.64.0.1", 443, &validation_config)
                .await
                .is_err()
        );
    }

    #[test]
    fn test_is_localhost() {
        assert!(is_localhost("localhost"));
        assert!(is_localhost("127.0.0.1"));
        assert!(is_localhost("::1"));
        assert!(is_localhost("[::1]"));
        assert!(!is_localhost("example.com"));
        assert!(!is_localhost("192.168.1.1"));
    }

    #[tokio::test]
    async fn test_builder_requires_backend() {
        let result = RuntimeProxyBuilder::new()
            .with_http_frontend("127.0.0.1:3000")
            .build()
            .await;

        assert!(result.is_err());
        match result {
            Err(ProxyError::Configuration { message, .. }) => {
                assert!(message.contains("Backend configuration is required"));
            }
            _ => panic!("Expected Configuration error"),
        }
    }

    #[tokio::test]
    async fn test_builder_requires_frontend() {
        let result = RuntimeProxyBuilder::new()
            .with_stdio_backend("python", vec!["server.py".to_string()])
            .build()
            .await;

        assert!(result.is_err());
        match result {
            Err(ProxyError::Configuration { message, .. }) => {
                assert!(message.contains("Frontend type is required"));
            }
            _ => panic!("Expected Configuration error"),
        }
    }

    #[test]
    fn test_validate_working_dir_nonexistent() {
        let config = BackendConfig::Stdio {
            command: "python".to_string(),
            args: vec![],
            working_dir: Some("/nonexistent/path/that/does/not/exist".to_string()),
        };

        let result = RuntimeProxyBuilder::validate_working_dir(&config);
        assert!(result.is_err());
    }

    #[test]
    fn test_constants() {
        assert_eq!(MAX_REQUEST_SIZE, 10 * 1024 * 1024);
        assert_eq!(DEFAULT_TIMEOUT_MS, 30_000);
        assert_eq!(MAX_TIMEOUT_MS, 300_000);
        assert_eq!(DEFAULT_BIND_ADDRESS, "127.0.0.1:3000");
        assert!(ALLOWED_COMMANDS.contains(&"python"));
        assert!(ALLOWED_COMMANDS.contains(&"node"));
    }
}