panda-rpki-oss/src/validation/tree.rs

use crate::model::rc::{AsResourceSet, IpResourceSet};
use crate::output::audit::{DiscoveredFrom, PublicationPointAudit};
use crate::output::report::Warning;
use crate::validation::manifest::PublicationPointSource;
use crate::validation::objects::{AspaAttestation, ObjectsOutput, RouterKeyPayload, Vrp};
use crate::validation::publication_point::PublicationPointSnapshot;
use std::borrow::Cow;
use std::sync::{Arc, OnceLock};

pub const DEFAULT_MAX_CA_DEPTH: usize = 32;

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TreeRunConfig {
    /// Max CA instance depth to process (0 = root only).
    pub max_depth: Option<usize>,
    /// Max number of CA instances to process.
    pub max_instances: Option<usize>,
    /// Drop per-object audit payload when the caller only needs validation outputs.
    pub compact_audit: bool,
    /// Build online CCR manifest projections during phase2 validation.
    pub build_ccr_accumulator: bool,
}

impl Default for TreeRunConfig {
    fn default() -> Self {
        Self {
            max_depth: Some(DEFAULT_MAX_CA_DEPTH),
            max_instances: None,
            compact_audit: false,
            build_ccr_accumulator: true,
        }
    }
}

pub(crate) fn ca_depth_is_allowed(config: &TreeRunConfig, depth: usize) -> bool {
    config.max_depth.is_none_or(|max_depth| depth <= max_depth)
}

pub(crate) fn next_allowed_ca_depth(config: &TreeRunConfig, parent_depth: usize) -> Option<usize> {
    let child_depth = parent_depth.checked_add(1)?;
    ca_depth_is_allowed(config, child_depth).then_some(child_depth)
}

#[derive(Clone, Debug)]
pub enum CaCertificateRef {
    InlineDer(Vec<u8>),
    RepoBytes {
        sha256_hex: String,
        cached_der: Arc<OnceLock<Arc<[u8]>>>,
    },
}

impl PartialEq for CaCertificateRef {
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            (Self::InlineDer(left), Self::InlineDer(right)) => left == right,
            (
                Self::RepoBytes {
                    sha256_hex: left, ..
                },
                Self::RepoBytes {
                    sha256_hex: right, ..
                },
            ) => left == right,
            _ => false,
        }
    }
}

impl Eq for CaCertificateRef {}

impl CaCertificateRef {
    pub fn inline_der(bytes: Vec<u8>) -> Self {
        Self::InlineDer(bytes)
    }

    pub fn repo_bytes(sha256_hex: String) -> Self {
        Self::RepoBytes {
            sha256_hex,
            cached_der: Arc::new(OnceLock::new()),
        }
    }

    pub fn sha256_hex(&self) -> Option<&str> {
        match self {
            Self::InlineDer(_) => None,
            Self::RepoBytes { sha256_hex, .. } => Some(sha256_hex.as_str()),
        }
    }

