1 files changed, 76 insertions, 47 deletions

diff --git a/src/stun_actor.rs b/src/stun_actor.rs
index 6740c83..bf4da73 100644
--- a/src/stun_actor.rs
+++ b/src/stun_actor.rs
@@ -8,12 +8,24 @@ use serde::{Deserialize, Serialize};
 use crate::config::{RuntimeConfigConsul, RuntimeConfigStun};
 use crate::consul;
 
+/// If autodiscovery returns None but an address was obtained less than
+/// this number of seconds ago (here 15 minutes), we keep that address
+/// in the Consul db instead of insterting a None.
+const PERSIST_SOME_RESULT_DURATION_SECS: u64 = 900;
+
 pub struct StunActor {
-    node: String,
     consul: consul::Consul,
-    stun_server_v4: Option<SocketAddr>,
-    stun_server_v6: SocketAddr,
     refresh_time: Duration,
+
+    autodiscovery_v4: StunAutodiscovery,
+    autodiscovery_v6: StunAutodiscovery,
+}
+
+pub struct StunAutodiscovery {
+    consul_key: String,
+    is_ipv4: bool,
+    stun_server: Option<SocketAddr>,
+    last_result: Option<AutodiscoverResult>,
 }
 
 #[derive(Serialize, Deserialize, Debug)]
@@ -34,79 +46,96 @@ impl StunActor {
             .unwrap_or(true));
         assert!(stun_config.stun_server_v6.is_ipv6());
 
+        let autodiscovery_v4 = StunAutodiscovery {
+            consul_key: format!("diplonat/autodiscovery/ipv4/{}", node),
+            is_ipv4: true,
+            stun_server: stun_config.stun_server_v4,
+            last_result: None,
+        };
+
+        let autodiscovery_v6 = StunAutodiscovery {
+            consul_key: format!("diplonat/autodiscovery/ipv6/{}", node),
+            is_ipv4: false,
+            stun_server: Some(stun_config.stun_server_v6),
+            last_result: None,
+        };
+
         Self {
             consul: consul::Consul::new(consul_config),
-            node: node.to_string(),
-            stun_server_v4: stun_config.stun_server_v4,
-            stun_server_v6: stun_config.stun_server_v6,
+            autodiscovery_v4,
+            autodiscovery_v6,
             refresh_time: stun_config.refresh_time,
         }
     }
 
     pub async fn listen(&mut self) -> Result<()> {
         loop {
-            let ipv4_result = match self.stun_server_v4 {
-                Some(stun_server_v4) => self.autodiscover_ip(stun_server_v4).await,
-                None => self.autodiscover_none_ipv4().await,
-            };
-            if let Err(e) = ipv4_result {
+            if let Err(e) = self.autodiscovery_v4.do_iteration(&self.consul).await {
                 error!("Unable to autodiscover IPv4 address: {}", e);
             }
-            if let Err(e) = self.autodiscover_ip(self.stun_server_v6).await {
+
+            if let Err(e) = self.autodiscovery_v6.do_iteration(&self.consul).await {
                 error!("Unable to autodiscover IPv6 address: {}", e);
             }
             tokio::time::sleep(self.refresh_time).await;
         }
     }
+}
 
-    async fn autodiscover_ip(&self, stun_server: SocketAddr) -> Result<()> {
-        let binding_ip = match stun_server.is_ipv4() {
+impl StunAutodiscovery {
+    async fn do_iteration(&mut self, consul: &consul::Consul) -> Result<()> {
+        let binding_ip = match self.is_ipv4 {
             true => IpAddr::V4(Ipv4Addr::UNSPECIFIED),  // 0.0.0.0
             false => IpAddr::V6(Ipv6Addr::UNSPECIFIED), // [::]
         };
         let binding_addr = SocketAddr::new(binding_ip, 0);
 
-        let discovered_addr = get_mapped_addr(stun_server, binding_addr)
-            .await?
-            .map(|x| x.ip());
+        let discovered_addr = match self.stun_server {
+            Some(stun_server) => {
+                assert_eq!(self.is_ipv4, stun_server.is_ipv4());
 
-        let consul_key = match stun_server.is_ipv4() {
-            true => {
-                debug!("Autodiscovered IPv4: {:?}", discovered_addr);
-                format!("diplonat/autodiscovery/ipv4/{}", self.node)
-            }
-            false => {
-                debug!("Autodiscovered IPv6: {:?}", discovered_addr);
-                format!("diplonat/autodiscovery/ipv6/{}", self.node)
+                get_mapped_addr(stun_server, binding_addr)
+                    .await?
+                    .map(|x| x.ip())
             }
+            None => None,
         };
 
-        self.consul
-            .kv_put(
-                &consul_key,
-                serde_json::to_vec(&AutodiscoverResult {
-                    timestamp: timestamp(),
-                    address: discovered_addr,
-                })?,
-            )
-            .await?;
+        let now = timestamp();
+
+        if discovered_addr.is_none() {
+            if let Some(last_result) = &self.last_result {
+                if last_result.address.is_some()
+                    && now - last_result.timestamp <= PERSIST_SOME_RESULT_DURATION_SECS
+                {
+                    // Keep non-None result that was obtained before by not
+                    // writing/taking into account None result.
+                    return Ok(());
+                }
+            }
+        }
 
-        Ok(())
-    }
+        let current_result = AutodiscoverResult {
+            timestamp: now,
+            address: discovered_addr,
+        };
 
-    async fn autodiscover_none_ipv4(&self) -> Result<()> {
-        let consul_key = format!("diplonat/autodiscovery/ipv4/{}", self.node);
-
-        self.consul
-            .kv_put(
-                &consul_key,
-                serde_json::to_vec(&AutodiscoverResult {
-                    timestamp: timestamp(),
-                    address: None,
-                })?,
-            )
+        let msg = format!(
+            "STUN autodiscovery result: {} -> {:?}",
+            self.consul_key, discovered_addr
+        );
+        if self.last_result.as_ref().and_then(|x| x.address) != discovered_addr {
+            info!("{}", msg);
+        } else {
+            debug!("{}", msg);
+        }
+
+        consul
+            .kv_put(&self.consul_key, serde_json::to_vec(&current_result)?)
             .await?;
 
+        self.last_result = Some(current_result);
+
         Ok(())
     }
 }