lazaret/crates/lazaret-core/src/network.rs

use std::collections::{BTreeMap, BTreeSet};
use std::fmt;

use sha2::{Digest, Sha256};

use crate::codec::{CodecError, encode};
use crate::configuration::ResolvedDomainConfiguration;
use crate::identity::{
    ConfigurationFingerprint, DerivedPolicyDigest, DomainId, GenerationId, NamespaceIdentity,
    NetworkSidecarId, PidfsProcessIdentityV1, ReadinessFingerprint,
};
use crate::records::{
    AddressFamily, AddressFamilyDisposition, AddressFamilyDispositionValue,
    DERIVED_NETWORK_POLICY_SCHEMA_ID, DefaultAction, DerivedNetworkPolicy,
    DerivedNetworkPolicyFields, IpAddress, IpNetwork, NETWORK_READINESS_IDENTITY_SCHEMA_ID,
    NetworkObservation, NetworkReachabilityProbe, NetworkReadinessAttestation,
    NetworkReadinessEvidence, NetworkReadinessIdentity, NetworkReadinessIdentityFields,
    PolicyAction, PolicyDelivery, PrefixRule, Resolver,
};

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CanonicalDerivedNetworkPolicy {
    policy: DerivedNetworkPolicy,
    canonical_bytes: Vec<u8>,
    digest: DerivedPolicyDigest,
}

impl CanonicalDerivedNetworkPolicy {
    #[must_use]
    pub const fn policy(&self) -> &DerivedNetworkPolicy {
        &self.policy
    }

    #[must_use]
    pub fn canonical_bytes(&self) -> &[u8] {
        &self.canonical_bytes
    }

    #[must_use]
    pub const fn digest(&self) -> DerivedPolicyDigest {
        self.digest
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ReachabilityExpectation {
    pub denied_destinations: Vec<IpAddress>,
    pub permitted_destinations: Vec<IpAddress>,
    pub dns_queries: Vec<DnsQueryExpectation>,
}

#[derive(Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
pub struct DnsQueryExpectation {
    pub resolver: IpAddress,
    pub query_name: String,
}

#[derive(Clone, Debug)]
pub struct ReadinessExpectation<'a> {
    pub domain_id: DomainId,
    pub generation_id: GenerationId,
    pub configuration_fingerprint: ConfigurationFingerprint,
    pub network_sidecar_id: NetworkSidecarId,
    pub enforcement_namespace: NamespaceIdentity,
    pub runtime_supervisor: &'a PidfsProcessIdentityV1,
    pub connectivity_helper: &'a PidfsProcessIdentityV1,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum NetworkValidationErrorKind {
    AddressFamilies,
    DefaultActions,
    Prefix,
    PrefixConflict,
    Resolver,
    CanonicalEncoding,
    Attestation,
    Reachability,
    Evidence,
}

#[derive(Debug)]
pub struct NetworkValidationError {
    kind: NetworkValidationErrorKind,
    detail: &'static str,
}

impl NetworkValidationError {
    const fn new(kind: NetworkValidationErrorKind, detail: &'static str) -> Self {
        Self { kind, detail }
    }

    #[must_use]
    pub const fn kind(&self) -> NetworkValidationErrorKind {
        self.kind
    }
}

impl fmt::Display for NetworkValidationError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(formatter, "{}", self.detail)
    }
}

impl std::error::Error for NetworkValidationError {}

impl From<CodecError> for NetworkValidationError {
    fn from(_: CodecError) -> Self {
        Self::new(
            NetworkValidationErrorKind::CanonicalEncoding,
            "network document cannot be canonically encoded",
        )
    }
}

