1 files changed, 68 insertions, 39 deletions
diff --git a/src/diplonat.rs b/src/diplonat.rs
index 22ebd6e..96bac3b 100644
--- a/src/diplonat.rs
+++ b/src/diplonat.rs
@@ -1,49 +1,78 @@
-use anyhow::Result;
+use anyhow::{Context, Result};
+use futures::future::FutureExt;
 use tokio::try_join;
 
 use crate::{
-  config::ConfigOpts, consul_actor::ConsulActor, fw_actor::FirewallActor, igd_actor::IgdActor,
+    config::ConfigOpts, consul_actor::ConsulActor, fw_actor::FirewallActor, igd_actor::IgdActor,
+    stun_actor::StunActor,
 };
 
 pub struct Diplonat {
-  consul: ConsulActor,
-  firewall: FirewallActor,
-  igd: IgdActor,
+    consul: ConsulActor,
+    firewall: FirewallActor,
+    igd: Option<IgdActor>,
+    stun: StunActor,
 }
 
 impl Diplonat {
-  pub async fn new() -> Result<Self> {
-    let rt_cfg = ConfigOpts::from_env()?;
-    println!("{:#?}", rt_cfg);
-
-    let ca = ConsulActor::new(&rt_cfg.consul, &rt_cfg.consul.node_name);
-
-    let fw = FirewallActor::new(rt_cfg.firewall.refresh_time, &ca.rx_open_ports).await?;
-
-    let ia = IgdActor::new(
-      rt_cfg.igd.private_ip.as_ref().map(String::as_str),
-      rt_cfg.igd.refresh_time,
-      rt_cfg.igd.expiration_time,
-      &ca.rx_open_ports,
-    )
-    .await?;
-
-    let ctx = Self {
-      consul: ca,
-      igd: ia,
-      firewall: fw,
-    };
-
-    return Ok(ctx);
-  }
-
-  pub async fn listen(&mut self) -> Result<()> {
-    try_join!(
-      self.consul.listen(),
-      self.igd.listen(),
-      self.firewall.listen()
-    )?;
-
-    return Ok(());
-  }
+    pub async fn new() -> Result<Self> {
+        let rt_cfg = ConfigOpts::from_env().context("Parse configuration")?;
+        println!("{:#?}", rt_cfg);
+
+        let ca = ConsulActor::new(&rt_cfg.consul, &rt_cfg.consul.node_name);
+
+        let fw = FirewallActor::new(
+            rt_cfg.firewall.ipv6_only,
+            rt_cfg.firewall.refresh_time,
+            &ca.rx_open_ports,
+        )
+        .await
+        .context("Setup fireall actor")?;
+
+        let ia = match rt_cfg.igd {
+            Some(igdc) => Some(
+                IgdActor::new(
+                    igdc.private_ip,
+                    igdc.refresh_time,
+                    igdc.expiration_time,
+                    &ca.rx_open_ports,
+                )
+                .await
+                .context("Setup IGD actor")?,
+            ),
+            None => None,
+        };
+
+        let sa = StunActor::new(&rt_cfg.consul, &rt_cfg.stun, &rt_cfg.consul.node_name);
+
+        let ctx = Self {
+            consul: ca,
+            igd: ia,
+            firewall: fw,
+            stun: sa,
+        };
+
+        Ok(ctx)
+    }
+
+    pub async fn listen(&mut self) -> Result<()> {
+        let igd_opt = &mut self.igd;
+
+        try_join!(
+            self.consul.listen().map(|x| x.context("Run consul actor")),
+            async {
+                if let Some(igd) = igd_opt {
+                    igd.listen().await.context("Run IGD actor")
+                } else {
+                    Ok(())
+                }
+            },
+            self.firewall
+                .listen()
+                .map(|x| x.context("Run firewall actor")),
+            self.stun.listen().map(|x| x.context("Run STUN actor")),
+        )?;
+
+        Ok(())
+    }
 }