liter-llm 1.11.2

//! Outbound HTTP request policy guard.
//!
//! Allows callers (the proxy server in particular) to constrain which upstream
//! URLs can be reached.  The library default is [`OutboundPolicy::Off`] —
//! preserves backward-compatibility for FFI consumers and embedded applications
//! where the application owner registers their own providers and is trusted.
//! The proxy switches the policy to [`OutboundPolicy::DenyPrivate`] at startup
//! so multi-tenant deployments cannot SSRF into cloud metadata services or
//! private networks.

use std::net::IpAddr;
use std::sync::{OnceLock, RwLock};

use url::Url;

use crate::error::LiterLlmError;

/// Controls which upstream URLs the library is allowed to connect to.
///
/// Set once at application startup via [`set_outbound_policy`].  Checked at
/// provider registration time and, when a custom DNS resolver is wired into
/// `reqwest`, at every TCP connection attempt (defense in depth against DNS
/// rebinding).
#[derive(Debug, Clone, Default)]
#[cfg_attr(alef, alef(skip))]
pub enum OutboundPolicy {
    /// No restrictions — library default.  Use only when the application is
    /// the sole registrar of provider URLs and trusts itself.
    #[default]
    Off,

    /// Reject URLs whose host resolves to any private / loopback / link-local
    /// / multicast / CGNAT address.  Recommended for multi-tenant proxies.
    DenyPrivate,

    /// Only allow URLs whose origin (scheme + host + port) matches one of the
    /// provided entries.
    Allowlist(Vec<Url>),
}

static GLOBAL_POLICY: OnceLock<RwLock<OutboundPolicy>> = OnceLock::new();

fn policy_lock() -> &'static RwLock<OutboundPolicy> {
    GLOBAL_POLICY.get_or_init(|| RwLock::new(OutboundPolicy::default()))
}

/// Set the global outbound policy.
///
/// Subsequent calls to [`validate_outbound_url`] and the per-connection DNS
/// resolver use this policy.  Intended to be called once at application
/// startup before any provider is registered.
#[cfg_attr(alef, alef(skip))]
pub fn set_outbound_policy(policy: OutboundPolicy) {
    *policy_lock().write().expect("outbound policy lock poisoned") = policy;
}

/// Read a snapshot of the current outbound policy.
#[cfg_attr(alef, alef(skip))]
pub fn current_policy() -> OutboundPolicy {
    policy_lock().read().expect("outbound policy lock poisoned").clone()
}

/// Validate `raw_url` against the current outbound policy.
///
/// Under [`OutboundPolicy::DenyPrivate`] the host is resolved via DNS; if
/// *any* returned address is forbidden the call returns
/// [`LiterLlmError::OutboundForbidden`].  This defeats DNS rebinding at
/// registration time.
///
/// Under [`OutboundPolicy::Off`] the function is a no-op and always returns
/// `Ok(())`.
#[cfg_attr(alef, alef(skip))]
pub async fn validate_outbound_url(raw_url: &str) -> Result<(), LiterLlmError> {
    let policy = current_policy();
    if matches!(policy, OutboundPolicy::Off) {
        return Ok(());
    }

    let url = Url::parse(raw_url).map_err(|e| LiterLlmError::OutboundForbidden {
        url: raw_url.to_string(),
        reason: format!("invalid URL: {e}"),
    })?;

    match url.scheme() {
        "http" | "https" => {}
        other => {
            return Err(LiterLlmError::OutboundForbidden {
                url: raw_url.to_string(),
                reason: format!("scheme '{other}' is not allowed; only http/https"),
            });
        }
    }

    match policy {
        OutboundPolicy::Off => Ok(()),
        OutboundPolicy::DenyPrivate => check_deny_private(&url, raw_url).await,
        OutboundPolicy::Allowlist(allowed) => check_allowlist(&url, raw_url, &allowed),
    }
}

#[cfg(target_arch = "wasm32")]
async fn check_deny_private(_url: &Url, _raw: &str) -> Result<(), LiterLlmError> {
    // ~keep WASM has no Rust DNS hook; only native GuardedResolver can block hostname SSRF.
    Ok(())
}

