1 files changed, 36 insertions, 19 deletions

diff --git a/src/diplonat.rs b/src/diplonat.rs
index 7049530..6334e5b 100644
--- a/src/diplonat.rs
+++ b/src/diplonat.rs
@@ -1,4 +1,4 @@
-use anyhow::Result;
+use anyhow::{Result, anyhow};
 use tokio::try_join;
 
 use crate::config::ConfigOpts;
@@ -8,43 +8,60 @@ use crate::igd_actor::IgdActor;
 
 pub struct Diplonat {
   consul: ConsulActor,
-  firewall: FirewallActor,
-  igd: IgdActor,
+
+  firewall: Option<FirewallActor>,
+  igd: Option<IgdActor>,
 }
 
 impl Diplonat {
   pub async fn new() -> Result<Self> {
-    let rt_cfg = ConfigOpts::from_env()?;
-    println!("{:#?}", rt_cfg);
+    let config = ConfigOpts::from_env()?;
+    println!("{:#?}", config);
     
-    let ca = ConsulActor::new(&rt_cfg.consul.url, &rt_cfg.consul.node_name);
+    let consul_actor = ConsulActor::new(config.consul);
 
-    let fw = FirewallActor::new(
-        rt_cfg.firewall.refresh_time,
-        &ca.rx_open_ports
+    let firewall_actor = FirewallActor::new(
+        config.firewall,
+        &consul_actor.rx_open_ports
     ).await?;
     
-    let ia = IgdActor::new(
-      &rt_cfg.igd.private_ip, 
-      rt_cfg.igd.refresh_time, 
-      rt_cfg.igd.expiration_time, 
-      &ca.rx_open_ports
+    let igd_actor = IgdActor::new(
+      config.igd,
+      &consul_actor.rx_open_ports
     ).await?;
 
+    if firewall_actor.is_none() && igd_actor.is_none() {
+      return Err(anyhow!(
+        "At least enable *one* module, otherwise it's boring!"));
+    }
+
     let ctx = Self {
-      consul: ca,
-      igd: ia,
-      firewall: fw
+      consul: consul_actor,
+      firewall: firewall_actor,
+      igd: igd_actor,
     };
 
     return Ok(ctx);
   }
 
   pub async fn listen(&mut self) -> Result<()> {
+    let firewall = &mut self.firewall;
+    let igd = &mut self.igd;
+
     try_join!(
       self.consul.listen(),
-      self.igd.listen(),
-      self.firewall.listen()
+      async {
+        match firewall {
+          Some(x) => x.listen().await,
+          None => Ok(())
+        }
+      },
+      async {
+        match igd {
+          Some(x) => x.listen().await,
+          None => Ok(())
+        }
+      },
     )?;
 
     return Ok(());