microsandbox-network 0.7.7

//! Shared state between the NetWorker thread, smoltcp poll thread, and tokio
//! proxy tasks.
//!
//! All inter-thread communication flows through [`SharedState`], which holds
//! lock-free frame queues and cross-platform [`WakePipe`] notifications.

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::{
    Arc, Mutex, OnceLock,
    atomic::{AtomicU64, Ordering},
};
use std::time::{Duration, Instant};

use crossbeam_queue::ArrayQueue;
use ipnetwork::Ipv6Network;
use microsandbox_types::HttpConfig;
use microsandbox_utils::ttl_reverse_index::TtlReverseIndex;
pub use microsandbox_utils::wake_pipe::WakePipe;
use parking_lot::RwLock;

use crate::addr::normalize_ip_addr;
use crate::engine::http_deny::{self, DEFAULT_HTTP_DENY_MESSAGE};

//--------------------------------------------------------------------------------------------------
// Constants
//--------------------------------------------------------------------------------------------------

/// Default frame queue capacity. Matches libkrun's virtio queue size.
pub const DEFAULT_QUEUE_CAPACITY: usize = 1024;

/// Maximum live hostname/address bindings retained by one sandbox.
pub(crate) const MAX_RESOLVED_HOSTNAME_BINDINGS: usize = 16_384;

/// Maximum live bindings for one normalized hostname, across both families.
pub(crate) const MAX_BINDINGS_PER_HOSTNAME: usize = 1_024;

//--------------------------------------------------------------------------------------------------
// Types
//--------------------------------------------------------------------------------------------------

/// All shared state between the three threads:
///
/// - **NetWorker** (libkrun) — pushes guest frames to `tx_ring`, pops
///   response frames from `rx_ring`.
/// - **smoltcp poll thread** — pops from `tx_ring`, processes through smoltcp,
///   pushes responses to `rx_ring`.
/// - **tokio proxy tasks** — relay data between smoltcp sockets and real
///   network connections.
///
/// Queue naming follows the **guest's perspective** (matching libkrun's
/// convention): `tx_ring` = "transmit from guest", `rx_ring` = "receive at
/// guest".
pub struct SharedState {
    /// Frames from guest → smoltcp (NetWorker writes, smoltcp reads).
    pub tx_ring: ArrayQueue<Vec<u8>>,

    /// Frames from smoltcp → guest (smoltcp writes, NetWorker reads).
    pub rx_ring: ArrayQueue<Vec<u8>>,

    /// Wakes NetWorker: "rx_ring has frames for the guest."
    /// Written by `SmoltcpDevice::transmit()`. Read end polled by NetWorker's
    /// epoll loop.
    pub rx_wake: WakePipe,

    /// Wakes smoltcp poll thread: "tx_ring has frames from the guest."
    /// Written by `SmoltcpBackend::write_frame()`. Read end polled by the
    /// poll loop.
    pub tx_wake: WakePipe,

    /// Wakes smoltcp poll thread: "proxy task has data to write to a smoltcp
    /// socket." Written by proxy tasks via channels. Read end polled by the
    /// poll loop.
    pub proxy_wake: WakePipe,

    /// Optional host-side termination hook used for fatal policy violations.
    termination_hook: Mutex<Option<Arc<dyn Fn() + Send + Sync>>>,

    /// Resolved hostname index used to map destination IPs back to queried hostnames.
    resolved_hostnames: RwLock<TtlReverseIndex<ResolvedHostnameKey, IpAddr>>,

    /// Per-sandbox gateway IPv4. Set once at boot; used by
    /// `DestinationGroup::Host` rule matching and `host.microsandbox.internal`
    /// DNS synthesis. `None` in isolated unit tests.
    gateway_ipv4: OnceLock<Ipv4Addr>,

    /// Per-sandbox gateway IPv6. Set once at boot. See `gateway_ipv4`.
    gateway_ipv6: OnceLock<Ipv6Addr>,

    /// NAT64 `/96` prefixes used by policy classification.
    nat64_prefixes: OnceLock<Vec<Ipv6Network>>,

    /// Aggregate network byte counters at the guest/runtime boundary.
    metrics: NetworkMetrics,

    /// HTTP denial settings installed before the network starts.
    http: OnceLock<HttpConfig>,
}

