corrections after review

2 files changed, 40 insertions, 35 deletions

diff --git a/src/fw.rs b/src/fw.rs
index 8ee3e6b..bc4d119 100644
--- a/src/fw.rs
+++ b/src/fw.rs
@@ -2,61 +2,58 @@ use iptables;
 use regex::Regex;
 use std::collections::HashSet;
 use crate::messages;
+use anyhow::{Result,Context};
+use log::*;
 
-#[derive(Debug)]
-pub struct FirewallError(String);
+pub fn setup(ipt: &iptables::IPTables) -> Result<()> {
 
-impl From<iptables::error::IPTError> for FirewallError {
-    fn from(error: iptables::error::IPTError) -> Self {
-        FirewallError(error.to_string())
-    }
-}
-
-pub fn setup(ipt: &iptables::IPTables) -> Result<(), FirewallError> {
-
-    if !ipt.chain_exists("filter", "DIPLONAT")? {
-        ipt.new_chain("filter", "DIPLONAT")?;
-    }
-
-    ipt.insert_unique("filter", "INPUT", "-j DIPLONAT", 1)?;
+    // ensure we start from a clean state without any rule already set
+    cleanup(ipt)?;
+    
+    ipt.new_chain("filter", "DIPLONAT").context("Failed to create new chain")?;
+    ipt.insert_unique("filter", "INPUT", "-j DIPLONAT", 1).context("Failed to insert jump rule")?;
 
     Ok(())
 }
 
-pub fn open_ports(ipt: &iptables::IPTables, ports: messages::PublicExposedPorts) -> Result<(), FirewallError> {
+pub fn open_ports(ipt: &iptables::IPTables, ports: messages::PublicExposedPorts) -> Result<()> {
     for p in ports.tcp_ports {
-        ipt.append("filter", "DIPLONAT", &format!("-p tcp --dport {} -j ACCEPT", p))?;
+        ipt.append("filter", "DIPLONAT", &format!("-p tcp --dport {} -j ACCEPT", p)).context("Failed to insert port rule")?;
     }
 
     for p in ports.udp_ports {
-        ipt.append("filter", "DIPLONAT", &format!("-p udp --dport {} -j ACCEPT", p))?;
+        ipt.append("filter", "DIPLONAT", &format!("-p udp --dport {} -j ACCEPT", p)).context("Failed to insert port rule")?;
     }
 
     Ok(())
 }
 
-pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<messages::PublicExposedPorts, FirewallError> {
+pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<messages::PublicExposedPorts> {
     let mut ports = messages::PublicExposedPorts {
         tcp_ports: HashSet::new(),
         udp_ports: HashSet::new()
     };
 
     let list = ipt.list("filter", "DIPLONAT")?;
-    let re = Regex::new(r"\-A.*? \-p (\w+).*\-\-dport (\d+).*?\-j ACCEPT").unwrap();
+    let re = Regex::new(r"\-A.*? \-p (\w+).*\-\-dport (\d+).*?\-j ACCEPT").context("Regex matching open ports encountered an unexpected rule")?;
     for i in list {
         let caps = re.captures(&i);
         match caps {
             Some(c) => {
-                let raw_proto = c.get(1).unwrap();
-                let raw_port = c.get(2).unwrap();
+           
+                if let (Some(raw_proto), Some(raw_port)) = (c.get(1), c.get(2)) {
+
+                    let proto = String::from(raw_proto.as_str());
+                    let number = String::from(raw_port.as_str()).parse::<u16>()?;
 
-                let proto = String::from(raw_proto.as_str());
-                let number = String::from(raw_port.as_str()).parse::<u16>().unwrap();
+                    if proto == "tcp" {
+                        ports.tcp_ports.insert(number);
+                    } else {
+                        ports.udp_ports.insert(number);
+                    }
 
-                if proto == "tcp" {
-                    ports.tcp_ports.insert(number);
                 } else {
-                    ports.udp_ports.insert(number);
+                    error!("Unexpected rule found in DIPLONAT chain")
                 }
 
             },
@@ -67,10 +64,18 @@ pub fn get_opened_ports(ipt: &iptables::IPTables) -> Result<messages::PublicExpo
     Ok(ports)
 }
 
-pub fn cleanup(ipt: &iptables::IPTables) -> Result<(), FirewallError> {
-    ipt.flush_chain("filter", "DIPLONAT")?;
-    ipt.delete("filter", "INPUT", "-j DIPLONAT")?;
-    ipt.delete_chain("filter", "DIPLONAT")?;
+pub fn cleanup(ipt: &iptables::IPTables) -> Result<()> {
+    
+    if ipt.chain_exists("filter", "DIPLONAT")? {
+        ipt.flush_chain("filter", "DIPLONAT").context("Failed to flush the DIPLONAT chain")?;
+    
+        if ipt.exists("filter", "INPUT", "-j DIPLONAT")? {
+            ipt.delete("filter", "INPUT", "-j DIPLONAT").context("Failed to delete jump rule")?;
+        }
+    
+        ipt.delete_chain("filter", "DIPLONAT").context("Failed to delete chain")?;
+    }
+
     Ok(())
 }
 
diff --git a/src/fw_actor.rs b/src/fw_actor.rs
index 523bdaa..b5e4c7e 100644
--- a/src/fw_actor.rs
+++ b/src/fw_actor.rs
@@ -23,13 +23,13 @@ pub struct FirewallActor {
 impl FirewallActor {
   pub async fn new(_refresh: Duration, rxp: &watch::Receiver<messages::PublicExposedPorts>) -> Result<Self> {
     let ctx = Self { 
-      ipt: iptables::new(false).unwrap(),
+      ipt: iptables::new(false)?,
       rx_ports: rxp.clone(),
       last_ports: messages::PublicExposedPorts::new(),
       refresh: _refresh,
     };
 
-    fw::setup(&ctx.ipt).expect("iptables setup failed");
+    fw::setup(&ctx.ipt)?;
     
     return Ok(ctx);
   }
@@ -56,7 +56,7 @@ impl FirewallActor {
   }
 
   pub async fn do_fw_update(&self) -> Result<()> {
-    let curr_opened_ports = fw::get_opened_ports(&self.ipt).unwrap();
+    let curr_opened_ports = fw::get_opened_ports(&self.ipt)?;
 
     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>>();
@@ -66,7 +66,7 @@ impl FirewallActor {
         udp_ports: diff_udp
     };
 
-    fw::open_ports(&self.ipt, ports_to_open).unwrap();
+    fw::open_ports(&self.ipt, ports_to_open)?;
 
     return Ok(());
   }