/// Synchronous URL validation — parse + scheme check + literal-IP private range
/// check only.  Does not perform DNS resolution.
///
/// Used from synchronous registration paths.  Catches the obvious
/// `http://169.254.169.254/` literal-IP case without requiring an async
/// context.  DNS-based checks still happen at connect time via
/// [`GuardedResolver`].
#[cfg_attr(alef, alef(skip))]
pub fn validate_outbound_url_sync(raw_url: &str) -> Result<(), LiterLlmError> {
    let policy = current_policy();
    if matches!(policy, OutboundPolicy::Off) {
        return Ok(());
    }

    let url = Url::parse(raw_url).map_err(|e| LiterLlmError::OutboundForbidden {
        url: raw_url.to_string(),
        reason: format!("invalid URL: {e}"),
    })?;

    match url.scheme() {
        "http" | "https" => {}
        other => {
            return Err(LiterLlmError::OutboundForbidden {
                url: raw_url.to_string(),
                reason: format!("scheme '{other}' is not allowed; only http/https"),
            });
        }
    }

    match url.host() {
        Some(url::Host::Ipv4(v4)) if is_forbidden(IpAddr::V4(v4)) => {
            return Err(LiterLlmError::OutboundForbidden {
                url: raw_url.to_string(),
                reason: format!("host is a forbidden address {v4}"),
            });
        }
        Some(url::Host::Ipv6(v6)) if is_forbidden(IpAddr::V6(v6)) => {
            return Err(LiterLlmError::OutboundForbidden {
                url: raw_url.to_string(),
                reason: format!("host is a forbidden address {v6}"),
            });
        }
        _ => {
            // ~keep No sync DNS here; GuardedResolver handles hostname SSRF at connect time.
        }
    }

    if let OutboundPolicy::Allowlist(allowed) = policy {
        return check_allowlist(&url, raw_url, &allowed);
    }

    Ok(())
}

#[cfg(not(target_arch = "wasm32"))]
async fn check_deny_private(url: &Url, raw: &str) -> Result<(), LiterLlmError> {
    let host = url.host_str().ok_or_else(|| LiterLlmError::OutboundForbidden {
        url: raw.to_string(),
        reason: "URL has no host".into(),
    })?;

    let port = url.port_or_known_default().unwrap_or(0);

    let addrs =
        tokio::net::lookup_host(format!("{host}:{port}"))
            .await
            .map_err(|e| LiterLlmError::OutboundForbidden {
                url: raw.to_string(),
                reason: format!("DNS resolution failed: {e}"),
            })?;

    for sa in addrs {
        if is_forbidden(sa.ip()) {
            return Err(LiterLlmError::OutboundForbidden {
                url: raw.to_string(),
                reason: format!("host resolves to forbidden address {}", sa.ip()),
            });
        }
    }
    Ok(())
}

fn check_allowlist(url: &Url, raw: &str, allowed: &[Url]) -> Result<(), LiterLlmError> {
    let origin_match = allowed.iter().any(|a| {
        a.scheme() == url.scheme()
            && a.host_str() == url.host_str()
            && a.port_or_known_default() == url.port_or_known_default()
    });
    if origin_match {
        Ok(())
    } else {
        Err(LiterLlmError::OutboundForbidden {
            url: raw.to_string(),
            reason: "URL not in outbound allowlist".into(),
        })
    }
}

/// Returns `true` if `ip` is in a range that must not be reached from a
/// multi-tenant proxy.
///
/// Covers loopback, unspecified, private (RFC 1918), link-local, multicast,
/// broadcast, CGNAT (100.64/10), IPv4-mapped IPv6, ULA (fc00::/7), and IPv6
/// link-local (fe80::/10).
#[cfg_attr(alef, alef(skip))]
pub fn is_forbidden(ip: IpAddr) -> bool {
    match ip {
        IpAddr::V4(v4) => {
            v4.is_loopback()
                || v4.is_unspecified()
                || v4.is_private()
                || v4.is_link_local()
                || v4.is_multicast()
                || v4.is_broadcast()
                || is_cgnat(v4)
        }
        IpAddr::V6(v6) => {
            v6.is_loopback()
                || v6.is_unspecified()
                || v6.is_multicast()
                || is_unique_local_v6(v6)
                || is_link_local_v6(v6)
                || v6
                    .to_ipv4_mapped()
                    .map(|m| is_forbidden(IpAddr::V4(m)))
                    .unwrap_or(false)
        }
    }
}

