path: root/src/dns_config.rs

use std::collections::{HashMap, HashSet};
use std::fmt;
use std::net::{Ipv4Addr, Ipv6Addr};
use std::sync::Arc;
use std::time::{Duration, SystemTime};

use anyhow::Result;
use serde::{Deserialize, Serialize};
use tokio::{select, sync::watch};
use tracing::*;

use df_consul::*;

const IP_TARGET_METADATA_TAG_PREFIX: &str = "public_";
const CNAME_TARGET_METADATA_TAG: &str = "cname_target";

const AUTODISCOVERY_CACHE_DURATION: u64 = 600; // 10 minutes

// ---- 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);
        }
    }
}

// ---- fetcher and autodiscovery cache ----

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

    let fetcher = DnsConfigFetcher {
        consul,
        node_ipv4_cache: HashMap::new(),
        node_ipv6_cache: HashMap::new(),
    };

    tokio::spawn(fetcher.task(tx, must_exit));

    rx
}

struct DnsConfigFetcher {
    consul: Consul,
    node_ipv4_cache: HashMap<String, (u64, Option<Ipv4Addr>)>,
    node_ipv6_cache: HashMap<String, (u64, Option<Ipv6Addr>)>,
}

impl DnsConfigFetcher {
    async fn task(
        mut self,
        tx: watch::Sender<Arc<DnsConfig>>,
        mut must_exit: watch::Receiver<bool>,
    ) {
        let mut catalog_rx = self
            .consul
            .watch_all_service_health(Duration::from_secs(60));

        while !*must_exit.borrow() {
            select! {
                _ = catalog_rx.changed() => (),
                _ = must_exit.changed() => continue,
            };

            let services = catalog_rx.borrow_and_update().clone();
            match self.parse_catalog(&services).await {
                Ok(dns_config) => tx.send(Arc::new(dns_config)).expect("Internal error"),
                Err(e) => {
                    error!("Error when parsing tags: {}", e);
                }
            };
        }
    }

    async fn parse_catalog(&mut self, services: &catalog::AllServiceHealth) -> Result<DnsConfig> {
        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)) = self.parse_d53_tag(tag, &node.node).await? {
                        dns_config.add(k, v);
                    }
                }
            }
        }

        Ok(dns_config)
    }

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

        let (record_type, targets) = match splits[0] {
            "d53-a" => match self.get_node_ipv4(&node).await? {
                Some(tgt) => (DnsRecordType::A, [tgt.to_string()].into_iter().collect()),
                None => {
                    warn!("Got d53-a tag `{}` but node {} does not appear to have a known public IPv4 address. Tag is ignored.", tag, node.node);
                    return Ok(None);
                }
            },
            "d53-aaaa" => match self.get_node_ipv6(&node).await? {
                Some(tgt) => (DnsRecordType::AAAA, [tgt.to_string()].into_iter().collect()),
                None => {
                    warn!("Got d53-aaaa tag `{}` but node {} does not appear to have a known public IPv6 address. Tag is ignored.", tag, node.node);
                    return Ok(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 Ok(None);
                }
            },
            _ => return Ok(None),
        };

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

    async fn get_node_ipv4(&mut self, node: &catalog::Node) -> Result<Option<Ipv4Addr>> {
        Self::get_node_ip(&self.consul, "ipv4", &mut self.node_ipv4_cache, node).await
    }

    async fn get_node_ipv6(&mut self, node: &catalog::Node) -> Result<Option<Ipv6Addr>> {
        Self::get_node_ip(&self.consul, "ipv6", &mut self.node_ipv6_cache, node).await
    }

    async fn get_node_ip<A>(
        consul: &Consul,
        family: &str,
        cache: &mut HashMap<String, (u64, Option<A>)>,
        node: &catalog::Node,
    ) -> Result<Option<A>>
    where
        A: Serialize + for<'de> Deserialize<'de> + std::fmt::Debug + std::str::FromStr + Copy + Eq,
        <A as std::str::FromStr>::Err: Send + Sync + std::error::Error + 'static,
    {
        match cache.get(&node.node) {
            Some((t, a)) if timestamp() <= t + AUTODISCOVERY_CACHE_DURATION => Ok(*a),
            _ => {
                let kv_key = format!("diplonat/autodiscovery/{}/{}", family, node.node);
                let autodiscovery = consul.kv_get(&kv_key).await?;

                if let Some(json) = autodiscovery {
                    let a = serde_json::from_slice::<DiplonatAutodiscoveryResult<A>>(&json)?;
                    if timestamp() <= a.timestamp + AUTODISCOVERY_CACHE_DURATION {
                        if cache.get(&node.node).map(|x| x.1) != Some(a.address) {
                            info!(
                                "Got {} address for {} from diplonat autodiscovery: {:?}",
                                family, node.node, a.address
                            );
                        }
                        cache.insert(node.node.clone(), (a.timestamp, a.address));
                        return Ok(a.address);
                    } else {
                        warn!("{} address for {} from diplonat autodiscovery is outdated (value: {:?}), falling back on value from Consul node meta", family, node.node, a.address);
                    }
                }

                let meta_tag = format!("{}{}", IP_TARGET_METADATA_TAG_PREFIX, family);
                let a = node.meta.get(&meta_tag).map(|x| x.parse()).transpose()?;

                if cache.get(&node.node).map(|x| x.1) != Some(a) {
                    info!(
                        "Got {} address for {} from Consul node meta: {:?}",
                        family, node.node, a
                    );
                }
                cache.insert(node.node.clone(), (timestamp(), a));
                Ok(a)
            }
        }
    }
}

// ---- util for interaction with diplonat ----

#[derive(Serialize, Deserialize, Debug)]
pub struct DiplonatAutodiscoveryResult<A> {
    pub timestamp: u64,
    pub address: Option<A>,
}

fn timestamp() -> u64 {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .expect("clock error")
        .as_secs()
}

// ---- 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, "]")
    }
}