add timeout & backoff, some refactoring
This commit is contained in:
Keenan Tims
parent
1560535fbe
commit
b92c9f5503
80
Cargo.lock
generated
80
Cargo.lock
generated
@ -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"
|
||||
|
||||
@ -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"
|
||||
|
||||
134
src/main.rs
134
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<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,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));
|
||||
}
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user