fn is_cgnat(ip: std::net::Ipv4Addr) -> bool {
    let [a, b, _, _] = ip.octets();
    a == 100 && (64..=127).contains(&b)
}

fn is_unique_local_v6(ip: std::net::Ipv6Addr) -> bool {
    (ip.segments()[0] & 0xfe00) == 0xfc00
}

fn is_link_local_v6(ip: std::net::Ipv6Addr) -> bool {
    (ip.segments()[0] & 0xffc0) == 0xfe80
}

/// A `reqwest` DNS resolver that filters resolved addresses through the
/// current [`OutboundPolicy`].
///
/// Install via `reqwest::Client::builder().dns_resolver(Arc::new(GuardedResolver))`.
/// Only active when the policy is not [`OutboundPolicy::Off`].  When the
/// policy is `Off` the resolver skips filtering entirely, falling back to
/// standard system behaviour.
#[cfg_attr(alef, alef(skip))]
pub struct GuardedResolver;

#[cfg(all(feature = "native-http", not(target_arch = "wasm32")))]
mod resolver_impl {
    use std::collections::HashMap;
    use std::net::SocketAddr;
    use std::sync::Arc;
    use std::sync::Mutex;
    use std::time::{Duration, Instant};

    use reqwest::dns::{Addrs, Name, Resolve, Resolving};

    use super::{GuardedResolver, OutboundPolicy, current_policy, is_forbidden};

    #[derive(Clone)]
    struct DnsCacheEntry {
        expires_at: Instant,
        addrs: Vec<SocketAddr>,
    }

    struct CachedGuardedResolver {
        cache_ttl: Option<Duration>,
        cache: Arc<Mutex<HashMap<String, DnsCacheEntry>>>,
    }

    impl CachedGuardedResolver {
        fn new(cache_ttl: Option<Duration>) -> Self {
            Self {
                cache_ttl,
                cache: Arc::new(Mutex::new(HashMap::new())),
            }
        }
    }

    fn cached_addrs(
        cache: &Mutex<HashMap<String, DnsCacheEntry>>,
        cache_ttl: Option<Duration>,
        host: &str,
    ) -> Option<Vec<SocketAddr>> {
        cache_ttl?;

        let now = Instant::now();
        let mut guard = cache.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
        let entry = guard.get(host)?;
        if entry.expires_at > now {
            return Some(entry.addrs.clone());
        }
        guard.remove(host);
        None
    }

    fn store_addrs(
        cache: &Mutex<HashMap<String, DnsCacheEntry>>,
        cache_ttl: Option<Duration>,
        host: String,
        addrs: Vec<SocketAddr>,
    ) {
        let Some(ttl) = cache_ttl else {
            return;
        };
        if ttl.is_zero() {
            return;
        }

        let expires_at = Instant::now().checked_add(ttl).unwrap_or_else(Instant::now);
        let mut guard = cache.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
        guard.insert(host, DnsCacheEntry { expires_at, addrs });
    }

    async fn resolve_system(host: &str) -> Result<Vec<SocketAddr>, Box<dyn std::error::Error + Send + Sync>> {
        let addrs = tokio::net::lookup_host((host, 0))
            .await
            .map_err(|e| {
                let err: Box<dyn std::error::Error + Send + Sync> = Box::new(e);
                err
            })?
            .collect();
        Ok(addrs)
    }

    fn validate_addrs(
        policy: OutboundPolicy,
        host: &str,
        addrs: &[SocketAddr],
    ) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
        if matches!(policy, OutboundPolicy::Off) {
            return Ok(());
        }

        for sa in addrs {
            if is_forbidden(sa.ip()) {
                let err: Box<dyn std::error::Error + Send + Sync> = format!(
                    "outbound DNS resolution for '{host}' produced forbidden address {}",
                    sa.ip()
                )
                .into();
                return Err(err);
            }
        }

