nazrin/nzrdns/src/dns.rs

use nzr_api::config::DNSConfig;
use std::borrow::Borrow;
use std::collections::{BTreeMap, HashMap};
use std::ops::Deref;
use std::str::FromStr;
use std::sync::Arc;
use tokio::sync::{Mutex, RwLock};

use nzr_api::model::{Instance, SubnetData};

use hickory_proto::rr::Name;
use hickory_server::authority::{AuthorityObject, Catalog};
use hickory_server::proto::rr::{rdata::soa, RData, RecordSet};
use hickory_server::proto::rr::{LowerName, RrKey};
use hickory_server::server::{Request, RequestHandler, ResponseHandler, ResponseInfo};
use hickory_server::{proto::rr::Record, store::in_memory::InMemoryAuthority};

#[derive(Clone)]
pub struct CatalogRef(Arc<RwLock<Catalog>>);

macro_rules! make_serial {
    () => {{
        use std::time::{SystemTime, UNIX_EPOCH};

        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .expect("You know what? Fuck you. *unepochs your time*")
            .as_secs() as u32
    }};
}

// hickory_dns 0.24 still requires async_trait
#[async_trait::async_trait]
impl RequestHandler for CatalogRef {
    async fn handle_request<R: ResponseHandler>(
        &self,
        request: &Request,
        response_handle: R,
    ) -> ResponseInfo {
        self.0
            .read()
            .await
            .handle_request(request, response_handle)
            .await
    }
}

#[derive(Clone)]
pub struct ZoneData(Arc<InnerZD>);

impl Deref for ZoneData {
    type Target = InnerZD;
    fn deref(&self) -> &Self::Target {
        self.0.as_ref()
    }
}

impl ZoneData {
    pub fn new(dns_config: &DNSConfig) -> Self {
        ZoneData(Arc::new(InnerZD::new(dns_config)))
    }
}

pub struct InnerZD {
    default_zone: Arc<InMemoryAuthority>,
    map: Mutex<HashMap<String, Arc<InMemoryAuthority>>>,
    catalog: CatalogRef,
    config: DNSConfig,
}

pub fn make_rectree_with_soa(name: &Name, config: &DNSConfig) -> BTreeMap<RrKey, RecordSet> {
    tracing::debug!("Creating initial SOA for {}", &name);
    let mut records: BTreeMap<RrKey, RecordSet> = BTreeMap::new();
    let soa_key = RrKey::new(
        LowerName::from(name),
        hickory_server::proto::rr::RecordType::SOA,
    );
    let soa_rec = Record::from_rdata(
        name.clone(),
        3600,
        RData::SOA(soa::SOA::new(
            name.clone(),
            config.soa.contact.clone(),
            make_serial!(),
            config.soa.refresh,
            config.soa.retry,
            config.soa.expire,
            3600,
        )),
    );
    records.insert(soa_key, RecordSet::from(soa_rec));
    records
}

impl InnerZD {
    pub fn new(dns_config: &DNSConfig) -> Self {
        let default_zone = Arc::new({
            let records = make_rectree_with_soa(&dns_config.default_zone, dns_config);
            InMemoryAuthority::new(
                dns_config.default_zone.clone(),
                records,
                hickory_server::authority::ZoneType::Primary,
                false,
            )
            .unwrap()
        });
        let mut catalog = Catalog::new();
        catalog.upsert(
            dns_config.default_zone.borrow().into(),
            Box::new(default_zone.clone()),
        );
        Self {
            default_zone,
            map: Mutex::new(HashMap::new()),
            catalog: CatalogRef(Arc::new(RwLock::new(catalog))),
            config: dns_config.clone(),
        }
    }

    /// Creates a new DNS zone for the given subnet.
    pub async fn new_zone(
        &self,
        zone_id: impl AsRef<str>,
        subnet: &SubnetData,
    ) -> Result<(), Box<dyn std::error::Error>> {
        if let Some(name) = &subnet.domain_name {
            let rectree = make_rectree_with_soa(name, &self.config);
            let auth = InMemoryAuthority::new(
                name.clone(),
                rectree,
                hickory_server::authority::ZoneType::Primary,
                false,
            )?;
            self.import(zone_id.as_ref(), auth).await;
        }
        Ok(())
    }

    /// Generates a zone with the given records.
    async fn import(&self, name: &str, auth: InMemoryAuthority) {
        let auth_arc = Arc::new(auth);
        tracing::debug!(
            "Importing {} with {} records...",
            name,
            auth_arc.records().await.len()
        );
        self.map
            .lock()
            .await
            .insert(name.to_owned(), auth_arc.clone());

        self.catalog
            .0
            .write()
            .await
            .upsert(auth_arc.origin().clone(), Box::new(auth_arc.clone()));
    }

    /// Deletes the DNS zone.
    pub async fn delete_zone(&self, domain_name: &str) -> bool {
        self.map.lock().await.remove(domain_name).is_some()
    }

    /// Adds a new host record in the DNS zone.
    pub async fn new_record(&self, inst: &Instance) -> Result<(), Box<dyn std::error::Error>> {
        let hostname = Name::from_str(&inst.name)?;
        let zones = self.map.lock().await;
        let zone = zones.get(&inst.lease.subnet).unwrap_or(&self.default_zone);
        let fqdn = {
            let origin: Name = zone.origin().into();
            hostname.append_domain(&origin)?
        };

        tracing::debug!(
            "Creating new host entry {} in zone {}...",
            &fqdn,
            zone.origin()
        );

        let addr = inst.lease.addr.addr;

        let record = Record::from_rdata(fqdn, 3600, RData::A(addr.into()));
        zone.upsert(record, 0).await;
        self.catalog()
            .0
            .write()
            .await
            .upsert(zone.origin().clone(), Box::new(zone.clone()));
        Ok(())
    }

    pub async fn delete_record(&self, inst: &Instance) -> Result<bool, Box<dyn std::error::Error>> {
        let hostname = Name::from_str(&inst.name)?;
        let mut zones = self.map.lock().await;
        if let Some(zone) = zones.get_mut(&inst.lease.subnet) {
            let hostname: LowerName = hostname.into();
            self.catalog.0.write().await.remove(&hostname);
            let key = RrKey::new(hostname, hickory_server::proto::rr::RecordType::A);
            Ok(zone.records_mut().await.remove(&key).is_some())
        } else {
            Ok(false)
        }
    }

    pub fn catalog(&self) -> CatalogRef {
        self.catalog.clone()
    }
}