//! Host-side network policy. `<dir>/host/network.toml` accepts `deny = ["address:port"]`.
//! Host listeners register their ports with `deny_host_port`. Guests cannot bind or connect there.
use std::{
    collections::HashSet,
    net::{IpAddr, SocketAddr},
    path::Path,
    sync::RwLock,
};

fn normalize(address: SocketAddr) -> SocketAddr {
    match address {
        SocketAddr::V6(v6) => v6
            .ip()
            .to_ipv4_mapped()
            .map(|ip| SocketAddr::new(ip.into(), v6.port()))
            .unwrap_or(address),
        _ => address,
    }
}
pub struct NetPolicy {
    denied: HashSet<SocketAddr>,
    local: HashSet<IpAddr>,
    ports: RwLock<HashSet<u16>>,
}
impl NetPolicy {
    pub fn open(host_dir: &Path) -> anyhow::Result<Self> {
        #[derive(serde::Deserialize, Default)]
        #[serde(deny_unknown_fields)]
        struct Config {
            #[serde(default)]
            deny: Vec<SocketAddr>,
        }
        let config = match std::fs::read_to_string(host_dir.join("network.toml")) {
            Ok(text) => toml::from_str::<Config>(&text)?,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Config::default(),
            Err(e) => return Err(e.into()),
        };
        Ok(Self {
            denied: config.deny.into_iter().map(normalize).collect(),
            local: if_addrs::get_if_addrs()?
                .into_iter()
                .map(|a| a.ip())
                .collect(),
            ports: RwLock::default(),
        })
    }
    pub fn deny_host_port(&self, port: u16) {
        self.ports.write().unwrap().insert(port);
    }
    pub fn allows(&self, address: SocketAddr) -> bool {
        let address = normalize(address);
        let ip = address.ip();
        !self.denied.contains(&address)
            && !(self.ports.read().unwrap().contains(&address.port())
                && (ip.is_loopback() || ip.is_unspecified() || self.local.contains(&ip)))
    }
    pub(super) async fn resolve(
        &self,
        uri: &http::Uri,
    ) -> wasmtime_wasi_http::Result<Vec<SocketAddr>> {
        use wasmtime_wasi_http::Error;
        let host = uri
            .host()
            .ok_or(Error::HttpRequestUriInvalid)?
            .trim_matches(['[', ']']);
        let port = uri
            .port_u16()
            .unwrap_or(if uri.scheme_str() == Some("https") {
                443
            } else {
                80
            });
        let addresses: Vec<_> = tokio::net::lookup_host((host, port))
            .await
            .map_err(|_| Error::DnsError {
                rcode: None,
                info_code: None,
            })?
            .collect();
        if addresses.is_empty() || addresses.iter().any(|a| !self.allows(*a)) {
            return Err(Error::HttpRequestDenied);
        }
        Ok(addresses)
    }
}