        Ok(())
    }

    impl Resolve for GuardedResolver {
        fn resolve(&self, name: Name) -> Resolving {
            Box::pin(async move {
                let policy = current_policy();
                let host = name.as_str().to_string();

                let addrs = resolve_system(&host).await?;
                validate_addrs(policy, &host, &addrs)?;

                let iter: Addrs = Box::new(addrs.into_iter());
                Ok(iter)
            })
        }
    }

    impl Resolve for CachedGuardedResolver {
        fn resolve(&self, name: Name) -> Resolving {
            let policy = current_policy();
            let host = name.as_str().to_string();
            let cache_ttl = self.cache_ttl;
            let cache = Arc::clone(&self.cache);

            Box::pin(async move {
                if let Some(addrs) = cached_addrs(&cache, cache_ttl, &host) {
                    validate_addrs(policy, &host, &addrs)?;
                    let iter: Addrs = Box::new(addrs.into_iter());
                    return Ok(iter);
                }

                let addrs = resolve_system(&host).await?;
                validate_addrs(policy, &host, &addrs)?;
                store_addrs(&cache, cache_ttl, host, addrs.clone());

                let iter: Addrs = Box::new(addrs.into_iter());
                Ok(iter)
            })
        }
    }

    /// Build an [`Arc`]-wrapped [`GuardedResolver`] ready for use with
    /// `reqwest::Client::builder().dns_resolver(...)`.
    #[cfg_attr(alef, alef(skip))]
    pub fn guarded_resolver() -> Arc<GuardedResolver> {
        Arc::new(GuardedResolver)
    }

    /// Build a resolver that applies the active outbound policy and optionally
    /// caches successful DNS lookups for the configured TTL.
    #[cfg_attr(alef, alef(skip))]
    pub fn cached_guarded_resolver(cache_ttl: Option<Duration>) -> Arc<dyn Resolve> {
        Arc::new(CachedGuardedResolver::new(cache_ttl))
    }
}

#[cfg(all(feature = "native-http", not(target_arch = "wasm32")))]
pub use resolver_impl::{cached_guarded_resolver, guarded_resolver};

#[cfg(test)]
mod tests {
    use serial_test::serial;

    use super::*;

    /// Helper for sync tests — `#[serial(outbound_policy)]` on the test fn
    /// guarantees no other policy-mutating test is running concurrently.
    fn with_policy<F: FnOnce()>(policy: OutboundPolicy, f: F) {
        set_outbound_policy(policy);
        f();
        set_outbound_policy(OutboundPolicy::Off);
    }

    #[test]
    fn is_forbidden_recognizes_private_ranges() {
        let cases: &[(&str, bool)] = &[
            ("10.0.0.1", true),
            ("172.16.0.1", true),
            ("192.168.1.1", true),
            ("127.0.0.1", true),
            ("169.254.0.1", true),
            ("100.100.0.1", true),
            ("0.0.0.0", true),
            ("255.255.255.255", true),
            ("224.0.0.1", true),
            ("8.8.8.8", false),
            ("1.1.1.1", false),
        ];
        for (addr, expected) in cases {
            let ip: IpAddr = addr.parse().expect("valid IP");
            assert_eq!(is_forbidden(ip), *expected, "is_forbidden({addr}) should be {expected}");
        }
    }

    #[test]
    fn is_forbidden_ipv6_loopback() {
        let ip: IpAddr = "::1".parse().expect("::1 is a valid IPv6 address");
        assert!(is_forbidden(ip));
    }

    #[test]
    fn is_forbidden_ipv6_ula() {
        let ip: IpAddr = "fc00::1".parse().expect("fc00::1 is a valid IPv6 address");
        assert!(is_forbidden(ip));
    }

    #[test]
    fn is_forbidden_ipv6_link_local() {
        let ip: IpAddr = "fe80::1".parse().expect("fe80::1 is a valid IPv6 address");
        assert!(is_forbidden(ip));
    }

