use super::PublicIpError;
use super::stun::{
STUN_SERVERS, StunError, StunFamily, get_public_ip_stun_family_with_servers_and_cache,
get_public_ip_stun_with_servers_and_cache,
};
use super::stun_cache::StunCache;
use serde::{Deserialize, Serialize};
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::RwLock;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct PublicIps {
pub v4: Option<Ipv4Addr>,
pub v6: Option<Ipv6Addr>,
}
fn map_stun_error(e: StunError) -> PublicIpError {
match e {
StunError::Timeout => PublicIpError::Timeout,
StunError::IoError(err) => PublicIpError::HttpError(err.to_string()),
StunError::InvalidResponse => {
PublicIpError::ParseError("Invalid STUN response".to_string())
}
StunError::NoMappedAddress => {
PublicIpError::ParseError("No mapped address in STUN response".to_string())
}
}
}
#[derive(Clone, Debug)]
pub struct StunClient {
cache: Arc<RwLock<StunCache>>,
servers: Vec<String>,
timeout: Duration,
verbose: u8,
}
impl StunClient {
pub fn new() -> Self {
Self::with_servers(STUN_SERVERS.iter().map(|s| (*s).to_string()).collect())
}
pub fn with_servers(servers: Vec<String>) -> Self {
Self {
cache: Arc::new(RwLock::new(StunCache::new())),
servers,
timeout: Duration::from_millis(500),
verbose: 0,
}
}
pub fn with_timeout(mut self, timeout: Duration) -> Self {
self.timeout = timeout;
self
}
pub fn with_verbose(mut self, verbose: u8) -> Self {
self.verbose = verbose;
self
}
pub fn with_cache(cache: StunCache, servers: Vec<String>) -> Self {
Self {
cache: Arc::new(RwLock::new(cache)),
servers,
timeout: Duration::from_millis(500),
verbose: 0,
}
}
pub async fn get_public_ip(&self) -> Result<IpAddr, PublicIpError> {
get_public_ip_stun_with_servers_and_cache(
&self.servers,
self.timeout,
&self.cache,
self.verbose,
)
.await
.map_err(map_stun_error)
}
pub async fn get_public_ip_v4(&self) -> Result<Ipv4Addr, PublicIpError> {
let ip = get_public_ip_stun_family_with_servers_and_cache(
&self.servers,
self.timeout,
&self.cache,
self.verbose,
StunFamily::V4,
)
.await
.map_err(map_stun_error)?;
match ip {
IpAddr::V4(v4) => Ok(v4),
IpAddr::V6(v6) => Err(PublicIpError::ParseError(format!(
"STUN IPv4 query returned IPv6 address {v6}"
))),
}
}
pub async fn get_public_ip_v6(&self) -> Result<Ipv6Addr, PublicIpError> {
let ip = get_public_ip_stun_family_with_servers_and_cache(
&self.servers,
self.timeout,
&self.cache,
self.verbose,
StunFamily::V6,
)
.await
.map_err(map_stun_error)?;
match ip {
IpAddr::V6(v6) => Ok(v6),
IpAddr::V4(v4) => Err(PublicIpError::ParseError(format!(
"STUN IPv6 query returned IPv4 address {v4}"
))),
}
}
pub async fn get_public_ips(&self) -> PublicIps {
let (v4, v6) = tokio::join!(self.get_public_ip_v4(), self.get_public_ip_v6());
PublicIps {
v4: v4.ok(),
v6: v6.ok(),
}
}
pub fn servers(&self) -> &[String] {
&self.servers
}
pub async fn clear_cache(&self) {
let cache = self.cache.write().await;
cache.clear();
}
pub async fn prewarm_cache(&self) -> Result<(), PublicIpError> {
let cache = self.cache.read().await;
for server in &self.servers {
let _ = cache.get_stun_server_addrs(server).await;
}
Ok(())
}
pub async fn is_server_cached(&self, server: &str) -> bool {
let cache = self.cache.read().await;
cache.get_stun_server_addrs(server).await.is_ok()
}
}
impl Default for StunClient {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct CacheStats {
pub servers_cached: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_stun_client_default() {
let client = StunClient::new();
assert_eq!(client.servers().len(), STUN_SERVERS.len());
assert!(client.servers()[0].contains("google.com"));
}
#[tokio::test]
async fn test_stun_client_custom_servers() {
let servers = vec![
"stun.example.com:3478".to_string(),
"stun2.example.com:3478".to_string(),
];
let client = StunClient::with_servers(servers.clone());
assert_eq!(client.servers(), &servers);
}
#[tokio::test]
async fn test_stun_client_timeout() {
let client = StunClient::new().with_timeout(Duration::from_secs(2));
assert_eq!(client.servers().len(), STUN_SERVERS.len());
}
#[tokio::test]
async fn test_custom_servers_are_queried() {
let client = StunClient::with_servers(vec!["stun.does-not-exist.invalid:3478".to_string()])
.with_timeout(Duration::from_millis(200));
let result = client.get_public_ip().await;
assert!(
result.is_err(),
"custom unresolvable server must not fall back to default servers"
);
}
#[tokio::test]
async fn test_cache_operations() {
let client = StunClient::new();
client.clear_cache().await;
let _ = client.prewarm_cache().await;
}
#[tokio::test]
async fn test_public_ip_detection() {
let client = StunClient::new();
match client.get_public_ip().await {
Ok(ip) => {
assert!(!ip.is_unspecified());
assert!(!ip.is_loopback());
}
Err(_) => {
}
}
}
}