tailstun/src/main.rs
2025-02-14 18:55:57 -08:00

127 lines
3.6 KiB
Rust

use clap::ValueEnum;
use log::{debug, info};
use std::net::{IpAddr, ToSocketAddrs, UdpSocket};
use tailstun::{AddrPort, StunMessage, TxId};
#[derive(Debug, Clone, ValueEnum)]
enum OutputFormat {
Text,
Json,
Yaml,
}
impl OutputFormat {
fn format_stun(&self, msg: &StunMessage) -> String {
match self {
OutputFormat::Text => format!("{}", msg),
OutputFormat::Json => serde_json::to_string_pretty(msg).unwrap(),
OutputFormat::Yaml => serde_yaml::to_string(msg).unwrap(),
}
}
fn format_address(&self, a: &AddrPort) -> String {
let a = match a.address {
IpAddr::V4(_) => a.address,
IpAddr::V6(v6) => {
if let Some(v4) = v6.to_ipv4_mapped() {
IpAddr::V4(v4)
} else {
a.address
}
}
};
match self {
OutputFormat::Text => format!("{}", a),
OutputFormat::Json => serde_json::to_string_pretty(&a).unwrap(),
OutputFormat::Yaml => serde_yaml::to_string(&a).unwrap(),
}
}
}
#[derive(clap::Parser)]
#[command(version, about, long_about = None)]
#[command(about = "Test a Tailscale derp node's stun service")]
struct Cli {
host: String,
#[clap(default_value = "3478")]
port: u16,
#[clap(short = '4', conflicts_with = "v6_only", default_value = "false")]
v4_only: bool,
#[clap(short = '6', conflicts_with = "v4_only", default_value = "false")]
v6_only: bool,
#[clap(short, long, default_value = "text")]
format: OutputFormat,
#[clap(
short,
long,
default_value = "false",
help = "Only output the first mapped address & convert IPv6-mapped to IPv4"
)]
address_only: bool,
#[command(flatten)]
verbose: clap_verbosity_flag::Verbosity,
}
fn main() {
let cli = <Cli as clap::Parser>::parse();
env_logger::Builder::new()
.filter_level(cli.verbose.log_level_filter())
.init();
let dest = (cli.host.as_str(), cli.port)
.to_socket_addrs()
.expect("Unable to resolve host")
.find(|a| {
if cli.v4_only {
a.is_ipv4()
} else if cli.v6_only {
a.is_ipv6()
} else {
true
}
})
.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 = TxId::new().make_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));
}
}
Err(e) => {
println!("recv function failed: {e:?}");
}
}
}