add timeout & backoff, some refactoring

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_chacha 0.9.0",
- "rand_core",
+ "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"

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"

src/main.rs

@@ -1,6 +1,10 @@
+use backoff::ExponentialBackoff;
 use clap::ValueEnum;
-use log::{debug, info};
+use log::{debug, error, info};
-use std::net::{IpAddr, SocketAddr, ToSocketAddrs, UdpSocket};
+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<Duration, std::num::ParseFloatError> {
+    let secs = s.parse()?;
+    Ok(Duration::from_secs_f64(secs))
+}
+
+fn stun_query(
+    target: &SocketAddr,
+    timeout: Duration,
+) -> Result<StunMessage, backoff::Error<std::io::Error>> {
+    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 = <Cli as clap::Parser>::parse();
     env_logger::Builder::new()
@@ -81,32 +159,7 @@ fn main() {
         })
         .expect("No address found for host");
 
-    let socket = UdpSocket::bind("[::]:0").expect("Unable to bind a UDP socket");
+    let msg = stun_query(&dest, cli.timeout).expect("Failed to query STUN server");
-    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)),
@@ -118,9 +171,4 @@ fn main() {
     } else {
         println!("{}", cli.format.format_stun(&msg));
     }
-        }
-        Err(e) => {
-            println!("recv function failed: {e:?}");
-        }
-    }
 }