1 files changed, 35 insertions, 23 deletions

diff --git a/src/fw_actor.rs b/src/fw_actor.rs
index fe68381..02d8bcb 100644
--- a/src/fw_actor.rs
+++ b/src/fw_actor.rs
@@ -12,7 +12,7 @@ use tokio::{
 use crate::{fw, messages};
 
 pub struct FirewallActor {
-    pub ipt_v4: iptables::IPTables,
+    pub ipt_v4: Option<iptables::IPTables>,
     pub ipt_v6: iptables::IPTables,
     rx_ports: watch::Receiver<messages::PublicExposedPorts>,
     last_ports: messages::PublicExposedPorts,
@@ -21,18 +21,24 @@ pub struct FirewallActor {
 
 impl FirewallActor {
     pub async fn new(
+        ipv6_only: bool,
         refresh: Duration,
         rxp: &watch::Receiver<messages::PublicExposedPorts>,
     ) -> Result<Self> {
         let ctx = Self {
-            ipt_v4: iptables::new(false)?,
+            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)?;
+        if let Some(ipt_v4) = &ctx.ipt_v4 {
+            fw::setup(ipt_v4)?;
+        }
         fw::setup(&ctx.ipt_v6)?;
 
         return Ok(ctx);
@@ -62,29 +68,35 @@ impl FirewallActor {
     }
 
     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)?;
+        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(());
     }