From b92c9f550322112116bdb033f5cea018634d5c80 Mon Sep 17 00:00:00 2001 From: Keenan Tims Date: Mon, 26 May 2025 15:49:35 -0700 Subject: [PATCH] add timeout & backoff, some refactoring --- Cargo.lock | 80 ++++++++++++++++++++++++++++--- Cargo.toml | 1 + src/main.rs | 134 +++++++++++++++++++++++++++++++++++----------------- 3 files changed, 166 insertions(+), 49 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8dfaf84..8387425 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -61,6 +61,17 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "backoff" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b62ddb9cb1ec0a098ad4bbf9344d0713fa193ae1a80af55febcff2627b6a00c1" +dependencies = [ + "getrandom 0.2.16", + "instant", + "rand 0.8.5", +] + [[package]] name = "bitflags" version = "2.8.0" @@ -173,6 +184,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" +[[package]] +name = "getrandom" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "335ff9f135e4384c8150d6f27c6daed433577f86b4750418338c01a1a2528592" +dependencies = [ + "cfg-if", + "libc", + "wasi 0.11.0+wasi-snapshot-preview1", +] + [[package]] name = "getrandom" version = "0.3.1" @@ -181,7 +203,7 @@ checksum = "43a49c392881ce6d5c3b8cb70f98717b7c07aabbdff06687b9030dbfbe2725f8" dependencies = [ "cfg-if", "libc", - "wasi", + "wasi 0.13.3+wasi-0.2.2", "windows-targets", ] @@ -213,6 +235,15 @@ dependencies = [ "hashbrown", ] +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + [[package]] name = "is_terminal_polyfill" version = "1.70.1" @@ -285,17 +316,38 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3779b94aeb87e8bd4e834cee3650289ee9e0d5677f976ecdb6d219e5f4f6cd94" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.0", "zerocopy 0.8.17", ] +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + [[package]] name = "rand_chacha" version = "0.9.0" @@ -303,7 +355,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.0", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.16", ] [[package]] @@ -312,7 +373,7 @@ version = "0.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b08f3c9802962f7e1b25113931d94f43ed9725bebc59db9d0c3e9a23b67e15ff" dependencies = [ - "getrandom", + "getrandom 0.3.1", "zerocopy 0.8.17", ] @@ -417,13 +478,14 @@ dependencies = [ name = "tailstun" version = "0.1.0" dependencies = [ + "backoff", "clap", "clap-verbosity-flag", "crc32fast", "env_logger", "log", "nom", - "rand", + "rand 0.9.0", "serde", "serde_json", "serde_yaml", @@ -447,6 +509,12 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "wasi" +version = "0.11.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c8d87e72b64a3b4db28d11ce29237c246188f4f51057d65a7eab63b7987e423" + [[package]] name = "wasi" version = "0.13.3+wasi-0.2.2" diff --git a/Cargo.toml b/Cargo.toml index 329524b..0643ae3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition = "2021" [dependencies] +backoff = "0.4.0" clap = { version = "4.5.29", features = ["derive"] } clap-verbosity-flag = "3.0.2" crc32fast = "1.4.2" diff --git a/src/main.rs b/src/main.rs index 03aff65..ef1c3b4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,6 +1,10 @@ +use backoff::ExponentialBackoff; use clap::ValueEnum; -use log::{debug, info}; -use std::net::{IpAddr, SocketAddr, ToSocketAddrs, UdpSocket}; +use log::{debug, error, info}; +use std::{ + net::{IpAddr, SocketAddr, ToSocketAddrs, UdpSocket}, + time::Duration, +}; use tailstun::StunMessage; #[derive(Debug, Clone, ValueEnum)] @@ -44,9 +48,19 @@ struct Cli { host: String, #[clap(default_value = "3478")] port: u16, - #[clap(short = '4', conflicts_with = "v6_only", default_value = "false")] + #[clap( + short = '4', + conflicts_with = "v6_only", + default_value = "false", + help = "Only use IPv4" + )] v4_only: bool, - #[clap(short = '6', conflicts_with = "v4_only", default_value = "false")] + #[clap( + short = '6', + conflicts_with = "v4_only", + default_value = "false", + help = "Only use IPv6" + )] v6_only: bool, #[clap(short, long, default_value = "text")] format: OutputFormat, @@ -57,10 +71,74 @@ struct Cli { help = "Only output the first mapped address & convert IPv6-mapped to IPv4" )] address_only: bool, + #[clap( + short, + long, + default_value = "5.0", + value_parser = parse_duration, + help = "Timeout in seconds" + )] + timeout: Duration, #[command(flatten)] verbose: clap_verbosity_flag::Verbosity, } +fn parse_duration(s: &str) -> Result { + let secs = s.parse()?; + Ok(Duration::from_secs_f64(secs)) +} + +fn stun_query( + target: &SocketAddr, + timeout: Duration, +) -> Result> { + let socket = UdpSocket::bind("[::]:0").expect("Unable to bind a UDP socket"); + socket + .connect(target) + .expect("Unable to connect to the destination"); + socket + .set_read_timeout(Some(timeout)) + .expect("Unable to set read timeout"); + debug!( + "Connected UDP socket to {:?} from {:?}", + socket.peer_addr(), + socket.local_addr() + ); + + debug!("Building request packet"); + let req = tailstun::rand_request(); + let backoff = ExponentialBackoff::default(); + + debug!("request {:?}", &req); + let op = || { + info!("Sending STUN request to {target} with timeout {timeout:?}"); + socket.send(&req).expect("Unable to send request"); + + let mut buf = [0u8; 1500]; + + match socket.recv(&mut buf) { + Ok(received) => { + let buf = &buf[..received]; + debug!("Received response: {:?}", buf); + + let msg = StunMessage::parse(buf).unwrap(); + info!("Parsed message from {}:", socket.peer_addr().unwrap()); + Ok(msg) + } + Err(e) => { + if e.kind() == std::io::ErrorKind::WouldBlock { + error!("Timed out waiting for response"); + Err(backoff::Error::transient(e)) + } else { + error!("recv function failed: {e:?}"); + Err(backoff::Error::permanent(e)) + } + } + } + }; + backoff::retry(backoff, op) +} + fn main() { let cli = ::parse(); env_logger::Builder::new() @@ -81,46 +159,16 @@ fn main() { }) .expect("No address found for host"); - let socket = UdpSocket::bind("[::]:0").expect("Unable to bind a UDP socket"); - socket - .connect(dest) - .expect("Unable to connect to the destination"); - - info!( - "Connected to {:?} from {:?}", - socket.peer_addr(), - socket.local_addr() - ); - - let req = tailstun::rand_request(); - - debug!("Sending request {:?}", &req); - - socket.send(&req).expect("Unable to send request"); - - let mut buf = [0u8; 1500]; - - match socket.recv(&mut buf) { - Ok(received) => { - let buf = &buf[..received]; - debug!("Received response: {:?}", buf); - - let msg = StunMessage::parse(buf).unwrap(); - info!("Parsed message from {}:", socket.peer_addr().unwrap()); - if cli.address_only { - match msg.attributes.mapped_address() { - Some(addr) => println!("{}", cli.format.format_address(addr)), - None => { - // No mapped address - std::process::exit(1); - } - } - } else { - println!("{}", cli.format.format_stun(&msg)); + let msg = stun_query(&dest, cli.timeout).expect("Failed to query STUN server"); + if cli.address_only { + match msg.attributes.mapped_address() { + Some(addr) => println!("{}", cli.format.format_address(addr)), + None => { + // No mapped address + std::process::exit(1); } } - Err(e) => { - println!("recv function failed: {e:?}"); - } + } else { + println!("{}", cli.format.format_stun(&msg)); } }