turbomcp-auth 3.0.12

//! # SSRF (Server-Side Request Forgery) Protection
//!
//! This module provides comprehensive protection against SSRF attacks,
//! which is critical for implementing Client ID Metadata Documents (CIMD)
//! in the MCP 2025-11-25 specification.
//!
//! ## Security Requirements (from MCP spec)
//!
//! Authorization servers that fetch client metadata documents **MUST**:
//! - Validate URLs and resolved IP addresses before fetching
//! - Implement response size limits (recommended: 5 kilobytes)
//! - Implement request timeouts
//! - Use aggressive caching to minimize repeated fetches
//! - Never cache errors
//! - Implement rate limiting per-client
//! - Monitor for unusual metadata fetch patterns
//! - Only fetch metadata after user authentication
//!
//! ## Attack Vectors Prevented
//!
//! - **Private Network Access**: Blocks access to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16)
//! - **Localhost Access**: Blocks 127.0.0.0/8, ::1
//! - **Link-Local**: Blocks 169.254.0.0/16 (IPv4) and fe80::/10 (IPv6)
//! - **Cloud Metadata**: Blocks 169.254.169.254 (AWS/Azure/GCP metadata endpoints)
//! - **DNS Rebinding**: Validates IP address at connection time (not just URL validation)
//! - **HTTP Redirects**: Optionally restricts redirect following
//!
//! ## Usage
//!
//! ```rust,ignore
//! use turbomcp_auth::ssrf::{SsrfValidator, SsrfPolicy};
//!
//! // Create validator with default policy (blocks all private networks)
//! let validator = SsrfValidator::default();
//!
//! // Validate a URL before fetching
//! validator.validate_url("https://example.com/.well-known/oauth-client")?;
//!
//! // Create custom policy
//! let policy = SsrfPolicy::builder()
//!     .allow_private_networks(false)
//!     .allow_localhost(false)
//!     .allow_cloud_metadata(false)
//!     .max_response_size(5 * 1024) // 5 KB
//!     .request_timeout(std::time::Duration::from_secs(5))
//!     .build();
//! ```

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
use std::time::Duration;
use thiserror::Error;
use tracing::{debug, warn};
use url::Url;

/// SSRF protection errors
#[derive(Debug, Clone, Error)]
pub enum SsrfError {
    /// URL validation failed
    #[error("Invalid URL: {0}")]
    InvalidUrl(String),

    /// URL scheme not allowed
    #[error("URL scheme not allowed: {0} (only https is permitted)")]
    InvalidScheme(String),

    /// IP address is blocked by policy
    #[error("IP address blocked: {0} ({1})")]
    BlockedIpAddress(IpAddr, String),

    /// Hostname resolution failed
    #[error("Failed to resolve hostname: {0}")]
    ResolutionFailed(String),

    /// Multiple IP addresses resolved (potential DNS rebinding)
    #[error("Multiple IP addresses resolved for hostname (potential DNS rebinding): {0}")]
    MultipleIpAddresses(String),

    /// Response size limit exceeded
    #[error("Response size limit exceeded: {0} bytes (max: {1} bytes)")]
    ResponseSizeLimitExceeded(usize, usize),

    /// Request timeout
    #[error("Request timeout after {0:?}")]
    Timeout(Duration),

    /// Cloud metadata endpoint access attempt
    #[error("Access to cloud metadata endpoint blocked: {0}")]
    CloudMetadataBlocked(IpAddr),

    /// Rate limit exceeded
    #[error("Rate limit exceeded for URL: {0}")]
    RateLimitExceeded(String),
}

/// SSRF protection policy configuration
#[derive(Debug, Clone)]
pub struct SsrfPolicy {
    /// Allow access to private network ranges (RFC 1918)
    pub allow_private_networks: bool,

    /// Allow access to localhost (127.0.0.0/8, ::1)
    pub allow_localhost: bool,

    /// Allow access to link-local addresses (169.254.0.0/16, fe80::/10)
    pub allow_link_local: bool,

    /// Allow access to cloud metadata endpoints (169.254.169.254)
    pub allow_cloud_metadata: bool,

    /// Maximum response size in bytes
    pub max_response_size: usize,

    /// Request timeout duration
    pub request_timeout: Duration,

    /// Require HTTPS scheme
    pub require_https: bool,

    /// Allow HTTP redirects
    pub allow_redirects: bool,