/// Derives the v1 policy from one resolved trusted-host configuration.
///
/// # Errors
///
/// Returns an error when resolver or canonical policy validation fails.
pub fn derive_ipv4_policy(
    configuration: &ResolvedDomainConfiguration,
    resolvers: impl IntoIterator<Item = IpAddress>,
) -> Result<CanonicalDerivedNetworkPolicy, NetworkValidationError> {
    let prefix_rules = configuration
        .effective_allow_prefixes()
        .iter()
        .map(|prefix| PrefixRule {
            family: AddressFamily::Ipv4,
            network: IpNetwork {
                address: IpAddress::V4(prefix.address()),
                prefix_length: prefix.prefix_length(),
            },
            action: PolicyAction::Allow,
        })
        .chain(
            configuration
                .effective_deny_prefixes()
                .iter()
                .map(|prefix| PrefixRule {
                    family: AddressFamily::Ipv4,
                    network: IpNetwork {
                        address: IpAddress::V4(prefix.address()),
                        prefix_length: prefix.prefix_length(),
                    },
                    action: PolicyAction::Deny,
                }),
        )
        .collect();
    derive_policy(DerivedNetworkPolicyFields {
        address_families: vec![
            AddressFamilyDisposition {
                family: AddressFamily::Ipv4,
                disposition: AddressFamilyDispositionValue::Enabled,
            },
            AddressFamilyDisposition {
                family: AddressFamily::Ipv6,
                disposition: AddressFamilyDispositionValue::Disabled,
            },
        ],
        default_actions: vec![DefaultAction {
            family: AddressFamily::Ipv4,
            action: PolicyAction::Allow,
        }],
        prefix_rules,
        resolvers: resolvers
            .into_iter()
            .map(|address| Resolver {
                family: AddressFamily::Ipv4,
                address,
            })
            .collect(),
    })
}

/// Validates, normalizes, and fingerprints a derived policy document.
///
/// # Errors
///
/// Returns an error for an invalid family, prefix, resolver, conflict, or encoding.
pub fn derive_policy(
    mut fields: DerivedNetworkPolicyFields,
) -> Result<CanonicalDerivedNetworkPolicy, NetworkValidationError> {
    validate_address_families(&fields.address_families)?;
    validate_default_actions(&fields.default_actions)?;
    validate_prefix_rules(&fields.prefix_rules)?;
    validate_resolvers(&fields.resolvers)?;

    fields
        .address_families
        .sort_unstable_by_key(|item| item.family);
    fields
        .default_actions
        .sort_unstable_by_key(|item| item.family);
    fields.prefix_rules.sort_unstable();
    fields.resolvers.sort_unstable();

    let policy = DerivedNetworkPolicy::v1(fields);
    let canonical_bytes = encode(&policy)?;
    let digest = domain_digest(DERIVED_NETWORK_POLICY_SCHEMA_ID, &canonical_bytes);
    Ok(CanonicalDerivedNetworkPolicy {
        policy,
        canonical_bytes,
        digest: DerivedPolicyDigest::new(digest),
    })
}

#[must_use]
pub fn readiness_identity(attestation: &NetworkReadinessAttestation) -> NetworkReadinessIdentity {
    NetworkReadinessIdentity::v1(NetworkReadinessIdentityFields {
        domain_id: attestation.domain_id,
        generation_id: attestation.generation_id,
        configuration_fingerprint: attestation.configuration_fingerprint,
        network_sidecar_id: attestation.network_sidecar_id,
        derived_policy_digest: attestation.derived_policy_digest,
        enforcement_namespace: attestation.enforcement_namespace,
    })
}

/// Encodes the v1 readiness identity in its canonical form.
///
/// # Errors
///
/// Returns an error when canonical encoding fails.
pub fn canonical_readiness_identity_bytes(
    identity: &NetworkReadinessIdentity,
) -> Result<Vec<u8>, NetworkValidationError> {
    Ok(encode(identity)?)
}

/// Computes the domain-separated readiness fingerprint.
///
/// # Errors
///
/// Returns an error when canonical identity encoding fails.
pub fn readiness_fingerprint(
    identity: &NetworkReadinessIdentity,
) -> Result<ReadinessFingerprint, NetworkValidationError> {
    let canonical_bytes = canonical_readiness_identity_bytes(identity)?;
    Ok(ReadinessFingerprint::new(domain_digest(
        NETWORK_READINESS_IDENTITY_SCHEMA_ID,
        &canonical_bytes,
    )))
}

/// Binds validated attestation and reachability results into readiness evidence.
///
/// # Errors
///
/// Returns an error for any identity, observation, policy, or probe mismatch.
pub fn accept_network_readiness(
    policy: &CanonicalDerivedNetworkPolicy,
    expectation: &ReadinessExpectation<'_>,
    attestation: NetworkReadinessAttestation,
    reachability: NetworkReachabilityProbe,
    expected_reachability: &ReachabilityExpectation,
) -> Result<NetworkReadinessEvidence, NetworkValidationError> {
    validate_attestation(policy, expectation, &attestation)?;
    validate_reachability(
        policy.policy(),
        expectation,
        &reachability,
        expected_reachability,
    )?;
    let identity = readiness_identity(&attestation);
    let fingerprint = readiness_fingerprint(&identity)?;
    Ok(NetworkReadinessEvidence {
        domain_id: expectation.domain_id,
        generation_id: expectation.generation_id,
        network_sidecar_id: expectation.network_sidecar_id,
        readiness_fingerprint: fingerprint,
        attestation,
        readiness_identity: identity,
        reachability_probe: reachability,
    })
}

