2022-12-07 14:35:12 +00:00
|
|
|
use std::collections::{HashMap, HashSet};
|
|
|
|
use std::fmt;
|
|
|
|
use std::sync::Arc;
|
|
|
|
use std::{cmp, time::Duration};
|
|
|
|
|
|
|
|
use anyhow::Result;
|
|
|
|
|
|
|
|
use futures::future::BoxFuture;
|
|
|
|
use futures::stream::{FuturesUnordered, StreamExt};
|
|
|
|
|
|
|
|
use tokio::{select, sync::watch, time::sleep};
|
2022-12-11 14:46:48 +00:00
|
|
|
use tracing::*;
|
2022-12-07 14:35:12 +00:00
|
|
|
|
|
|
|
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)]
|
|
|
|
pub struct DnsConfig {
|
|
|
|
pub entries: HashMap<DnsEntryKey, DnsEntryValue>,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug, Hash, PartialEq, Eq)]
|
|
|
|
pub struct DnsEntryKey {
|
2022-12-11 15:29:06 +00:00
|
|
|
pub dns_path: String,
|
2022-12-07 14:35:12 +00:00
|
|
|
pub record_type: DnsRecordType,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug, PartialEq, Eq)]
|
|
|
|
pub struct DnsEntryValue {
|
|
|
|
pub targets: HashSet<String>,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(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: &ConsulNode) -> Option<(DnsEntryKey, DnsEntryValue)> {
|
|
|
|
let splits = tag.split(' ').collect::<Vec<_>>();
|
2022-12-11 15:29:06 +00:00
|
|
|
if splits.len() != 2 {
|
2022-12-07 14:35:12 +00:00
|
|
|
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 {
|
2022-12-11 15:29:06 +00:00
|
|
|
dns_path: splits[1].to_string(),
|
2022-12-07 14:35:12 +00:00
|
|
|
record_type,
|
|
|
|
},
|
|
|
|
DnsEntryValue { targets },
|
|
|
|
))
|
|
|
|
}
|
|
|
|
|
|
|
|
fn parse_consul_catalog(catalog: &ConsulNodeCatalog, dns_config: &mut DnsConfig) {
|
|
|
|
trace!("Parsing node catalog: {:#?}", catalog);
|
|
|
|
|
|
|
|
for (_, svc) in catalog.services.iter() {
|
|
|
|
for tag in svc.tags.iter() {
|
|
|
|
if let Some((k, v)) = parse_d53_tag(tag, &catalog.node) {
|
|
|
|
dns_config.add(k, v);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Default)]
|
|
|
|
struct NodeWatchState {
|
|
|
|
last_idx: Option<usize>,
|
|
|
|
last_catalog: Option<ConsulNodeCatalog>,
|
|
|
|
retries: u32,
|
|
|
|
}
|
|
|
|
|
|
|
|
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 consul = Arc::new(consul);
|
|
|
|
|
|
|
|
tokio::spawn(async move {
|
|
|
|
let mut nodes = HashMap::new();
|
|
|
|
let mut watches = FuturesUnordered::<BoxFuture<'static, (String, Result<_>)>>::new();
|
|
|
|
|
|
|
|
let mut node_site = HashMap::new();
|
|
|
|
|
|
|
|
while !*must_exit.borrow() {
|
|
|
|
let list_nodes = select! {
|
|
|
|
ln = consul.list_nodes() => ln,
|
|
|
|
_ = must_exit.changed() => continue,
|
|
|
|
};
|
|
|
|
|
|
|
|
match list_nodes {
|
|
|
|
Ok(consul_nodes) => {
|
|
|
|
debug!("Watched consul nodes: {:?}", consul_nodes);
|
|
|
|
for consul_node in consul_nodes {
|
|
|
|
let node = &consul_node.node;
|
|
|
|
if !nodes.contains_key(node) {
|
|
|
|
nodes.insert(node.clone(), NodeWatchState::default());
|
|
|
|
|
|
|
|
let node = node.to_string();
|
|
|
|
let consul = consul.clone();
|
|
|
|
|
|
|
|
watches.push(Box::pin(async move {
|
|
|
|
let res = consul.watch_node(&node, None).await;
|
|
|
|
(node, res)
|
|
|
|
}));
|
|
|
|
}
|
|
|
|
if let Some(site) = consul_node.meta.get("site") {
|
|
|
|
node_site.insert(node.clone(), site.clone());
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|
|
|
|
Err(e) => {
|
|
|
|
error!("Could not get Consul node list: {}", e);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
let next_watch = select! {
|
|
|
|
nw = watches.next() => nw,
|
|
|
|
_ = must_exit.changed() => continue,
|
|
|
|
};
|
|
|
|
|
|
|
|
let (node, res): (String, Result<_>) = match next_watch {
|
|
|
|
Some(v) => v,
|
|
|
|
None => {
|
|
|
|
warn!("No nodes currently watched in dns_config.rs");
|
|
|
|
sleep(Duration::from_secs(10)).await;
|
|
|
|
continue;
|
|
|
|
}
|
|
|
|
};
|
|
|
|
|
|
|
|
match res {
|
|
|
|
Ok((catalog, new_idx)) => {
|
|
|
|
let mut watch_state = nodes.get_mut(&node).unwrap();
|
|
|
|
watch_state.last_idx = Some(new_idx);
|
|
|
|
watch_state.last_catalog = Some(catalog);
|
|
|
|
watch_state.retries = 0;
|
|
|
|
|
|
|
|
let idx = watch_state.last_idx;
|
|
|
|
let consul = consul.clone();
|
|
|
|
watches.push(Box::pin(async move {
|
|
|
|
let res = consul.watch_node(&node, idx).await;
|
|
|
|
(node, res)
|
|
|
|
}));
|
|
|
|
}
|
|
|
|
Err(e) => {
|
|
|
|
let mut watch_state = nodes.get_mut(&node).unwrap();
|
|
|
|
watch_state.retries += 1;
|
|
|
|
watch_state.last_idx = None;
|
|
|
|
|
|
|
|
let will_retry_in =
|
|
|
|
retry_to_time(watch_state.retries, Duration::from_secs(600));
|
|
|
|
error!(
|
|
|
|
"Failed to query consul for node {}. Will retry in {}s. {}",
|
|
|
|
node,
|
|
|
|
will_retry_in.as_secs(),
|
|
|
|
e
|
|
|
|
);
|
|
|
|
|
|
|
|
let consul = consul.clone();
|
|
|
|
watches.push(Box::pin(async move {
|
|
|
|
sleep(will_retry_in).await;
|
|
|
|
let res = consul.watch_node(&node, None).await;
|
|
|
|
(node, res)
|
|
|
|
}));
|
|
|
|
continue;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
let mut dns_config = DnsConfig::new();
|
|
|
|
for (_, watch_state) in nodes.iter() {
|
|
|
|
if let Some(catalog) = &watch_state.last_catalog {
|
|
|
|
parse_consul_catalog(catalog, &mut dns_config);
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
tx.send(Arc::new(dns_config)).expect("Internal error");
|
|
|
|
}
|
|
|
|
});
|
|
|
|
|
|
|
|
rx
|
|
|
|
}
|
|
|
|
|
|
|
|
fn retry_to_time(retries: u32, max_time: Duration) -> Duration {
|
|
|
|
// 1.2^x seems to be a good value to exponentially increase time at a good pace
|
|
|
|
// eg. 1.2^32 = 341 seconds ~= 5 minutes - ie. after 32 retries we wait 5
|
|
|
|
// minutes
|
|
|
|
Duration::from_secs(cmp::min(
|
|
|
|
max_time.as_secs(),
|
|
|
|
1.2f64.powf(retries as f64) as u64,
|
|
|
|
))
|
|
|
|
}
|
|
|
|
|
|
|
|
// ---- 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 {
|
2022-12-11 15:29:06 +00:00
|
|
|
write!(f, "{} IN {}", self.dns_path, self.record_type)
|
2022-12-07 14:35:12 +00:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
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, "]")
|
|
|
|
}
|
|
|
|
}
|