path: root/src/stun_actor.rs

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

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

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

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

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

impl StunActor {
    pub fn new(
        consul_config: &RuntimeConfigConsul,
        stun_config: &RuntimeConfigStun,
        node: &str,
    ) -> Self {
        assert!(stun_config
            .stun_server_v4
            .map(|x| x.is_ipv4())
            .unwrap_or(true));
        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 {
            let ipv4_result = match self.stun_server_v4 {
                Some(stun_server_v4) => self.autodiscover_ip(stun_server_v4).await,
                None => self.autodiscover_none_ipv4().await,
            };
            if let Err(e) = ipv4_result {
                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: timestamp(),
                    address: Some(discovered_addr),
                })?,
            )
            .await?;

        Ok(())
    }

    async fn autodiscover_none_ipv4(&self) -> Result<()> {
        let consul_key = format!("diplonat/autodiscovery/ipv4/{}", self.node);

        self.consul
            .kv_put(
                &consul_key,
                serde_json::to_vec(&AutodiscoverResult {
                    timestamp: timestamp(),
                    address: None,
                })?,
            )
            .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)
}

fn timestamp() -> u64 {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .expect("clock error")
        .as_secs()
}