Add STUN actor that saves autodiscovered IPv4/IPv6 to Consul

1 files changed, 95 insertions, 0 deletions

diff --git a/src/stun_actor.rs b/src/stun_actor.rs
new file mode 100644
index 0000000..dee60f0
--- /dev/null
+++ b/src/stun_actor.rs
@@ -0,0 +1,95 @@
+use std::net::{SocketAddr, IpAddr};
+use std::time::{Duration, SystemTime};
+
+use anyhow::{Result, bail, anyhow};
+use log::*;
+use serde::{Serialize, Deserialize};
+
+use crate::config::{RuntimeConfigConsul, RuntimeConfigStun};
+use crate::consul;
+
+pub struct StunActor {
+  node: String,
+  consul: consul::Consul,
+  stun_server_v4: SocketAddr,
+  stun_server_v6: SocketAddr,
+  refresh_time: Duration,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct AutodiscoverResult {
+  pub timestamp: u64,
+  pub address: IpAddr,
+}
+
+impl StunActor {
+  pub fn new(consul_config: &RuntimeConfigConsul, stun_config: &RuntimeConfigStun, node: &str) -> Self {
+    assert!(stun_config.stun_server_v4.is_ipv4());
+    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 {
+      if let Err(e) = self.autodiscover_ip(self.stun_server_v4).await {
+        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_addr = match stun_server.is_ipv4() {
+      true => "0.0.0.0:34791".parse().unwrap(),
+      false => "[::]:34792".parse().unwrap(),
+    };
+
+    let discovered_addr = get_mapped_addr(stun_server, binding_addr).await?.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: SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_secs(),
+      address: discovered_addr,
+    })?).await?;
+
+    Ok(())
+  }
+}
+
+async fn get_mapped_addr(stun_server: SocketAddr, binding_addr: SocketAddr) -> Result<SocketAddr> {
+  use stun_client::*;
+
+  let mut client = Client::new(binding_addr, None).await.unwrap();
+  let res = client
+    .binding_request(stun_server, None)
+    .await
+    .unwrap();
+
+  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(xor_mapped_addr)
+}