    /// Maximum number of redirects to follow
    pub max_redirects: u32,

    /// Custom IP address allowlist (if Some, only these IPs are allowed)
    pub ip_allowlist: Option<Vec<IpAddr>>,

    /// Custom IP address denylist (these IPs are always blocked)
    pub ip_denylist: Vec<IpAddr>,

    /// Custom hostname allowlist (if Some, only these hostnames are allowed)
    pub hostname_allowlist: Option<Vec<String>>,
}

impl Default for SsrfPolicy {
    fn default() -> Self {
        Self {
            allow_private_networks: false,
            allow_localhost: false,
            allow_link_local: false,
            allow_cloud_metadata: false,
            max_response_size: 5 * 1024, // 5 KB (MCP spec recommendation)
            request_timeout: Duration::from_secs(5),
            require_https: true,
            allow_redirects: false, // Disabled by default for security
            max_redirects: 0,
            ip_allowlist: None,
            ip_denylist: vec![
                // AWS metadata endpoint
                IpAddr::V4(Ipv4Addr::new(169, 254, 169, 254)),
                // Localhost variations
                IpAddr::V4(Ipv4Addr::LOCALHOST),
                IpAddr::V6(Ipv6Addr::LOCALHOST),
            ],
            hostname_allowlist: None,
        }
    }
}

impl SsrfPolicy {
    /// Create a builder for constructing policies
    pub fn builder() -> SsrfPolicyBuilder {
        SsrfPolicyBuilder::default()
    }

    /// Create a permissive policy (for testing - NOT for production)
    #[cfg(test)]
    pub fn permissive() -> Self {
        Self {
            allow_private_networks: true,
            allow_localhost: true,
            allow_link_local: true,
            allow_cloud_metadata: false,    // Still block cloud metadata
            max_response_size: 1024 * 1024, // 1 MB
            request_timeout: Duration::from_secs(30),
            require_https: false,
            allow_redirects: true,
            max_redirects: 5,
            ip_allowlist: None,
            ip_denylist: vec![],
            hostname_allowlist: None,
        }
    }
}

/// Builder for SSRF policies
#[derive(Debug, Default)]
pub struct SsrfPolicyBuilder {
    allow_private_networks: Option<bool>,
    allow_localhost: Option<bool>,
    allow_link_local: Option<bool>,
    allow_cloud_metadata: Option<bool>,
    max_response_size: Option<usize>,
    request_timeout: Option<Duration>,
    require_https: Option<bool>,
    allow_redirects: Option<bool>,
    max_redirects: Option<u32>,
    ip_allowlist: Option<Option<Vec<IpAddr>>>,
    ip_denylist: Option<Vec<IpAddr>>,
    hostname_allowlist: Option<Option<Vec<String>>>,
}

impl SsrfPolicyBuilder {
    /// Allow or deny access to private networks
    pub fn allow_private_networks(mut self, allow: bool) -> Self {
        self.allow_private_networks = Some(allow);
        self
    }

    /// Allow or deny access to localhost
    pub fn allow_localhost(mut self, allow: bool) -> Self {
        self.allow_localhost = Some(allow);
        self
    }

    /// Allow or deny access to link-local addresses
    pub fn allow_link_local(mut self, allow: bool) -> Self {
        self.allow_link_local = Some(allow);
        self
    }

    /// Allow or deny access to cloud metadata endpoints
    pub fn allow_cloud_metadata(mut self, allow: bool) -> Self {
        self.allow_cloud_metadata = Some(allow);
        self
    }

    /// Set maximum response size in bytes
    pub fn max_response_size(mut self, size: usize) -> Self {
        self.max_response_size = Some(size);
        self
    }

    /// Set request timeout duration
    pub fn request_timeout(mut self, timeout: Duration) -> Self {
        self.request_timeout = Some(timeout);
        self
    }

    /// Require HTTPS scheme
    pub fn require_https(mut self, require: bool) -> Self {
        self.require_https = Some(require);
        self
    }

    /// Allow HTTP redirects
    pub fn allow_redirects(mut self, allow: bool) -> Self {
        self.allow_redirects = Some(allow);
        self
    }

    /// Set maximum number of redirects
    pub fn max_redirects(mut self, max: u32) -> Self {
        self.max_redirects = Some(max);
        self
    }

