monovm-whois-rust 1.0.0

//! [`LayeredRegistry`]: stack several providers into one.

use std::sync::Arc;

use crate::domain::Tld;
use crate::registry::{Registry, RegistryProvider};

/// How a higher layer combines with a lower one.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum LayerStrategy {
    /// The higher layer replaces every field it defines, endpoints included.
    ///
    /// The right choice for a local override file: when you correct a WHOIS host
    /// because the old one stopped answering, you do not want the old one kept as
    /// a fallback.
    #[default]
    Override,

    /// Endpoints from every layer are kept, lower layers first; all other fields
    /// follow [`Override`](LayerStrategy::Override) rules.
    ///
    /// The right choice for stacking sources that each know about a different
    /// protocol — curated WHOIS hosts under IANA's RDAP bootstrap, say — where
    /// dropping one source's endpoints would lose real coverage.
    Union,
}

/// Several providers consulted as one, later layers taking precedence.
///
/// ```
/// use monovm_whois::registry::{
///     BootstrapRegistry, JsonRegistry, LayerStrategy, LayeredRegistry, RegistryProvider,
/// };
/// use monovm_whois::Tld;
///
/// // Curated WHOIS hosts and markers, plus IANA's RDAP endpoints on top.
/// let registry = LayeredRegistry::new(LayerStrategy::Union)
///     .layer(JsonRegistry::bundled())
///     .layer(BootstrapRegistry::bundled());
///
/// let com = registry.get(&Tld::parse("com").unwrap()).unwrap();
/// assert!(com.endpoints().iter().any(|e| e.is_whois()));
/// assert!(com.endpoints().iter().any(|e| e.is_rdap()));
/// ```
#[derive(Debug, Clone, Default)]
pub struct LayeredRegistry {
    layers: Vec<Arc<dyn RegistryProvider>>,
    strategy: LayerStrategy,
}

impl LayeredRegistry {
    /// An empty stack. Add layers with [`layer`](LayeredRegistry::layer).
    pub fn new(strategy: LayerStrategy) -> Self {
        LayeredRegistry {
            layers: Vec::new(),
            strategy,
        }
    }

    /// Add a layer on top of the existing ones.
    pub fn layer(mut self, provider: impl RegistryProvider + 'static) -> Self {
        self.layers.push(Arc::new(provider));
        self
    }

    /// Add an already-shared layer on top of the existing ones.
    pub fn shared_layer(mut self, provider: Arc<dyn RegistryProvider>) -> Self {
        self.layers.push(provider);
        self
    }

    /// How layers are combined.
    pub fn strategy(&self) -> LayerStrategy {
        self.strategy
    }

    /// How many layers are stacked.
    pub fn layer_count(&self) -> usize {
        self.layers.len()
    }
}

impl RegistryProvider for LayeredRegistry {
    fn get(&self, tld: &Tld) -> Option<Arc<Registry>> {
        // Walk bottom-up so each higher layer is applied over what is already
        // there; that way "higher wins" falls out of the fold rather than needing
        // a separate precedence pass.
        let mut merged: Option<Arc<Registry>> = None;

        for layer in &self.layers {
            let Some(found) = layer.get(tld) else {
                continue;
            };
            merged = Some(match merged {
                None => found,
                Some(base) => Arc::new(match self.strategy {
                    LayerStrategy::Override => base.overlay(&found),
                    LayerStrategy::Union => base.union(&found),
                }),
            });
        }

        merged
    }

    fn tlds(&self) -> Vec<Tld> {
        let mut tlds: Vec<Tld> = self.layers.iter().flat_map(|layer| layer.tlds()).collect();
        tlds.sort();
        tlds.dedup();
        tlds
    }

    fn describe(&self) -> String {
        let layers: Vec<String> = self.layers.iter().map(|layer| layer.describe()).collect();
        format!("layered[{:?}]: {}", self.strategy, layers.join(" < "))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::registry::{Endpoint, JsonRegistry};

    fn provider(json: &str) -> JsonRegistry {
        JsonRegistry::from_json(json, "test").unwrap()
    }

    const BASE: &str = r#"{"registries":[{
        "tlds":["example"],
        "whois":"old.example",
        "available":["No match"],
        "note":"bundled"
    }]}"#;

    #[test]
    fn override_replaces_endpoints() {
        let registry = LayeredRegistry::new(LayerStrategy::Override)
            .layer(provider(BASE))
            .layer(provider(
                r#"{"registries":[{"tlds":["example"],"whois":"new.example"}]}"#,
            ));

        let entry = registry.get(&Tld::parse("example").unwrap()).unwrap();
        assert_eq!(entry.endpoints().len(), 1);
        assert_eq!(entry.endpoints()[0].address(), "new.example");
        // Fields the patch said nothing about survive.
        assert_eq!(entry.available_markers(), ["No match"]);
        assert_eq!(entry.note(), Some("bundled"));
    }

    #[test]
    fn union_keeps_both_sets_of_endpoints_lower_first() {
        let registry = LayeredRegistry::new(LayerStrategy::Union)
            .layer(provider(BASE))
            .layer(provider(
                r#"{"registries":[{"tlds":["example"],"rdap":["https://rdap.example/"]}]}"#,
            ));

        let entry = registry.get(&Tld::parse("example").unwrap()).unwrap();
        assert_eq!(entry.endpoints().len(), 2);
        assert_eq!(entry.endpoints()[0], Endpoint::whois("old.example"));
        assert!(entry.endpoints()[1].is_rdap());
    }

    #[test]
    fn union_does_not_duplicate_a_shared_endpoint() {
        let registry = LayeredRegistry::new(LayerStrategy::Union)
            .layer(provider(BASE))
            .layer(provider(BASE));

        let entry = registry.get(&Tld::parse("example").unwrap()).unwrap();
        assert_eq!(entry.endpoints().len(), 1);
    }

    #[test]
    fn a_tld_only_the_upper_layer_knows_is_still_found() {
        let registry = LayeredRegistry::new(LayerStrategy::Override)
            .layer(provider(BASE))
            .layer(provider(
                r#"{"registries":[{"tlds":["other"],"whois":"w.other"}]}"#,
            ));

        assert!(registry.get(&Tld::parse("other").unwrap()).is_some());
        assert!(registry.get(&Tld::parse("example").unwrap()).is_some());
        assert!(registry.get(&Tld::parse("absent").unwrap()).is_none());
    }

    #[test]
    fn tlds_are_the_deduplicated_union() {
        let registry = LayeredRegistry::new(LayerStrategy::Override)
            .layer(provider(BASE))
            .layer(provider(
                r#"{"registries":[{"tlds":["example","other"],"whois":"w.other"}]}"#,
            ));

        let listed: Vec<String> = registry
            .tlds()
            .iter()
            .map(|t| t.ascii().to_string())
            .collect();
        assert_eq!(listed, ["example", "other"]);
    }

    #[test]
    fn an_empty_stack_resolves_nothing() {
        let registry = LayeredRegistry::new(LayerStrategy::Override);
        assert!(registry.get(&Tld::parse("com").unwrap()).is_none());
        assert!(registry.tlds().is_empty());
    }
}