path: root/src/stun_actor.rs

use std::net::{SocketAddr, IpAddr};
use std::time::{Duration, SystemTime};

use anyhow::{Result, bail, anyhow};
use log::*;
use serde::{Serialize, Deserialize};

use crate::config::{RuntimeConfigConsul, RuntimeConfigStun};
use crate::consul;

pub struct StunActor {
  node: String,
  consul: consul::Consul,
  stun_server_v4: SocketAddr,
  stun_server_v6: SocketAddr,
  refresh_time: Duration,
}

#[derive(Serialize, Deserialize, Debug)]
pub struct AutodiscoverResult {
  pub timestamp: u64,
  pub address: IpAddr,
}

impl StunActor {
  pub fn new(consul_config: &RuntimeConfigConsul, stun_config: &RuntimeConfigStun, node: &str) -> Self {
    assert!(stun_config.stun_server_v4.is_ipv4());
    assert!(stun_config.stun_server_v6.is_ipv6());

    Self {
      consul: consul::Consul::new(consul_config),
      node: node.to_string(),
      stun_server_v4: stun_config.stun_server_v4,
      stun_server_v6: stun_config.stun_server_v6,
      refresh_time: stun_config.refresh_time,
    }
  }

  pub async fn listen(&mut self) -> Result<()> {
    loop {
      if let Err(e) = self.autodiscover_ip(self.stun_server_v4).await {
        error!("Unable to autodiscover IPv4 address: {}", e);
      }
      if let Err(e) = self.autodiscover_ip(self.stun_server_v6).await {
        error!("Unable to autodiscover IPv6 address: {}", e);
      }
      tokio::time::sleep(self.refresh_time).await;
    }
  }

  async fn autodiscover_ip(&self, stun_server: SocketAddr) -> Result<()> {
    let binding_addr = match stun_server.is_ipv4() {
      true => "0.0.0.0:34791".parse().unwrap(),
      false => "[::]:34792".parse().unwrap(),
    };

    let discovered_addr = get_mapped_addr(stun_server, binding_addr).await?.ip();

    let consul_key = match stun_server.is_ipv4() {
      true => {
        debug!("Autodiscovered IPv4: {}", discovered_addr);
        format!("diplonat/autodiscovery/ipv4/{}", self.node)
      }
      false => {
        debug!("Autodiscovered IPv6: {}", discovered_addr);
        format!("diplonat/autodiscovery/ipv6/{}", self.node)
      }
    };

    self.consul.kv_put(&consul_key, serde_json::to_vec(&AutodiscoverResult{
      timestamp: SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_secs(),
      address: discovered_addr,
    })?).await?;

    Ok(())
  }
}

async fn get_mapped_addr(stun_server: SocketAddr, binding_addr: SocketAddr) -> Result<SocketAddr> {
  use stun_client::*;

  let mut client = Client::new(binding_addr, None).await.unwrap();
  let res = client
    .binding_request(stun_server, None)
    .await
    .unwrap();

  if res.get_class() != Class::SuccessResponse {
    bail!("STUN server did not responde with a success response");
  }

  let xor_mapped_addr = Attribute::get_xor_mapped_address(&res)
      .ok_or(anyhow!("no XorMappedAddress found in STUN response"))?;
  Ok(xor_mapped_addr)
}