/// Validates Runtime Supervisor attestation before the sandbox-view probe is available.
///
/// # Errors
/// Returns an error for any identity, observation, policy, or process mismatch.
pub fn validate_network_attestation(
    policy: &CanonicalDerivedNetworkPolicy,
    expectation: &ReadinessExpectation<'_>,
    attestation: &NetworkReadinessAttestation,
) -> Result<(), NetworkValidationError> {
    validate_attestation(policy, expectation, attestation)
}

/// Revalidates persisted readiness evidence against current generation authority.
///
/// # Errors
///
/// Returns an error for any identity, projection, fingerprint, or probe mismatch.
pub fn validate_network_readiness_evidence(
    policy: &CanonicalDerivedNetworkPolicy,
    expectation: &ReadinessExpectation<'_>,
    evidence: &NetworkReadinessEvidence,
    expected_reachability: &ReachabilityExpectation,
) -> Result<(), NetworkValidationError> {
    if evidence.domain_id != expectation.domain_id
        || evidence.generation_id != expectation.generation_id
        || evidence.network_sidecar_id != expectation.network_sidecar_id
    {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::Evidence,
            "network readiness evidence identity mismatch",
        ));
    }
    validate_attestation(policy, expectation, &evidence.attestation)?;
    let projected = readiness_identity(&evidence.attestation);
    let fingerprint = readiness_fingerprint(&projected)?;
    if evidence.readiness_identity != projected || evidence.readiness_fingerprint != fingerprint {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::Evidence,
            "network readiness projection or fingerprint mismatch",
        ));
    }
    validate_reachability(
        policy.policy(),
        expectation,
        &evidence.reachability_probe,
        expected_reachability,
    )
}

fn validate_address_families(
    dispositions: &[AddressFamilyDisposition],
) -> Result<(), NetworkValidationError> {
    let expected = BTreeMap::from([
        (AddressFamily::Ipv4, AddressFamilyDispositionValue::Enabled),
        (AddressFamily::Ipv6, AddressFamilyDispositionValue::Disabled),
    ]);
    let actual: BTreeMap<_, _> = dispositions
        .iter()
        .map(|item| (item.family, item.disposition))
        .collect();
    if dispositions.len() != expected.len() || actual != expected {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::AddressFamilies,
            "v1 requires IPv4 enabled and IPv6 disabled exactly once",
        ));
    }
    Ok(())
}

fn validate_default_actions(actions: &[DefaultAction]) -> Result<(), NetworkValidationError> {
    if actions
        != [DefaultAction {
            family: AddressFamily::Ipv4,
            action: PolicyAction::Allow,
        }]
    {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::DefaultActions,
            "v1 requires one IPv4 default-allow action",
        ));
    }
    Ok(())
}

fn validate_prefix_rules(rules: &[PrefixRule]) -> Result<(), NetworkValidationError> {
    let mut actions = BTreeMap::new();
    for rule in rules {
        validate_network(rule.family, rule.network)?;
        if rule.family != AddressFamily::Ipv4 {
            return Err(NetworkValidationError::new(
                NetworkValidationErrorKind::AddressFamilies,
                "disabled address family has a prefix rule",
            ));
        }
        match actions.insert((rule.family, rule.network), rule.action) {
            Some(previous) if previous != rule.action => {
                return Err(NetworkValidationError::new(
                    NetworkValidationErrorKind::PrefixConflict,
                    "equal prefix has conflicting actions",
                ));
            }
            Some(_) => {
                return Err(NetworkValidationError::new(
                    NetworkValidationErrorKind::Prefix,
                    "duplicate prefix rule",
                ));
            }
            None => {}
        }
    }
    for network in [
        ([10, 0, 0, 0], 8),
        ([172, 16, 0, 0], 12),
        ([192, 168, 0, 0], 16),
    ] {
        let required = IpNetwork {
            address: IpAddress::V4(network.0),
            prefix_length: network.1,
        };
        if actions.get(&(AddressFamily::Ipv4, required)) != Some(&PolicyAction::Deny) {
            return Err(NetworkValidationError::new(
                NetworkValidationErrorKind::Prefix,
                "v1 policy is missing a default private-network deny",
            ));
        }
    }
    Ok(())
}

