use std::collections::{HashMap, HashSet}; use std::fmt; use std::net::{Ipv4Addr, Ipv6Addr}; use std::sync::Arc; use std::time::Duration; use anyhow::Result; use serde::{Deserialize, Serialize}; use tokio::{select, sync::watch}; use tracing::*; use df_consul::*; use crate::autodiscovery::*; 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, } #[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, } #[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, ) -> watch::Receiver> { let (tx, rx) = watch::channel(Arc::new(DnsConfig::new())); let fetcher = DnsConfigTask { consul }; tokio::spawn(fetcher.task(tx, must_exit)); rx } struct DnsConfigTask { consul: Consul, } impl DnsConfigTask { async fn task( mut self, tx: watch::Sender>, mut must_exit: watch::Receiver, ) { let mut autodiscovery_rx = watch_autodiscovered_ips(self.consul.clone(), must_exit.clone()); let mut catalog_rx = self .consul .watch_all_service_health(Duration::from_secs(60)); while !*must_exit.borrow() { select! { _ = catalog_rx.changed() => (), _ = autodiscovery_rx.changed() => (), _ = must_exit.changed() => continue, }; let services = catalog_rx.borrow_and_update().clone(); let autodiscovery = autodiscovery_rx.borrow_and_update().clone(); match self.parse_catalog(&services, &autodiscovery) { Ok(dns_config) => tx.send(Arc::new(dns_config)).expect("Internal error"), Err(e) => { error!("Error when parsing tags: {}", e); } }; } } fn parse_catalog( &mut self, services: &catalog::AllServiceHealth, autodiscovery: &AutodiscoveredAddresses, ) -> Result { 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, autodiscovery)? { dns_config.add(k, v); } } } } Ok(dns_config) } fn parse_d53_tag( &mut self, tag: &str, node: &catalog::Node, autodiscovery: &AutodiscoveredAddresses, ) -> Result> { let splits = tag.split(' ').collect::>(); if splits.len() != 2 { return Ok(None); } let (record_type, target) = match splits[0] { "d53-a" => match self.get_node_ipv4(&autodiscovery, &node)? { Some(tgt) => (DnsRecordType::A, tgt.to_string()), 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(&autodiscovery, &node)? { Some(tgt) => (DnsRecordType::AAAA, tgt.to_string()), 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()), 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: [target].into_iter().collect(), }, ))) } fn get_node_ipv4( &mut self, autodiscovery: &AutodiscoveredAddresses, node: &catalog::Node, ) -> Result> { Self::get_node_ip("ipv4", &autodiscovery.ipv4, node) } fn get_node_ipv6( &mut self, autodiscovery: &AutodiscoveredAddresses, node: &catalog::Node, ) -> Result> { Self::get_node_ip("ipv6", &autodiscovery.ipv6, node) } fn get_node_ip( family: &'static str, autodiscovery: &HashMap>, node: &catalog::Node, ) -> Result> where A: Serialize + for<'de> Deserialize<'de> + std::fmt::Debug + std::str::FromStr + Copy + Eq, ::Err: Send + Sync + std::error::Error + 'static, { match autodiscovery.get(&node.node) { Some(ar) if timestamp() <= ar.timestamp + AUTODISCOVERY_CACHE_DURATION => { Ok(ar.address) } x => { if let Some(ar) = x { warn!("{} address for {} from diplonat autodiscovery is outdated (value: {:?}), falling back on value from Consul node meta", family, node.node, ar.address); } let meta_tag = format!("{}{}", IP_TARGET_METADATA_TAG_PREFIX, family); let addr = node.meta.get(&meta_tag).map(|x| x.parse()).transpose()?; Ok(addr) } } } } // ---- 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, "]") } }