From 097a2a029d8e06355ff24a5dcdfb2fcabfca839c Mon Sep 17 00:00:00 2001
From: Alex Auvolat <alex@adnab.me>
Date: Tue, 16 Jan 2024 15:24:34 +0100
Subject: [PATCH] refactor stun_actor

---
 src/stun_actor.rs | 85 +++++++++++++++++++++++------------------------
 1 file changed, 42 insertions(+), 43 deletions(-)

diff --git a/src/stun_actor.rs b/src/stun_actor.rs
index 6740c83..3182662 100644
--- a/src/stun_actor.rs
+++ b/src/stun_actor.rs
@@ -9,11 +9,18 @@ use crate::config::{RuntimeConfigConsul, RuntimeConfigStun};
 use crate::consul;
 
 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,56 +41,64 @@ 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(&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
+        consul
             .kv_put(
-                &consul_key,
+                &self.consul_key,
                 serde_json::to_vec(&AutodiscoverResult {
                     timestamp: timestamp(),
                     address: discovered_addr,
@@ -93,22 +108,6 @@ impl StunActor {
 
         Ok(())
     }
-
-    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,
-                })?,
-            )
-            .await?;
-
-        Ok(())
-    }
 }
 
 async fn get_mapped_addr(