D53/src/dns_config.rs

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::sync::Arc;
use std::time::Duration;

use tokio::{select, sync::watch};
use tracing::*;

use df_consul::*;

const IPV4_TARGET_METADATA_TAG: &str = "public_ipv4";
const IPV6_TARGET_METADATA_TAG: &str = "public_ipv6";
const CNAME_TARGET_METADATA_TAG: &str = "cname_target";

// ---- Extract DNS config from Consul catalog ----

#[derive(Debug, Eq, PartialEq, Default)]
pub struct DnsConfig {
    pub entries: HashMap<DnsEntryKey, DnsEntryValue>,
}

#[derive(Clone, Debug, Hash, PartialEq, Eq)]
pub struct DnsEntryKey {
    pub dns_path: String,
    pub record_type: DnsRecordType,
}

#[derive(Debug, PartialEq, Eq)]
pub struct DnsEntryValue {
    pub targets: HashSet<String>,
}

#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq)]
#[allow(clippy::upper_case_acronyms)]
pub enum DnsRecordType {
    A,
    AAAA,
    CNAME,
}

impl DnsConfig {
    pub fn new() -> Self {
        Self {
            entries: HashMap::new(),
        }
    }

    fn add(&mut self, k: DnsEntryKey, v: DnsEntryValue) {
        if let Some(ent) = self.entries.get_mut(&k) {
            ent.targets.extend(v.targets);
        } else {
            self.entries.insert(k, v);
        }
    }
}

fn parse_d53_tag(tag: &str, node: &catalog::Node) -> Option<(DnsEntryKey, DnsEntryValue)> {
    let splits = tag.split(' ').collect::<Vec<_>>();
    if splits.len() != 2 {
        return None;
    }

    let (record_type, targets) = match splits[0] {
        "d53-a" => match node.meta.get(IPV4_TARGET_METADATA_TAG) {
            Some(tgt) => (DnsRecordType::A, [tgt.to_string()].into_iter().collect()),
            None => {
                warn!("Got d53-a tag `{}` but node {} does not have a {} metadata value. Tag is ignored.", tag, node.node, IPV4_TARGET_METADATA_TAG);
                return None;
            }
        },
        "d53-aaaa" => match node.meta.get(IPV6_TARGET_METADATA_TAG) {
            Some(tgt) => (DnsRecordType::AAAA, [tgt.to_string()].into_iter().collect()),
            None => {
                warn!("Got d53-aaaa tag `{}` but node {} does not have a {} metadata value. Tag is ignored.", tag, node.node, IPV6_TARGET_METADATA_TAG);
                return None;
            }
        },
        "d53-cname" => match node.meta.get(CNAME_TARGET_METADATA_TAG) {
            Some(tgt) => (
                DnsRecordType::CNAME,
                [tgt.to_string()].into_iter().collect(),
            ),
            None => {
                warn!("Got d53-cname tag `{}` but node {} does not have a {} metadata value. Tag is ignored.", tag, node.node, CNAME_TARGET_METADATA_TAG);
                return None;
            }
        },
        _ => return None,
    };

    Some((
        DnsEntryKey {
            dns_path: splits[1].to_string(),
            record_type,
        },
        DnsEntryValue { targets },
    ))
}

pub fn spawn_dns_config_task(
    consul: &Consul,
    mut must_exit: watch::Receiver<bool>,
) -> watch::Receiver<Arc<DnsConfig>> {
    let (tx, rx) = watch::channel(Arc::new(DnsConfig::new()));

    let mut catalog_rx = consul.watch_all_service_health(Duration::from_secs(60));

    tokio::spawn(async move {
        while !*must_exit.borrow() {
            select! {
                _ = catalog_rx.changed() => (),
                _ = must_exit.changed() => continue,
            };

            let services = catalog_rx.borrow_and_update();

            let mut dns_config = DnsConfig::new();
            for (_svc, nodes) in services.iter() {
                for node in nodes.iter() {
                    // Do not take into account backends if any have status critical
                    if node.checks.iter().any(|x| x.status == "critical") {
                        continue;
                    }
                    for tag in node.service.tags.iter() {
                        if let Some((k, v)) = parse_d53_tag(tag, &node.node) {
                            dns_config.add(k, v);
                        }
                    }
                }
            }

            tx.send(Arc::new(dns_config)).expect("Internal error");
        }
    });

    rx
}

// ---- Display impls ----

impl std::fmt::Display for DnsRecordType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            DnsRecordType::A => write!(f, "A"),
            DnsRecordType::AAAA => write!(f, "AAAA"),
            DnsRecordType::CNAME => write!(f, "CNAME"),
        }
    }
}

impl std::fmt::Display for DnsEntryKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{} IN {}", self.dns_path, self.record_type)
    }
}

impl std::fmt::Display for DnsEntryValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "[")?;
        for (i, tgt) in self.targets.iter().enumerate() {
            if i > 0 {
                write!(f, " ")?;
            }
            write!(f, "{}", tgt)?;
        }
        write!(f, "]")
    }
}