fn validate_resolvers(resolvers: &[Resolver]) -> Result<(), NetworkValidationError> {
    if resolvers.is_empty() {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::Resolver,
            "derived policy requires a resolver",
        ));
    }
    let mut seen = BTreeSet::new();
    for resolver in resolvers {
        if resolver.family != AddressFamily::Ipv4
            || address_family(resolver.address) != resolver.family
        {
            return Err(NetworkValidationError::new(
                NetworkValidationErrorKind::Resolver,
                "resolver address family is disabled or mismatched",
            ));
        }
        if !seen.insert(*resolver) {
            return Err(NetworkValidationError::new(
                NetworkValidationErrorKind::Resolver,
                "duplicate resolver",
            ));
        }
    }
    Ok(())
}

fn validate_network(
    family: AddressFamily,
    network: IpNetwork,
) -> Result<(), NetworkValidationError> {
    if address_family(network.address) != family {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::Prefix,
            "prefix address family mismatch",
        ));
    }
    let (bytes, maximum) = match network.address {
        IpAddress::V4(bytes) => (bytes.to_vec(), 32),
        IpAddress::V6(bytes) => (bytes.to_vec(), 128),
    };
    if network.prefix_length > maximum || !is_network_address(&bytes, network.prefix_length) {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::Prefix,
            "prefix is not a canonical network",
        ));
    }
    Ok(())
}

fn is_network_address(bytes: &[u8], prefix_length: u8) -> bool {
    let complete = usize::from(prefix_length / 8);
    let remainder = prefix_length % 8;
    if remainder != 0 && bytes[complete] & (u8::MAX >> remainder) != 0 {
        return false;
    }
    let start = complete + usize::from(remainder != 0);
    bytes[start..].iter().all(|byte| *byte == 0)
}

fn validate_attestation(
    policy: &CanonicalDerivedNetworkPolicy,
    expected: &ReadinessExpectation<'_>,
    attestation: &NetworkReadinessAttestation,
) -> Result<(), NetworkValidationError> {
    if attestation.domain_id != expected.domain_id
        || attestation.generation_id != expected.generation_id
        || attestation.configuration_fingerprint != expected.configuration_fingerprint
        || attestation.network_sidecar_id != expected.network_sidecar_id
        || attestation.derived_policy_digest != policy.digest()
        || attestation.enforcement_namespace != expected.enforcement_namespace
        || attestation.runtime_supervisor != *expected.runtime_supervisor
        || attestation.connectivity_helper != *expected.connectivity_helper
        || !valid_observation(&attestation.observation)
    {
        return Err(NetworkValidationError::new(
            NetworkValidationErrorKind::Attestation,
            "network readiness attestation mismatch",
        ));
    }
    for route in &attestation.observation.observed_routes {
        validate_network(address_family(route.address), *route)?;
    }
    Ok(())
}

fn valid_observation(observation: &NetworkObservation) -> bool {
    observation.policy_delivery == PolicyDelivery::SealedMemfd
        && observation.helper_ready
        && observation.routes_attested
}