/// Aggregate network byte counters shared with the runtime metrics sampler.
pub struct NetworkMetrics {
    tx_bytes: AtomicU64,
    rx_bytes: AtomicU64,
}

/// Address family for resolved hostname entries.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum ResolvedHostnameFamily {
    Ipv4,
    Ipv6,
}

/// Composite cache key for a single DNS resolution.
///
/// `family` partitions entries so that `A` and `AAAA` responses for the
/// same hostname refresh independently instead of overwriting each other.
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct ResolvedHostnameKey {
    hostname: String,
    family: ResolvedHostnameFamily,
}

//--------------------------------------------------------------------------------------------------
// Methods
//--------------------------------------------------------------------------------------------------

impl SharedState {
    /// Create shared state with the given queue capacity.
    pub fn new(queue_capacity: usize) -> Self {
        Self {
            tx_ring: ArrayQueue::new(queue_capacity),
            rx_ring: ArrayQueue::new(queue_capacity),
            rx_wake: WakePipe::new(),
            tx_wake: WakePipe::new(),
            proxy_wake: WakePipe::new(),
            termination_hook: Mutex::new(None),
            resolved_hostnames: RwLock::new(TtlReverseIndex::default()),
            gateway_ipv4: OnceLock::new(),
            gateway_ipv6: OnceLock::new(),
            nat64_prefixes: OnceLock::new(),
            metrics: NetworkMetrics::default(),
            http: OnceLock::new(),
        }
    }

    /// Set the per-sandbox gateway IPs. Called once at boot. Each family is
    /// only published when active for this sandbox.
    pub fn set_gateway_ips(&self, ipv4: Option<Ipv4Addr>, ipv6: Option<Ipv6Addr>) {
        if let Some(ipv4) = ipv4 {
            let _ = self.gateway_ipv4.set(ipv4);
        }
        if let Some(ipv6) = ipv6 {
            let _ = self.gateway_ipv6.set(ipv6);
        }
    }

    /// Gateway IPv4 address, if set.
    pub fn gateway_ipv4(&self) -> Option<Ipv4Addr> {
        self.gateway_ipv4.get().copied()
    }

    /// Gateway IPv6 address, if set.
    pub fn gateway_ipv6(&self) -> Option<Ipv6Addr> {
        self.gateway_ipv6.get().copied()
    }

    /// Install HTTP denial settings. The first call wins.
    pub fn set_http_config(&self, config: HttpConfig) {
        let _ = self.http.set(config);
    }

    /// Whether readable denial responses were explicitly enabled.
    pub fn http_deny_response_enabled(&self) -> bool {
        self.http.get().is_some_and(|http| http.deny_response)
    }

    /// Render the HTTP/HTTPS deny body for `host`.
    pub fn http_deny_body(&self, host: &str) -> String {
        let template = self
            .http
            .get()
            .and_then(|http| http.deny_message.as_deref())
            .unwrap_or(DEFAULT_HTTP_DENY_MESSAGE);
        http_deny::render_http_deny_message(template, host)
    }

    /// Set NAT64 prefixes. Called once before policy evaluation starts.
    pub fn set_nat64_prefixes(&self, prefixes: Vec<Ipv6Network>) {
        let _ = self.nat64_prefixes.set(prefixes);
    }

    /// NAT64 prefixes used by policy classification.
    pub fn nat64_prefixes(&self) -> &[Ipv6Network] {
        self.nat64_prefixes.get().map(Vec::as_slice).unwrap_or(&[])
    }

    /// Install a host-side termination hook.
    pub fn set_termination_hook(&self, hook: Arc<dyn Fn() + Send + Sync>) {
        *self.termination_hook.lock().unwrap() = Some(hook);
    }

    /// Trigger host-side termination if a hook is installed.
    pub fn trigger_termination(&self) {
        let hook = self.termination_hook.lock().unwrap().clone();
        if let Some(hook) = hook {
            hook();
        }
    }