    pub fn der<'a>(
        &'a self,
        store: &crate::repository::storage::RocksStore,
    ) -> Result<Cow<'a, [u8]>, String> {
        match self {
            Self::InlineDer(bytes) => Ok(Cow::Borrowed(bytes.as_slice())),
            Self::RepoBytes {
                sha256_hex,
                cached_der,
            } => {
                if cached_der.get().is_none() {
                    let bytes = store
                        .get_blob_bytes(sha256_hex)
                        .map_err(|e| format!("load CA certificate bytes failed: {e}"))?
                        .ok_or_else(|| {
                            format!("missing CA certificate repo bytes for sha256={sha256_hex}")
                        })?;
                    let _ = cached_der.set(Arc::from(bytes));
                }
                let bytes = cached_der.get().ok_or_else(|| {
                    format!("missing cached CA certificate bytes for sha256={sha256_hex}")
                })?;
                Ok(Cow::Borrowed(bytes.as_ref()))
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CaInstanceHandle {
    pub depth: usize,
    pub tal_id: String,
    pub parent_manifest_rsync_uri: Option<String>,
    /// CA certificate bytes or a lazy repo-bytes reference for this CA instance.
    pub ca_certificate: CaCertificateRef,
    /// rsync URI of this CA certificate object (where it is published).
    ///
    /// This is used for strict AIA binding checks (RFC 6487 §4.8.7) when validating
    /// subordinate CA and EE certificates. Trust anchor certificates downloaded from
    /// TAL URIs may not have an rsync publication URI, so this can be `None`.
    pub ca_certificate_rsync_uri: Option<String>,
    /// Effective (fully resolved) resources for this CA instance, used for resource-path
    /// validation of subordinate CA certificates.
    pub effective_ip_resources: Option<IpResourceSet>,
    pub effective_as_resources: Option<AsResourceSet>,
    pub rsync_base_uri: String,
    pub manifest_rsync_uri: String,
    pub publication_point_rsync_uri: String,
    pub rrdp_notification_uri: Option<String>,
}

impl CaInstanceHandle {
    pub fn with_depth(mut self, depth: usize) -> Self {
        self.depth = depth;
        self
    }

    pub fn ca_certificate_der<'a>(
        &'a self,
        store: &crate::repository::storage::RocksStore,
    ) -> Result<Cow<'a, [u8]>, String> {
        self.ca_certificate.der(store)
    }

    pub fn ca_certificate_sha256_hex(&self) -> Option<&str> {
        self.ca_certificate.sha256_hex()
    }

    pub fn ca_certificate_sha256_32(&self) -> Option<[u8; 32]> {
        match &self.ca_certificate {
            CaCertificateRef::InlineDer(bytes) => {
                use sha2::Digest as _;
                let digest = sha2::Sha256::digest(bytes);
                Some(digest.into())
            }
            CaCertificateRef::RepoBytes { sha256_hex, .. } => {
                let mut out = [0u8; 32];
                hex::decode_to_slice(sha256_hex, &mut out).ok()?;
                Some(out)
            }
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PublicationPointRunResult {
    pub source: PublicationPointSource,
    pub snapshot: Option<PublicationPointSnapshot>,
    pub warnings: Vec<Warning>,
    pub objects: ObjectsOutput,
    pub audit: PublicationPointAudit,
    /// Candidate child CA instances to enqueue after this publication point completes.
    ///
    /// - For `Fresh`, these are discovered from the current validated publication point.
    /// - For failed synchronization, the result contains no child instances;
    ///   callers may retry the repository on a later run.
    pub discovered_children: Vec<DiscoveredChildCaInstance>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DiscoveredChildEntryProjection {
    pub child_ski: String,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DiscoveredChildCaInstance {
    pub handle: CaInstanceHandle,
    pub discovered_from: DiscoveredFrom,
    pub child_entry_projection: Option<DiscoveredChildEntryProjection>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TreeRunOutput {
    pub instances_processed: usize,
    pub instances_failed: usize,
    pub warnings: Vec<Warning>,
    pub vrps: Vec<Vrp>,
    pub aspas: Vec<AspaAttestation>,
    pub router_keys: Vec<RouterKeyPayload>,
}

#[derive(Debug, thiserror::Error)]
pub enum TreeRunError {
    #[error("publication point runner failed: {0}")]
    Runner(String),
}

pub trait PublicationPointRunner {
    fn run_publication_point(
        &self,
        ca: &CaInstanceHandle,
    ) -> Result<PublicationPointRunResult, String>;
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TreeRunAuditOutput {
    pub tree: TreeRunOutput,
    pub publication_points: Vec<PublicationPointAudit>,
}

pub fn run_tree_serial(
    root: CaInstanceHandle,
    runner: &dyn PublicationPointRunner,
    config: &TreeRunConfig,
) -> Result<TreeRunOutput, TreeRunError> {
    Ok(run_tree_serial_audit(root, runner, config)?.tree)
}

pub fn run_tree_serial_audit(
    root: CaInstanceHandle,
    runner: &dyn PublicationPointRunner,
    config: &TreeRunConfig,
) -> Result<TreeRunAuditOutput, TreeRunError> {
    run_tree_serial_audit_multi_root(vec![root], runner, config)
}

pub fn run_tree_serial_audit_multi_root(
    roots: Vec<CaInstanceHandle>,
    runner: &dyn PublicationPointRunner,
    config: &TreeRunConfig,
) -> Result<TreeRunAuditOutput, TreeRunError> {
    #[derive(Clone, Debug)]
    struct QueuedCaInstance {
        id: u64,
        handle: CaInstanceHandle,
        parent_id: Option<u64>,
        discovered_from: Option<DiscoveredFrom>,
    }

    let mut next_id: u64 = 0;
    let mut queue: std::collections::VecDeque<QueuedCaInstance> = std::collections::VecDeque::new();
    for root in roots {
        queue.push_back(QueuedCaInstance {
            id: next_id,
            handle: root,
            parent_id: None,
            discovered_from: None,
        });
        next_id += 1;
    }

    let mut visited_manifest_uris: std::collections::HashSet<String> =
        std::collections::HashSet::new();
    let mut instances_processed = 0usize;
    let mut instances_failed = 0usize;
    let mut warnings: Vec<Warning> = Vec::new();
    let mut vrps: Vec<Vrp> = Vec::new();
    let mut aspas: Vec<AspaAttestation> = Vec::new();
    let mut router_keys: Vec<RouterKeyPayload> = Vec::new();
    let mut publication_points: Vec<PublicationPointAudit> = Vec::new();

    while let Some(node) = queue.pop_front() {
        let ca = &node.handle;
        if !visited_manifest_uris.insert(ca.manifest_rsync_uri.clone()) {
            continue;
        }

        if let Some(max_instances) = config.max_instances
            && instances_processed >= max_instances
        {
            break;
        }

        if !ca_depth_is_allowed(config, ca.depth) {
            continue;
        }

        let res = match runner.run_publication_point(ca) {
            Ok(v) => v,
            Err(e) => {
                instances_failed += 1;
                warnings.push(
                    Warning::new(format!("publication point failed: {e}"))
                        .with_context(&ca.manifest_rsync_uri),
                );
                continue;
            }
        };

        instances_processed += 1;
        crate::logging::observer::update(instances_processed as u64, || {
            serde_json::json!({
                "publication_points_completed": instances_processed,
                "publication_points_failed": instances_failed,
            })
        });
        warnings.extend(res.warnings);
        warnings.extend(res.objects.warnings);
        vrps.extend(res.objects.vrps);
        aspas.extend(res.objects.aspas);
        router_keys.extend(res.objects.router_keys);

        let mut audit = res.audit;
        audit.node_id = Some(node.id);
        audit.parent_node_id = node.parent_id;
        audit.discovered_from = node.discovered_from.clone();
        if config.compact_audit {
            audit.objects.clear();
            audit.warnings.clear();
        }
        publication_points.push(audit);

        let Some(child_depth) = next_allowed_ca_depth(config, ca.depth) else {
            continue;
        };
        let mut children = res.discovered_children;
        children.sort_by(|a, b| {
            a.handle
                .manifest_rsync_uri
                .cmp(&b.handle.manifest_rsync_uri)
                .then_with(|| {
                    a.discovered_from
                        .child_ca_certificate_rsync_uri
                        .cmp(&b.discovered_from.child_ca_certificate_rsync_uri)
                })
        });
        for child in children {
            queue.push_back(QueuedCaInstance {
                id: next_id,
                handle: child.handle.with_depth(child_depth),
                parent_id: Some(node.id),
                discovered_from: Some(child.discovered_from),
            });
            next_id += 1;
        }
    }

    Ok(TreeRunAuditOutput {
        tree: TreeRunOutput {
            instances_processed,
            instances_failed,
            warnings,
            vrps,
            aspas,
            router_keys,
        },
        publication_points,
    })
}