    #[test]
    fn is_forbidden_ipv6_public() {
        let ip: IpAddr = "2001:4860:4860::8888"
            .parse()
            .expect("Google DNS is a valid IPv6 address");
        assert!(!is_forbidden(ip));
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_off_passes_everything() {
        with_policy(OutboundPolicy::Off, || {
            assert!(validate_outbound_url_sync("http://127.0.0.1/").is_ok());
            assert!(validate_outbound_url_sync("http://169.254.169.254/").is_ok());
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_deny_private_rejects_loopback() {
        with_policy(OutboundPolicy::DenyPrivate, || {
            let result = validate_outbound_url_sync("http://127.0.0.1/");
            assert!(result.is_err(), "loopback should be rejected");
            let err = result.unwrap_err().to_string();
            assert!(
                err.contains("forbidden"),
                "error message should mention 'forbidden': {err}"
            );
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_deny_private_rejects_metadata_ip() {
        with_policy(OutboundPolicy::DenyPrivate, || {
            let result = validate_outbound_url_sync("http://169.254.169.254/");
            assert!(result.is_err(), "metadata IP should be rejected");
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_deny_private_rejects_ula() {
        with_policy(OutboundPolicy::DenyPrivate, || {
            let result = validate_outbound_url_sync("http://[fc00::1]/");
            assert!(result.is_err(), "ULA address should be rejected");
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_deny_private_rejects_link_local_v6() {
        with_policy(OutboundPolicy::DenyPrivate, || {
            let result = validate_outbound_url_sync("http://[fe80::1]/");
            assert!(result.is_err(), "IPv6 link-local should be rejected");
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_deny_private_rejects_unknown_scheme() {
        with_policy(OutboundPolicy::DenyPrivate, || {
            let result = validate_outbound_url_sync("ftp://example.com/");
            assert!(result.is_err(), "ftp:// scheme should be rejected");
            let err = result.unwrap_err().to_string();
            assert!(err.contains("scheme"), "error should mention 'scheme': {err}");
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_allowlist_accepts_exact_origin() {
        let allowed = vec![Url::parse("https://api.openai.com").expect("openai URL should be valid")];
        with_policy(OutboundPolicy::Allowlist(allowed), || {
            let result = validate_outbound_url_sync("https://api.openai.com/v1/chat/completions");
            assert!(result.is_ok(), "same-origin with different path should pass");
        });
    }

    #[test]
    #[serial(outbound_policy)]
    fn validate_sync_allowlist_rejects_other_host() {
        let allowed = vec![Url::parse("https://api.openai.com").expect("openai URL should be valid")];
        with_policy(OutboundPolicy::Allowlist(allowed), || {
            let result = validate_outbound_url_sync("https://api.anthropic.com/");
            assert!(result.is_err(), "different host should be rejected");
        });
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_off_passes_everything() {
        set_outbound_policy(OutboundPolicy::Off);
        assert!(validate_outbound_url("http://127.0.0.1/").await.is_ok());
        assert!(validate_outbound_url("http://169.254.169.254/").await.is_ok());
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_deny_private_rejects_loopback() {
        set_outbound_policy(OutboundPolicy::DenyPrivate);
        let result = validate_outbound_url("http://127.0.0.1/").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_err(), "loopback should be rejected by DenyPrivate");
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_deny_private_rejects_metadata_ip() {
        set_outbound_policy(OutboundPolicy::DenyPrivate);
        let result = validate_outbound_url("http://169.254.169.254/").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_err(), "AWS metadata IP should be rejected");
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_deny_private_rejects_ula() {
        set_outbound_policy(OutboundPolicy::DenyPrivate);
        let result = validate_outbound_url("http://[fc00::1]/").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_err(), "ULA address should be rejected");
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_deny_private_rejects_link_local_v6() {
        set_outbound_policy(OutboundPolicy::DenyPrivate);
        let result = validate_outbound_url("http://[fe80::1]/").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_err(), "IPv6 link-local should be rejected");
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_deny_private_rejects_unknown_scheme() {
        set_outbound_policy(OutboundPolicy::DenyPrivate);
        let result = validate_outbound_url("ftp://example.com/").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_err(), "ftp:// scheme should be rejected");
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_allowlist_accepts_exact_origin() {
        let allowed = vec![Url::parse("https://api.openai.com").expect("openai URL should be valid")];
        set_outbound_policy(OutboundPolicy::Allowlist(allowed));
        let result = validate_outbound_url("https://api.openai.com/v1/chat/completions").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_ok(), "same-origin with different path should pass");
    }

    #[tokio::test]
    #[serial(outbound_policy)]
    async fn validate_async_allowlist_rejects_other_host() {
        let allowed = vec![Url::parse("https://api.openai.com").expect("openai URL should be valid")];
        set_outbound_policy(OutboundPolicy::Allowlist(allowed));
        let result = validate_outbound_url("https://api.anthropic.com/").await;
        set_outbound_policy(OutboundPolicy::Off);
        assert!(result.is_err(), "different host should be rejected");
    }
}