path: root/src/stun_actor.rs

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, 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_ip = match stun_server.is_ipv4() {
            true => IpAddr::V4(Ipv4Addr::UNSPECIFIED),  // 0.0.0.0
            false => IpAddr::V6(Ipv6Addr::UNSPECIFIED), // [::]
        };
        let binding_addr = SocketAddr::new(binding_ip, 0);

        let discovered_addr = get_mapped_addr(stun_server, binding_addr)
            .await?
            .map(|x| x.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: 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<Option<SocketAddr>> {
    use stun_client::*;

    let mut client = Client::new(binding_addr, None).await?;
    let res = match client.binding_request(stun_server, None).await {
        Err(e) => {
            info!(
                "STUN binding request to {} failed, assuming no address (error: {})",
                binding_addr, e
            );
            return Ok(None);
        }
        Ok(r) => r,
    };

    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(Some(xor_mapped_addr))
}

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