    /// Set custom IP allowlist
    pub fn ip_allowlist(mut self, ips: Vec<IpAddr>) -> Self {
        self.ip_allowlist = Some(Some(ips));
        self
    }

    /// Set custom IP denylist
    pub fn ip_denylist(mut self, ips: Vec<IpAddr>) -> Self {
        self.ip_denylist = Some(ips);
        self
    }

    /// Set custom hostname allowlist
    pub fn hostname_allowlist(mut self, hostnames: Vec<String>) -> Self {
        self.hostname_allowlist = Some(Some(hostnames));
        self
    }

    /// Build the policy
    pub fn build(self) -> SsrfPolicy {
        let default = SsrfPolicy::default();
        SsrfPolicy {
            allow_private_networks: self
                .allow_private_networks
                .unwrap_or(default.allow_private_networks),
            allow_localhost: self.allow_localhost.unwrap_or(default.allow_localhost),
            allow_link_local: self.allow_link_local.unwrap_or(default.allow_link_local),
            allow_cloud_metadata: self
                .allow_cloud_metadata
                .unwrap_or(default.allow_cloud_metadata),
            max_response_size: self.max_response_size.unwrap_or(default.max_response_size),
            request_timeout: self.request_timeout.unwrap_or(default.request_timeout),
            require_https: self.require_https.unwrap_or(default.require_https),
            allow_redirects: self.allow_redirects.unwrap_or(default.allow_redirects),
            max_redirects: self.max_redirects.unwrap_or(default.max_redirects),
            ip_allowlist: self.ip_allowlist.unwrap_or(default.ip_allowlist),
            ip_denylist: self.ip_denylist.unwrap_or(default.ip_denylist),
            hostname_allowlist: self
                .hostname_allowlist
                .unwrap_or(default.hostname_allowlist),
        }
    }
}

/// SSRF validator
#[derive(Debug, Clone)]
pub struct SsrfValidator {
    policy: SsrfPolicy,
}

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

impl SsrfValidator {
    /// Create a new SSRF validator with the given policy
    pub fn new(policy: SsrfPolicy) -> Self {
        Self { policy }
    }

    /// Validate a URL before fetching
    ///
    /// # Errors
    ///
    /// Returns [`SsrfError`] if the URL fails validation
    pub fn validate_url(&self, url_str: &str) -> Result<(), SsrfError> {
        // Parse URL
        let url = Url::parse(url_str)
            .map_err(|e| SsrfError::InvalidUrl(format!("Failed to parse URL: {}", e)))?;

        // Validate scheme
        if self.policy.require_https && url.scheme() != "https" {
            return Err(SsrfError::InvalidScheme(url.scheme().to_string()));
        }

        // Check hostname allowlist
        if let Some(ref allowlist) = self.policy.hostname_allowlist
            && let Some(host) = url.host_str()
            && !allowlist.iter().any(|allowed| host == allowed)
        {
            debug!("Hostname not in allowlist: {}", host);
            return Err(SsrfError::InvalidUrl(format!(
                "Hostname not in allowlist: {}",
                host
            )));
        }

        // Resolve hostname and validate IP address
        if let Some(host) = url.host_str() {
            self.validate_hostname(host)?;
        } else {
            return Err(SsrfError::InvalidUrl("URL has no host".to_string()));
        }

        Ok(())
    }

    /// Validate a hostname by resolving it and checking the IP address
    ///
    /// # Errors
    ///
    /// Returns [`SsrfError`] if resolution fails or IP is blocked
    fn validate_hostname(&self, hostname: &str) -> Result<(), SsrfError> {
        // Resolve hostname to IP address(es)
        let addr_str = format!("{}:443", hostname); // Use port 443 for resolution
        let addrs: Vec<_> = addr_str
            .to_socket_addrs()
            .map_err(|e| SsrfError::ResolutionFailed(format!("{}: {}", hostname, e)))?
            .collect();

        if addrs.is_empty() {
            return Err(SsrfError::ResolutionFailed(format!(
                "No IP addresses resolved for: {}",
                hostname
            )));
        }

        // Check for multiple IPs (potential DNS rebinding)
        if addrs.len() > 1 {
            warn!(
                "Multiple IP addresses resolved for hostname (potential DNS rebinding): {} -> {:?}",
                hostname, addrs
            );
            // Don't fail, but log warning - this is common for load-balanced services
        }

        // Validate each resolved IP
        for socket_addr in addrs {
            let ip = socket_addr.ip();
            self.validate_ip_address(&ip)?;
        }

        Ok(())
    }

