1 files changed, 147 insertions, 0 deletions
diff --git a/src/stun_actor.rs b/src/stun_actor.rs
new file mode 100644
index 0000000..6740c83
--- /dev/null
+++ b/src/stun_actor.rs
@@ -0,0 +1,147 @@
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
+use std::time::{Duration, SystemTime};
+
+use anyhow::{anyhow, bail, Result};
+use log::*;
+use serde::{Deserialize, Serialize};
+
+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,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct AutodiscoverResult {
+    pub timestamp: u64,
+    pub address: Option<IpAddr>,
+}
+
+impl StunActor {
+    pub fn new(
+        consul_config: &RuntimeConfigConsul,
+        stun_config: &RuntimeConfigStun,
+        node: &str,
+    ) -> Self {
+        assert!(stun_config
+            .stun_server_v4
+            .map(|x| x.is_ipv4())
+            .unwrap_or(true));
+        assert!(stun_config.stun_server_v6.is_ipv6());
+
+        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,
+            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 {
+                error!("Unable to autodiscover IPv4 address: {}", e);
+            }
+            if let Err(e) = self.autodiscover_ip(self.stun_server_v6).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() {
+            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 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)
+            }
+        };
+
+        self.consul
+            .kv_put(
+                &consul_key,
+                serde_json::to_vec(&AutodiscoverResult {
+                    timestamp: timestamp(),
+                    address: discovered_addr,
+                })?,
+            )
+            .await?;
+
+        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(
+    stun_server: SocketAddr,
+    binding_addr: SocketAddr,
+) -> Result<Option<SocketAddr>> {
+    use stun_client::*;
+
+    let mut client = Client::new(binding_addr, None).await?;
+    let res = match client.binding_request(stun_server, None).await {
+        Err(e) => {
+            info!(
+                "STUN binding request to {} failed, assuming no address (error: {})",
+                binding_addr, e
+            );
+            return Ok(None);
+        }
+        Ok(r) => r,
+    };
+
+    if res.get_class() != Class::SuccessResponse {
+        bail!("STUN server did not responde with a success response");
+    }
+
+    let xor_mapped_addr = Attribute::get_xor_mapped_address(&res)
+        .ok_or(anyhow!("no XorMappedAddress found in STUN response"))?;
+
+    Ok(Some(xor_mapped_addr))
+}
+
+fn timestamp() -> u64 {
+    SystemTime::now()
+        .duration_since(SystemTime::UNIX_EPOCH)
+        .expect("clock error")
+        .as_secs()
+}