path: root/src/igd_actor.rs

use std::net::{Ipv4Addr, SocketAddrV4};

use anyhow::{Context, Result};
use igd::{aio::*, PortMappingProtocol};
use log::*;
use tokio::{
    select,
    sync::watch,
    time::{self, Duration},
};

use crate::messages;

pub struct IgdActor {
    last_ports: messages::PublicExposedPorts,
    rx_ports: watch::Receiver<messages::PublicExposedPorts>,
    gateway: Gateway,
    refresh: Duration,
    expire: Duration,
    private_ip: Ipv4Addr,
}

impl IgdActor {
    pub async fn new(
        priv_ip: Option<Ipv4Addr>,
        refresh: Duration,
        expire: Duration,
        rxp: &watch::Receiver<messages::PublicExposedPorts>,
    ) -> Result<Self> {
        let gw = search_gateway(Default::default())
            .await
            .context("Failed to find IGD gateway")?;
        info!("IGD gateway: {}", gw);

        let private_ip = if let Some(ip) = priv_ip {
            info!("Using private IP from config: {}", ip);
            ip
        } else {
            info!("Trying to automatically detect private IP");
            let gwa = gw.addr.ip().octets();
            let cmplen = match gwa {
                [192, 168, _, _] => 3,
                [10, _, _, _] => 2,
                _ => panic!(
                    "Gateway IP does not appear to be in a local network ({})",
                    gw.addr.ip()
                ),
            };
            #[allow(unused_parens)]
            let private_ip = get_if_addrs::get_if_addrs()?
                .into_iter()
                .map(|i| i.addr.ip())
                .filter_map(|a| match a {
                    std::net::IpAddr::V4(a4) if a4.octets()[..cmplen] == gwa[..cmplen] => Some(a4),
                    _ => None,
                })
                .next()
                .expect("No interface has an IP on same subnet as gateway");
            info!("Autodetected private IP: {}", private_ip);
            private_ip
        };

        let ctx = Self {
            gateway: gw,
            rx_ports: rxp.clone(),
            private_ip,
            refresh: refresh,
            expire: expire,
            last_ports: messages::PublicExposedPorts::new(),
        };

        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. Flush IGD requests
            match self.do_igd().await {
                Ok(()) => debug!("Successfully updated IGD"),
                Err(e) => error!("An error occured while updating IGD. {}", e),
            }
        }
    }

    pub async fn do_igd(&self) -> Result<()> {
        let actions = [
            (PortMappingProtocol::TCP, &self.last_ports.tcp_ports),
            (PortMappingProtocol::UDP, &self.last_ports.udp_ports),
        ];

        for (proto, list) in actions.iter() {
            for port in *list {
                let service = SocketAddrV4::new(self.private_ip, *port);
                self.gateway
                    .add_port(
                        *proto,
                        *port,
                        service,
                        self.expire.as_secs() as u32,
                        "diplonat",
                    )
                    .await?;
                debug!("IGD request successful for {:#?} {}", proto, service);
            }
        }

        return Ok(());
    }
}