    /// Validate an IP address against the policy
    ///
    /// # Errors
    ///
    /// Returns [`SsrfError`] if the IP is blocked by policy
    pub fn validate_ip_address(&self, ip: &IpAddr) -> Result<(), SsrfError> {
        // Check IP allowlist first (if configured)
        if let Some(ref allowlist) = self.policy.ip_allowlist {
            if !allowlist.contains(ip) {
                debug!("IP not in allowlist: {}", ip);
                return Err(SsrfError::BlockedIpAddress(
                    *ip,
                    "IP not in allowlist".to_string(),
                ));
            }
            // If in allowlist, skip other checks
            return Ok(());
        }

        // Check for cloud metadata endpoint BEFORE general denylist (more specific error)
        if !self.policy.allow_cloud_metadata
            && let IpAddr::V4(ipv4) = ip
            && *ipv4 == Ipv4Addr::new(169, 254, 169, 254)
        {
            warn!("Cloud metadata endpoint access attempt: {}", ip);
            return Err(SsrfError::CloudMetadataBlocked(*ip));
        }

        // Check IP denylist
        if self.policy.ip_denylist.contains(ip) {
            warn!("IP in denylist: {}", ip);
            return Err(SsrfError::BlockedIpAddress(
                *ip,
                "IP in denylist".to_string(),
            ));
        }

        match ip {
            IpAddr::V4(ipv4) => self.validate_ipv4(ipv4)?,
            IpAddr::V6(ipv6) => self.validate_ipv6(ipv6)?,
        }

        Ok(())
    }

