path: root/src/igd_actor.rs

use std::net::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: String,
}

impl IgdActor {
  pub async fn new(
    priv_ip: Option<&str>,
    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.to_string()
    } 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 public_ip = get_if_addrs::get_if_addrs()?
        .into_iter()
        .map(|i| i.addr.ip())
        .filter(|a| match a {
          std::net::IpAddr::V4(a4) => (a4.octets()[..cmplen] == gwa[..cmplen]),
          _ => false,
        })
        .next()
        .expect("No interface has an IP on same subnet as gateway")
        .to_string();
      info!("Found private IP: {}", public_ip);
      public_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_str = format!("{}:{}", self.private_ip, port);
        let service = service_str
          .parse::<SocketAddrV4>()
          .context("Invalid socket address")?;
        self
          .gateway
          .add_port(
            *proto,
            *port,
            service,
            self.expire.as_secs() as u32,
            "diplonat",
          )
          .await?;
        debug!("IGD request successful for {:#?} {}", proto, service);
      }
    }

    return Ok(());
  }
}