1 files changed, 77 insertions, 65 deletions
diff --git a/src/fw_actor.rs b/src/fw_actor.rs
index 518c1b8..02d8bcb 100644
--- a/src/fw_actor.rs
+++ b/src/fw_actor.rs
@@ -4,88 +4,100 @@ use anyhow::Result;
 use iptables;
 use log::*;
 use tokio::{
-  select,
-  sync::watch,
-  time::{self, Duration},
+    select,
+    sync::watch,
+    time::{self, Duration},
 };
 
 use crate::{fw, messages};
 
 pub struct FirewallActor {
-  pub ipt_v4: iptables::IPTables,
-  pub ipt_v6: iptables::IPTables,
-  rx_ports: watch::Receiver<messages::PublicExposedPorts>,
-  last_ports: messages::PublicExposedPorts,
-  refresh: Duration,
+    pub ipt_v4: Option<iptables::IPTables>,
+    pub ipt_v6: iptables::IPTables,
+    rx_ports: watch::Receiver<messages::PublicExposedPorts>,
+    last_ports: messages::PublicExposedPorts,
+    refresh: Duration,
 }
 
 impl FirewallActor {
-  pub async fn new(
-    refresh: Duration,
-    rxp: &watch::Receiver<messages::PublicExposedPorts>,
-  ) -> Result<Self> {
-    let ctx = Self {
-      ipt_v4: iptables::new(false)?,
-      ipt_v6: iptables::new(true)?,
-      rx_ports: rxp.clone(),
-      last_ports: messages::PublicExposedPorts::new(),
-      refresh,
-    };
+    pub async fn new(
+        ipv6_only: bool,
+        refresh: Duration,
+        rxp: &watch::Receiver<messages::PublicExposedPorts>,
+    ) -> Result<Self> {
+        let ctx = Self {
+            ipt_v4: match ipv6_only {
+                false => Some(iptables::new(false)?),
+                true => None,
+            },
+            ipt_v6: iptables::new(true)?,
+            rx_ports: rxp.clone(),
+            last_ports: messages::PublicExposedPorts::new(),
+            refresh,
+        };
 
-    fw::setup(&ctx.ipt_v4)?;
-    fw::setup(&ctx.ipt_v6)?;
+        if let Some(ipt_v4) = &ctx.ipt_v4 {
+            fw::setup(ipt_v4)?;
+        }
+        fw::setup(&ctx.ipt_v6)?;
 
-    return Ok(ctx);
-  }
+        return Ok(ctx);
+    }
 
-  pub async fn listen(&mut self) -> Result<()> {
-    let mut interval = time::interval(self.refresh);
-    loop {
-      // 1. Wait for an event
-      let new_ports = select! {
-        _ = self.rx_ports.changed() => Some(self.rx_ports.borrow().clone()),
-        _ = interval.tick() => None,
-        else => return Ok(()) // Sender dropped, terminate loop.
-      };
+    pub async fn listen(&mut self) -> Result<()> {
+        let mut interval = time::interval(self.refresh);
+        loop {
+            // 1. Wait for an event
+            let new_ports = select! {
+              _ = self.rx_ports.changed() => Some(self.rx_ports.borrow().clone()),
+              _ = interval.tick() => None,
+              else => return Ok(()) // Sender dropped, terminate loop.
+            };
 
-      // 2. Update last ports if needed
-      if let Some(p) = new_ports {
-        self.last_ports = p;
-      }
+            // 2. Update last ports if needed
+            if let Some(p) = new_ports {
+                self.last_ports = p;
+            }
 
-      // 3. Update firewall rules
-      match self.do_fw_update().await {
-        Ok(()) => debug!("Successfully updated firewall rules"),
-        Err(e) => error!("An error occured while updating firewall rules. {}", e),
-      }
+            // 3. Update firewall rules
+            match self.do_fw_update().await {
+                Ok(()) => debug!("Successfully updated firewall rules"),
+                Err(e) => error!("An error occured while updating firewall rules. {}", e),
+            }
+        }
     }
-  }
 
-  pub async fn do_fw_update(&self) -> Result<()> {
-    for ipt in [&self.ipt_v4, &self.ipt_v6] {
-      let curr_opened_ports = fw::get_opened_ports(ipt)?;
+    pub async fn do_fw_update(&self) -> Result<()> {
+        if let Some(ipt_v4) = &self.ipt_v4 {
+            self.do_fw_update_on(ipt_v4).await?;
+        }
+        self.do_fw_update_on(&self.ipt_v6).await?;
+        Ok(())
+    }
 
-      let diff_tcp = self
-        .last_ports
-        .tcp_ports
-        .difference(&curr_opened_ports.tcp_ports)
-        .copied()
-        .collect::<HashSet<u16>>();
-      let diff_udp = self
-        .last_ports
-        .udp_ports
-        .difference(&curr_opened_ports.udp_ports)
-        .copied()
-        .collect::<HashSet<u16>>();
+    pub async fn do_fw_update_on(&self, ipt: &iptables::IPTables) -> Result<()> {
+        let curr_opened_ports = fw::get_opened_ports(ipt)?;
 
-      let ports_to_open = messages::PublicExposedPorts {
-        tcp_ports: diff_tcp,
-        udp_ports: diff_udp,
-      };
+        let diff_tcp = self
+            .last_ports
+            .tcp_ports
+            .difference(&curr_opened_ports.tcp_ports)
+            .copied()
+            .collect::<HashSet<u16>>();
+        let diff_udp = self
+            .last_ports
+            .udp_ports
+            .difference(&curr_opened_ports.udp_ports)
+            .copied()
+            .collect::<HashSet<u16>>();
 
-      fw::open_ports(ipt, ports_to_open)?;
-    }
+        let ports_to_open = messages::PublicExposedPorts {
+            tcp_ports: diff_tcp,
+            udp_ports: diff_udp,
+        };
+
+        fw::open_ports(ipt, ports_to_open)?;
 
-    return Ok(());
-  }
+        return Ok(());
+    }
 }