    /// Record resolved addresses for a hostname within the given address family.
    ///
    /// Addresses from earlier answers stay bound until the TTL of the answer
    /// that returned them expires.
    /// Resolvers that rotate their answers hand concurrent lookups of one name
    /// different addresses, and a connection dialed on any of them must still
    /// match the name's domain rules.
    ///
    /// Returns `false` if the complete answer would exceed either the hostname
    /// or sandbox binding limit; existing bindings are preserved on rejection.
    pub fn cache_resolved_hostname(
        &self,
        domain: &str,
        family: ResolvedHostnameFamily,
        addrs: impl IntoIterator<Item = IpAddr>,
        ttl: Duration,
    ) -> bool {
        let hostname = normalize_hostname(domain);
        let key = ResolvedHostnameKey { hostname, family };
        let addrs = addrs.into_iter().map(normalize_ip_addr);

        let now = Instant::now();
        let mut index = self.resolved_hostnames.write();
        index.evict_expired(now);

        let bindings_for = |family| {
            index.member_count(&ResolvedHostnameKey {
                hostname: key.hostname.clone(),
                family,
            })
        };

        let ipv4_bindings = bindings_for(ResolvedHostnameFamily::Ipv4);
        let ipv6_bindings = bindings_for(ResolvedHostnameFamily::Ipv6);
        let hostname_bindings = ipv4_bindings + ipv6_bindings;

        let total_bindings = index.binding_count();
        let hostname_slots = MAX_BINDINGS_PER_HOSTNAME.saturating_sub(hostname_bindings);
        let sandbox_slots = MAX_RESOLVED_HOSTNAME_BINDINGS.saturating_sub(total_bindings);
        let available_slots = hostname_slots.min(sandbox_slots);

        // try_extend takes a total capacity; only new bindings consume slots.
        let capacity = total_bindings + available_slots;
        index.try_extend(key, addrs, ttl, now, capacity)
    }

    /// Clear the resolved addresses for a hostname within the given address family.
    pub fn clear_resolved_hostname(&self, domain: &str, family: ResolvedHostnameFamily) {
        let hostname = normalize_hostname(domain);
        let key = ResolvedHostnameKey { hostname, family };
        self.resolved_hostnames.write().remove(&key, Instant::now());
    }

    /// Returns `true` when any resolved hostname for `addr` satisfies `predicate`.
    pub fn any_resolved_hostname(
        &self,
        addr: IpAddr,
        mut predicate: impl FnMut(&str) -> bool,
    ) -> bool {
        let addr = normalize_ip_addr(addr);

        self.resolved_hostnames
            .read()
            .member_matches(&addr, Instant::now(), |key| predicate(&key.hostname))
    }

    /// Best-effort expiry maintenance for resolved hostnames.
    ///
    /// This runs outside the hot egress read path. If the index is currently
    /// busy, cleanup is skipped and retried on the next maintenance pass.
    pub fn cleanup_resolved_hostnames(&self) {
        if let Some(mut idx) = self.resolved_hostnames.try_write() {
            idx.evict_expired(Instant::now());
        }
    }

    /// Increment the guest -> runtime byte counter.
    pub fn add_tx_bytes(&self, bytes: usize) {
        self.metrics
            .tx_bytes
            .fetch_add(bytes as u64, Ordering::Relaxed);
    }

    /// Increment the runtime -> guest byte counter.
    pub fn add_rx_bytes(&self, bytes: usize) {
        self.metrics
            .rx_bytes
            .fetch_add(bytes as u64, Ordering::Relaxed);
    }

    /// Push a runtime -> guest ethernet frame and update RX metrics on success.
    pub(crate) fn push_rx_frame(&self, frame: Vec<u8>) -> bool {
        let frame_len = frame.len();
        if self.rx_ring.push(frame).is_err() {
            return false;
        }

        self.add_rx_bytes(frame_len);
        true
    }

    /// Push a runtime -> guest ethernet frame, update RX metrics, and wake libkrun.
    pub(crate) fn push_rx_frame_and_wake(&self, frame: Vec<u8>) -> bool {
        if !self.push_rx_frame(frame) {
            return false;
        }

        self.rx_wake.wake();
        true
    }

    /// Total bytes transmitted by the guest into the runtime.
    pub fn tx_bytes(&self) -> u64 {
        self.metrics.tx_bytes.load(Ordering::Relaxed)
    }

    /// Total bytes delivered by the runtime to the guest.
    pub fn rx_bytes(&self) -> u64 {
        self.metrics.rx_bytes.load(Ordering::Relaxed)
    }
}