    /// Validate an IPv4 address
    fn validate_ipv4(&self, ip: &Ipv4Addr) -> Result<(), SsrfError> {
        // Check for private networks (RFC 1918)
        if !self.policy.allow_private_networks && ip.is_private() {
            debug!("Private network access blocked: {}", ip);
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V4(*ip),
                "Private network (RFC 1918)".to_string(),
            ));
        }

        // Check for localhost
        if !self.policy.allow_localhost && ip.is_loopback() {
            debug!("Localhost access blocked: {}", ip);
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V4(*ip),
                "Localhost".to_string(),
            ));
        }

        // Check for link-local
        if !self.policy.allow_link_local && ip.is_link_local() {
            debug!("Link-local access blocked: {}", ip);
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V4(*ip),
                "Link-local".to_string(),
            ));
        }

        // Additional checks
        if ip.is_unspecified() {
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V4(*ip),
                "Unspecified address (0.0.0.0)".to_string(),
            ));
        }

        if ip.is_broadcast() {
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V4(*ip),
                "Broadcast address".to_string(),
            ));
        }

        if ip.is_documentation() {
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V4(*ip),
                "Documentation address range".to_string(),
            ));
        }

        Ok(())
    }

    /// Validate an IPv6 address
    fn validate_ipv6(&self, ip: &Ipv6Addr) -> Result<(), SsrfError> {
        // Check for localhost
        if !self.policy.allow_localhost && ip.is_loopback() {
            debug!("Localhost access blocked: {}", ip);
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V6(*ip),
                "Localhost (::1)".to_string(),
            ));
        }

        // Check for unspecified
        if ip.is_unspecified() {
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V6(*ip),
                "Unspecified address (::)".to_string(),
            ));
        }

        // Note: Rust std doesn't have is_private() for IPv6 yet
        // Check for common private ranges manually
        if !self.policy.allow_private_networks {
            // Unique local addresses (fc00::/7)
            if ip.segments()[0] & 0xfe00 == 0xfc00 {
                debug!("Private network access blocked: {}", ip);
                return Err(SsrfError::BlockedIpAddress(
                    IpAddr::V6(*ip),
                    "Unique local address (fc00::/7)".to_string(),
                ));
            }
        }

        // Check for link-local (fe80::/10)
        if !self.policy.allow_link_local && (ip.segments()[0] & 0xffc0 == 0xfe80) {
            debug!("Link-local access blocked: {}", ip);
            return Err(SsrfError::BlockedIpAddress(
                IpAddr::V6(*ip),
                "Link-local (fe80::/10)".to_string(),
            ));
        }

        Ok(())
    }

    /// Get the policy
    pub fn policy(&self) -> &SsrfPolicy {
        &self.policy
    }

    /// Create an HTTP client with pinned DNS resolution
    ///
    /// This prevents DNS rebinding attacks by:
    /// 1. Resolving hostname to IP addresses
    /// 2. Validating all resolved IPs against SSRF policy
    /// 3. Creating a reqwest client with DNS pinned to validated IPs
    /// 4. Setting Host header manually to preserve the original hostname
    ///
    /// # Arguments
    ///
    /// * `url` - The URL to fetch (must pass `validate_url()` first)
    ///
    /// # Returns
    ///
    /// A tuple of (reqwest::Client, final_url) where:
    /// - Client has DNS pinned to validated IP addresses
    /// - final_url is the URL to use with the client
    ///
    /// # Security
    ///
    /// The client will ONLY connect to the validated IP addresses, even if
    /// DNS returns different IPs during the actual fetch. This prevents
    /// time-of-check/time-of-use (TOCTOU) DNS rebinding attacks.
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let validator = SsrfValidator::default();
    /// validator.validate_url("https://example.com")?;
    ///
    /// let (client, url) = validator.create_pinned_client("https://example.com")?;
    /// let response = client.get(&url).send().await?;
    /// ```
    pub fn create_pinned_client(
        &self,
        url_str: &str,
    ) -> Result<(reqwest::Client, String), SsrfError> {
        // Parse URL
        let url = Url::parse(url_str)
            .map_err(|e| SsrfError::InvalidUrl(format!("Failed to parse URL: {}", e)))?;

        let hostname = url
            .host_str()
            .ok_or_else(|| SsrfError::InvalidUrl("URL has no host".to_string()))?;

        // Determine port
        let port = url
            .port()
            .unwrap_or_else(|| if url.scheme() == "https" { 443 } else { 80 });

        // Resolve hostname to IP addresses
        let addr_str = format!("{}:{}", hostname, port);
        let addrs: Vec<SocketAddr> = addr_str
            .to_socket_addrs()
            .map_err(|e| SsrfError::ResolutionFailed(format!("{}: {}", hostname, e)))?
            .collect();

        if addrs.is_empty() {
            return Err(SsrfError::ResolutionFailed(format!(
                "No IP addresses resolved for: {}",
                hostname
            )));
        }

        // Validate all resolved IPs
        for socket_addr in &addrs {
            let ip = socket_addr.ip();
            self.validate_ip_address(&ip)?;
        }

        // Create reqwest client with pinned DNS
        let mut client_builder = reqwest::Client::builder().timeout(self.policy.request_timeout);

        // Pin DNS resolution to the validated IPs
        // Note: reqwest's resolve() takes a hostname and a single SocketAddr
        // We'll use the first validated IP (they're all validated at this point)
        if let Some(first_addr) = addrs.first() {
            debug!(
                hostname = hostname,
                resolved_ip = %first_addr.ip(),
                "Pinning DNS resolution to validated IP"
            );
            client_builder = client_builder.resolve(hostname, *first_addr);
        }

        // Configure redirect policy
        if !self.policy.allow_redirects {
            client_builder = client_builder.redirect(reqwest::redirect::Policy::none());
        } else {
            client_builder = client_builder.redirect(reqwest::redirect::Policy::limited(
                self.policy.max_redirects as usize,
            ));
        }

        let client = client_builder
            .build()
            .map_err(|e| SsrfError::InvalidUrl(format!("Failed to create HTTP client: {}", e)))?;

        // Return the original URL - the client will use pinned DNS
        Ok((client, url_str.to_string()))
    }

    /// Fetch a URL with SSRF protection and DNS pinning
    ///
    /// This is a convenience method that combines validation, DNS pinning, and fetching.
    ///
    /// # Arguments
    ///
    /// * `url` - The URL to fetch
    ///
    /// # Returns
    ///
    /// The HTTP response body as bytes
    ///
    /// # Errors
    ///
    /// Returns error if:
    /// - URL validation fails
    /// - DNS resolution fails
    /// - Any resolved IP is blocked
    /// - HTTP request fails
    /// - Response size exceeds limit
    ///
    /// # Example
    ///
    /// ```rust,ignore
    /// let validator = SsrfValidator::default();
    /// let body = validator.fetch("https://example.com/.well-known/oauth-client").await?;
    /// ```
    pub async fn fetch(&self, url: &str) -> Result<Vec<u8>, SsrfError> {
        // Validate URL
        self.validate_url(url)?;

        // Create client with pinned DNS
        let (client, final_url) = self.create_pinned_client(url)?;

        // Make request
        let response = client.get(&final_url).send().await.map_err(|e| {
            if e.is_timeout() {
                SsrfError::Timeout(self.policy.request_timeout)
            } else {
                SsrfError::InvalidUrl(format!("HTTP request failed: {}", e))
            }
        })?;

        // Check response size
        let content_length = response.content_length().unwrap_or(0) as usize;
        if content_length > self.policy.max_response_size {
            return Err(SsrfError::ResponseSizeLimitExceeded(
                content_length,
                self.policy.max_response_size,
            ));
        }

        // Read response body with size limit
        let bytes = response
            .bytes()
            .await
            .map_err(|e| SsrfError::InvalidUrl(format!("Failed to read response: {}", e)))?;

        if bytes.len() > self.policy.max_response_size {
            return Err(SsrfError::ResponseSizeLimitExceeded(
                bytes.len(),
                self.policy.max_response_size,
            ));
        }

        Ok(bytes.to_vec())
    }
}

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

    #[test]
    fn test_default_policy_blocks_private_networks() {
        let validator = SsrfValidator::default();

        // RFC 1918 private ranges
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)))
                .is_err()
        );
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(172, 16, 0, 1)))
                .is_err()
        );
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)))
                .is_err()
        );
    }

    #[test]
    fn test_default_policy_blocks_localhost() {
        let validator = SsrfValidator::default();

        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
                .is_err()
        );
        assert!(
            validator
                .validate_ip_address(&IpAddr::V6(Ipv6Addr::LOCALHOST))
                .is_err()
        );
    }

    #[test]
    fn test_default_policy_blocks_cloud_metadata() {
        let validator = SsrfValidator::default();

        assert!(matches!(
            validator.validate_ip_address(&IpAddr::V4(Ipv4Addr::new(169, 254, 169, 254))),
            Err(SsrfError::CloudMetadataBlocked(_))
        ));
    }

    #[test]
    fn test_default_policy_allows_public_ip() {
        let validator = SsrfValidator::default();

        // Google DNS
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)))
                .is_ok()
        );

        // Cloudflare DNS
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)))
                .is_ok()
        );
    }

    #[test]
    fn test_url_validation_requires_https() {
        let validator = SsrfValidator::default();

        assert!(matches!(
            validator.validate_url("http://example.com"),
            Err(SsrfError::InvalidScheme(_))
        ));

        // Note: example.com resolves to public IP, so this should pass if DNS works
        // In tests, we might not have network access, so we'll skip actual resolution tests
    }

    #[test]
    fn test_custom_policy_builder() {
        let policy = SsrfPolicy::builder()
            .allow_private_networks(true)
            .allow_localhost(false)
            .max_response_size(10 * 1024)
            .build();

        let validator = SsrfValidator::new(policy);

        // Private network should now be allowed
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)))
                .is_ok()
        );

        // Localhost still blocked
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::LOCALHOST))
                .is_err()
        );
    }

    #[test]
    fn test_ip_allowlist() {
        let policy = SsrfPolicy::builder()
            .ip_allowlist(vec![IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100))])
            .build();

        let validator = SsrfValidator::new(policy);

        // Only allowlisted IP should pass
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100)))
                .is_ok()
        );

        // Other IPs should fail
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(192, 168, 1, 101)))
                .is_err()
        );
    }

    #[test]
    fn test_ipv6_unique_local_blocked() {
        let validator = SsrfValidator::default();

        // fd00::1 is unique local (private)
        let ipv6 = Ipv6Addr::new(0xfd00, 0, 0, 0, 0, 0, 0, 1);
        assert!(validator.validate_ip_address(&IpAddr::V6(ipv6)).is_err());
    }

    #[test]
    fn test_link_local_blocked() {
        let validator = SsrfValidator::default();

        // 169.254.1.1 is link-local
        assert!(
            validator
                .validate_ip_address(&IpAddr::V4(Ipv4Addr::new(169, 254, 1, 1)))
                .is_err()
        );

        // fe80::1 is link-local
        let ipv6 = Ipv6Addr::new(0xfe80, 0, 0, 0, 0, 0, 0, 1);
        assert!(validator.validate_ip_address(&IpAddr::V6(ipv6)).is_err());
    }
}