harn-vm 0.10.23

use std::sync::Arc;

use url::Url;

use super::{current_resolved_ip_rules, current_ssrf_client_settings, GuardedResolver};

/// Install the connect-time SSRF guard while allowing exact configured
/// provider hosts to resolve to loopback/private-LAN addresses.
///
/// This is intended for operator-configured LLM endpoints only. Generic
/// tool/web/MCP traffic must keep using `install_ssrf_guard` so attacker-
/// controlled URLs cannot opt into private network reachability.
pub fn install_ssrf_guard_with_private_host_allowlist(
    builder: reqwest::ClientBuilder,
    allow_private_for_hosts: &[String],
) -> reqwest::ClientBuilder {
    let (block_private, allow_loopback) = current_ssrf_client_settings();
    let rules = current_resolved_ip_rules();
    if block_private || !rules.deny.is_empty() {
        builder.dns_resolver(Arc::new(
            GuardedResolver::with_policy_and_private_host_allowlist(
                block_private,
                allow_loopback,
                &rules,
                allow_private_for_hosts,
            ),
        ))
    } else {
        builder
    }
}

pub fn configured_provider_private_allow_host(base_url: &str) -> Option<String> {
    Url::parse(base_url)
        .ok()
        .and_then(|url| url.host_str().map(normalize_private_allow_host))
        .filter(|host| !host.is_empty())
}

fn normalize_private_allow_host(host: &str) -> String {
    host.trim().trim_end_matches('.').to_ascii_lowercase()
}

pub fn ssrf_client_cache_key(allow_private_for_hosts: &[String]) -> String {
    let (block_private, allow_loopback) = current_ssrf_client_settings();
    let rules = current_resolved_ip_rules();
    let deny_nets = rules
        .deny
        .iter()
        .map(|net| net.to_string())
        .collect::<Vec<_>>()
        .join(",");
    let hosts = allow_private_for_hosts
        .iter()
        .map(|host| normalize_private_allow_host(host))
        .collect::<Vec<_>>()
        .join(",");
    format!(
        "ssrf={block_private};loopback={allow_loopback};deny_nets={deny_nets};allow_private_hosts={hosts}"
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::egress::{
        require_ssrf_guard_for_host, test_support::OneShotHttpServer, EgressTestConfigGuard,
    };

    #[tokio::test]
    async fn install_ssrf_guard_private_host_allowlist_permits_configured_provider_host_only() {
        let server = OneShotHttpServer::start();
        let client = {
            let _config = EgressTestConfigGuard::new();
            let _scope = require_ssrf_guard_for_host();
            install_ssrf_guard_with_private_host_allowlist(
                reqwest::Client::builder(),
                &[String::from("localhost")],
            )
            .build()
            .expect("guarded client builds")
        };
        let response = client
            .get(server.url())
            .send()
            .await
            .expect("configured provider host may reach local inference endpoint");
        assert_eq!(response.status().as_u16(), 200);
        server.join();

        let server = OneShotHttpServer::start();
        let client = {
            let _config = EgressTestConfigGuard::new();
            let _scope = require_ssrf_guard_for_host();
            install_ssrf_guard_with_private_host_allowlist(
                reqwest::Client::builder(),
                &[String::from("other.local")],
            )
            .build()
            .expect("guarded client builds")
        };
        let result = client.get(server.url()).send().await;
        assert!(
            result.is_err(),
            "unlisted loopback hostname must remain guarded, got {result:?}"
        );
        server.unblock_and_join();
    }
}