Support for loading IP addresses from DiploNAT autodiscovery

6 files changed, 184 insertions, 69 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 541ab6d..addbd5c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -106,6 +106,7 @@ dependencies = [
  "futures",
  "reqwest",
  "serde",
+ "serde_json",
  "structopt",
  "tokio",
  "tracing",
diff --git a/Cargo.nix b/Cargo.nix
index e153bb7..4b369b1 100644
--- a/Cargo.nix
+++ b/Cargo.nix
@@ -23,7 +23,7 @@ args@{
   ignoreLockHash,
 }:
 let
-  nixifiedLockHash = "66706604b5c3f270cb3a6b72cb23bf3abba94b9d03fed4048584d92e89851d54";
+  nixifiedLockHash = "f8fd87706eb1709f2cf3a695f4400e1a5e130e3c599cdfebc00bba51c80f494f";
   workspaceSrc = if args.workspaceSrc == null then ./. else args.workspaceSrc;
   currentLockHash = builtins.hashFile "sha256" (workspaceSrc + /Cargo.lock);
   lockHashIgnored = if ignoreLockHash
@@ -188,10 +188,11 @@ in
     dependencies = {
       anyhow = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".anyhow."1.0.66" { inherit profileName; }).out;
       async_trait = (buildRustPackages."registry+https://github.com/rust-lang/crates.io-index".async-trait."0.1.59" { profileName = "__noProfile"; }).out;
-      df_consul = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".df-consul."0.2.0" { inherit profileName; }).out;
+      df_consul = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".df-consul."0.3.3" { inherit profileName; }).out;
       futures = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".futures."0.3.25" { inherit profileName; }).out;
       reqwest = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".reqwest."0.11.13" { inherit profileName; }).out;
       serde = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".serde."1.0.149" { inherit profileName; }).out;
+      serde_json = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".serde_json."1.0.89" { inherit profileName; }).out;
       structopt = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".structopt."0.3.26" { inherit profileName; }).out;
       tokio = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".tokio."1.23.0" { inherit profileName; }).out;
       tracing = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".tracing."0.1.37" { inherit profileName; }).out;
@@ -199,17 +200,19 @@ in
     };
   });
   
-  "registry+https://github.com/rust-lang/crates.io-index".df-consul."0.2.0" = overridableMkRustCrate (profileName: rec {
+  "registry+https://github.com/rust-lang/crates.io-index".df-consul."0.3.3" = overridableMkRustCrate (profileName: rec {
     name = "df-consul";
-    version = "0.2.0";
+    version = "0.3.3";
     registry = "registry+https://github.com/rust-lang/crates.io-index";
-    src = fetchCratesIo { inherit name version; sha256 = "a351d00f138e768845cdefb9ae27b79aeed97c698745c73bb2805cad1167aa81"; };
+    src = fetchCratesIo { inherit name version; sha256 = "0e38cfbab431b53dfd2d09f2a9902510c636d3d7397645bac5cf1959cfde2999"; };
     dependencies = {
       anyhow = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".anyhow."1.0.66" { inherit profileName; }).out;
       bytes = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".bytes."1.3.0" { inherit profileName; }).out;
+      futures = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".futures."0.3.25" { inherit profileName; }).out;
       log = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".log."0.4.17" { inherit profileName; }).out;
       reqwest = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".reqwest."0.11.13" { inherit profileName; }).out;
       serde = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".serde."1.0.149" { inherit profileName; }).out;
+      tokio = (rustPackages."registry+https://github.com/rust-lang/crates.io-index".tokio."1.23.0" { inherit profileName; }).out;
     };
   });
   
diff --git a/Cargo.toml b/Cargo.toml
index 6af9ec2..0b26b63 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,5 +16,6 @@ structopt = "0.3"
 tokio = { version = "1.22", default-features = false, features = ["rt", "rt-multi-thread", "io-util", "net", "time", "macros", "sync", "signal", "fs"] }
 reqwest = { version = "0.11", default-features = false, features = ["json", "rustls-tls-webpki-roots" ] }
 serde = { version = "1.0.107", features = ["derive"] }
+serde_json = "1.0"
 
 df-consul = "0.3.3"
diff --git a/run_local.sh b/run_local.sh
index 150709a..ea4f534 100755
--- a/run_local.sh
+++ b/run_local.sh
@@ -1,6 +1,6 @@
 #!/bin/sh
 
-RUST_LOG=d53=info cargo run \
+RUST_LOG=d53=debug cargo run \
 	-- \
 	--consul-addr http://localhost:8500 \
 	--providers deuxfleurs.org:gandi \
diff --git a/src/dns_config.rs b/src/dns_config.rs
index f4a95be..9eb6075 100644
--- a/src/dns_config.rs
+++ b/src/dns_config.rs
@@ -1,17 +1,21 @@
 use std::collections::{HashMap, HashSet};
 use std::fmt;
+use std::net::{Ipv4Addr, Ipv6Addr};
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, SystemTime};
 
+use anyhow::Result;
+use serde::{Deserialize, Serialize};
 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 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)]
@@ -54,86 +58,192 @@ impl DnsConfig {
     }
 }
 
-fn parse_d53_tag(tag: &str, node: &catalog::Node) -> Option<(DnsEntryKey, DnsEntryValue)> {
-    let splits = tag.split(' ').collect::<Vec<_>>();
-    if splits.len() != 2 {
-        return None;
-    }
+// ---- fetcher and autodiscovery cache ----
 
-    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,
+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(),
     };
 
-    Some((
-        DnsEntryKey {
-            dns_path: splits[1].to_string(),
-            record_type,
-        },
-        DnsEntryValue { targets },
-    ))
+    tokio::spawn(fetcher.task(tx, must_exit));
+
+    rx
 }
 
-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()));
+struct DnsConfigFetcher {
+    consul: Consul,
+    node_ipv4_cache: HashMap<String, (u64, Option<Ipv4Addr>)>,
+    node_ipv6_cache: HashMap<String, (u64, Option<Ipv6Addr>)>,
+}
 
-    let mut catalog_rx = consul.watch_all_service_health(Duration::from_secs(60));
+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));
 
-    tokio::spawn(async move {
         while !*must_exit.borrow() {
             select! {
                 _ = catalog_rx.changed() => (),
                 _ = must_exit.changed() => continue,
             };
 
-            let services = catalog_rx.borrow_and_update();
+            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);
+                }
+            };
+        }
+    }
 
-            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;
+    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);
                     }
-                    for tag in node.service.tags.iter() {
-                        if let Some((k, v)) = parse_d53_tag(tag, &node.node) {
-                            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);
                     }
                 }
-            }
 
-            tx.send(Arc::new(dns_config)).expect("Internal error");
+                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)
+            }
         }
-    });
+    }
+}
 
-    rx
+// ---- 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 ----
diff --git a/src/main.rs b/src/main.rs
index 1a41c18..a8a6b36 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -111,7 +111,7 @@ async fn main() {
         .map(ToString::to_string)
         .collect::<Vec<_>>();
 
-    let rx_dns_config = dns_config::spawn_dns_config_task(&consul, exit_signal.clone());
+    let rx_dns_config = dns_config::spawn_dns_config_task(consul, exit_signal.clone());
 
     let updater_task = tokio::spawn(dns_updater::dns_updater_task(
         rx_dns_config.clone(),