fn validate_reachability(
    policy: &DerivedNetworkPolicy,
    expected: &ReadinessExpectation<'_>,
    probe: &NetworkReachabilityProbe,
    expected_results: &ReachabilityExpectation,
) -> Result<(), NetworkValidationError> {
    if probe.domain_id != expected.domain_id
        || probe.generation_id != expected.generation_id
        || probe.enforcement_namespace != expected.enforcement_namespace
        || expected_results.denied_destinations.is_empty()
        || expected_results.permitted_destinations.is_empty()
        || expected_results.dns_queries.is_empty()
    {
        return reachability_error();
    }

    let denied = exact_observations(
        &expected_results.denied_destinations,
        &probe.denied_destinations,
        |item| (item.destination, item.unreachable),
    )?;
    let permitted = exact_observations(
        &expected_results.permitted_destinations,
        &probe.permitted_destinations,
        |item| (item.destination, item.reachable),
    )?;
    if denied.iter().any(|(_, passed)| !passed)
        || permitted.iter().any(|(_, passed)| !passed)
        || expected_results
            .denied_destinations
            .iter()
            .any(|address| effective_action(policy, *address) != Some(PolicyAction::Deny))
        || expected_results
            .permitted_destinations
            .iter()
            .any(|address| effective_action(policy, *address) != Some(PolicyAction::Allow))
    {
        return reachability_error();
    }

    let expected_dns: BTreeSet<_> = expected_results.dns_queries.iter().cloned().collect();
    let actual_dns: BTreeSet<_> = probe
        .dns_observations
        .iter()
        .map(|item| DnsQueryExpectation {
            resolver: item.resolver,
            query_name: item.query_name.clone(),
        })
        .collect();
    let resolvers: BTreeSet<_> = policy
        .resolvers
        .iter()
        .map(|resolver| resolver.address)
        .collect();
    if expected_dns.len() != expected_results.dns_queries.len()
        || actual_dns.len() != probe.dns_observations.len()
        || actual_dns != expected_dns
        || probe.dns_observations.iter().any(|item| {
            item.query_name.is_empty()
                || !resolvers.contains(&item.resolver)
                || item.resolved_addresses.is_empty()
                || item
                    .resolved_addresses
                    .iter()
                    .any(|address| address_family(*address) != AddressFamily::Ipv4)
        })
    {
        return reachability_error();
    }
    Ok(())
}

fn exact_observations<T, F>(
    expected: &[IpAddress],
    actual: &[T],
    project: F,
) -> Result<BTreeSet<(IpAddress, bool)>, NetworkValidationError>
where
    F: Fn(&T) -> (IpAddress, bool),
{
    let expected_items: BTreeSet<_> = expected.iter().copied().collect();
    let actual_items: BTreeSet<_> = actual.iter().map(project).collect();
    let actual_addresses: BTreeSet<_> = actual_items.iter().map(|(address, _)| *address).collect();
    if expected_items.len() != expected.len()
        || actual_items.len() != actual.len()
        || expected_items.len() != actual_items.len()
        || expected_items.len() != actual_addresses.len()
        || actual_addresses != expected_items
    {
        return reachability_error();
    }
    Ok(actual_items)
}

fn effective_action(policy: &DerivedNetworkPolicy, address: IpAddress) -> Option<PolicyAction> {
    let family = address_family(address);
    let disposition = policy
        .address_families
        .iter()
        .find(|item| item.family == family)?;
    if disposition.disposition != AddressFamilyDispositionValue::Enabled {
        return None;
    }
    policy
        .prefix_rules
        .iter()
        .filter(|rule| rule.family == family && network_contains(rule.network, address))
        .max_by_key(|rule| rule.network.prefix_length)
        .map(|rule| rule.action)
        .or_else(|| {
            policy
                .default_actions
                .iter()
                .find(|item| item.family == family)
                .map(|item| item.action)
        })
}

fn network_contains(network: IpNetwork, address: IpAddress) -> bool {
    let prefix_length = network.prefix_length;
    let (network, address) = match (network.address, address) {
        (IpAddress::V4(network), IpAddress::V4(address)) => (network.to_vec(), address.to_vec()),
        (IpAddress::V6(network), IpAddress::V6(address)) => (network.to_vec(), address.to_vec()),
        _ => return false,
    };
    let complete = usize::from(prefix_length / 8);
    let remainder = prefix_length % 8;
    if network[..complete] != address[..complete] {
        return false;
    }
    remainder == 0
        || network[complete] & (u8::MAX << (8 - remainder))
            == address[complete] & (u8::MAX << (8 - remainder))
}

const fn address_family(address: IpAddress) -> AddressFamily {
    match address {
        IpAddress::V4(_) => AddressFamily::Ipv4,
        IpAddress::V6(_) => AddressFamily::Ipv6,
    }
}

fn domain_digest(schema_id: &str, canonical_bytes: &[u8]) -> [u8; 32] {
    let mut hasher = Sha256::new();
    hasher.update(schema_id.as_bytes());
    hasher.update([0]);
    hasher.update(canonical_bytes);
    hasher.finalize().into()
}

fn reachability_error<T>() -> Result<T, NetworkValidationError> {
    Err(NetworkValidationError::new(
        NetworkValidationErrorKind::Reachability,
        "network reachability evidence mismatch",
    ))
}