impl Default for NetworkMetrics {
    fn default() -> Self {
        Self {
            tx_bytes: AtomicU64::new(0),
            rx_bytes: AtomicU64::new(0),
        }
    }
}

pub(crate) fn normalize_hostname(domain: &str) -> String {
    domain.trim_end_matches('.').to_ascii_lowercase()
}

//--------------------------------------------------------------------------------------------------
// Tests
//--------------------------------------------------------------------------------------------------

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

    #[test]
    fn shared_state_queue_push_pop() {
        let state = SharedState::new(4);

        // Push frames to tx_ring.
        state.tx_ring.push(vec![1, 2, 3]).unwrap();
        state.tx_ring.push(vec![4, 5, 6]).unwrap();

        // Pop in FIFO order.
        assert_eq!(state.tx_ring.pop(), Some(vec![1, 2, 3]));
        assert_eq!(state.tx_ring.pop(), Some(vec![4, 5, 6]));
        assert_eq!(state.tx_ring.pop(), None);
    }

    #[test]
    fn shared_state_queue_full() {
        let state = SharedState::new(2);

        state.rx_ring.push(vec![1]).unwrap();
        state.rx_ring.push(vec![2]).unwrap();
        // Queue is full — push returns the frame back.
        assert!(state.rx_ring.push(vec![3]).is_err());
    }

    #[test]
    fn push_rx_frame_counts_only_successful_pushes() {
        let state = SharedState::new(1);

        assert!(state.push_rx_frame(vec![1, 2, 3]));
        assert_eq!(state.rx_bytes(), 3);

        assert!(!state.push_rx_frame(vec![4, 5]));
        assert_eq!(state.rx_bytes(), 3);
    }

    #[test]
    fn resolved_hostnames_are_isolated_per_family() {
        let state = SharedState::new(4);
        let v4: IpAddr = "1.1.1.1".parse().unwrap();
        let v6: IpAddr = "2606:4700:4700::1111".parse().unwrap();

        state.cache_resolved_hostname(
            "Example.com.",
            ResolvedHostnameFamily::Ipv4,
            [v4],
            Duration::from_secs(30),
        );
        state.cache_resolved_hostname(
            "example.com",
            ResolvedHostnameFamily::Ipv6,
            [v6],
            Duration::from_secs(30),
        );

        assert!(state.any_resolved_hostname(v4, |h| h == "example.com"));
        assert!(state.any_resolved_hostname(v6, |h| h == "example.com"));
        assert!(!state.any_resolved_hostname(v4, |h| h == "other.example"));
    }

    #[test]
    fn resolved_hostnames_keep_earlier_answers() {
        let state = SharedState::new(4);
        let first: IpAddr = "142.250.0.1".parse().unwrap();
        let second: IpAddr = "142.250.0.2".parse().unwrap();

        // Rotating answers must remain usable even beyond the old 64-address
        // limit, as long as their TTLs have not expired.
        for n in 1..=65 {
            let addr = IpAddr::V4(Ipv4Addr::new(142, 250, 0, n));
            state.cache_resolved_hostname(
                "fonts.example",
                ResolvedHostnameFamily::Ipv4,
                [addr],
                Duration::from_secs(30),
            );
        }

        for n in 1..=65 {
            let addr = IpAddr::V4(Ipv4Addr::new(142, 250, 0, n));
            assert!(state.any_resolved_hostname(addr, |h| h == "fonts.example"));
        }

        state.clear_resolved_hostname("fonts.example", ResolvedHostnameFamily::Ipv4);
        assert!(!state.any_resolved_hostname(first, |h| h == "fonts.example"));
        assert!(!state.any_resolved_hostname(second, |h| h == "fonts.example"));
    }

    #[test]
    fn resolved_hostnames_normalize_ipv4_mapped_ipv6() {
        let state = SharedState::new(4);
        let mapped: IpAddr = "::ffff:169.254.169.254".parse().unwrap();
        let embedded: IpAddr = "169.254.169.254".parse().unwrap();

        state.cache_resolved_hostname(
            "metadata.example",
            ResolvedHostnameFamily::Ipv6,
            [mapped],
            Duration::from_secs(30),
        );

        assert!(state.any_resolved_hostname(embedded, |h| h == "metadata.example"));
        assert!(state.any_resolved_hostname(mapped, |h| h == "metadata.example"));
    }
}