path: root/src/fw_actor.rs

use std::collections::HashSet;

use anyhow::Result;
use iptables;
use log::*;
use tokio::{
    select,
    sync::watch,
    time::{self, Duration},
};

use crate::{fw, messages};

pub struct FirewallActor {
    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(
        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,
        };

        if let Some(ipt_v4) = &ctx.ipt_v4 {
            fw::setup(ipt_v4)?;
        }
        fw::setup(&ctx.ipt_v6)?;

        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.
            };

            // 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),
            }
        }
    }

    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(())
    }

    pub async fn do_fw_update_on(&self, ipt: &iptables::IPTables) -> Result<()> {
        let curr_opened_ports = fw::get_opened_ports(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>>();

        let ports_to_open = messages::PublicExposedPorts {
            tcp_ports: diff_tcp,
            udp_ports: diff_udp,
        };

        fw::open_ports(ipt, ports_to_open)?;

        return Ok(());
    }
}