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