Rust: udp SO_REUSEADDR tester

Cargo.toml

[package]
name = "udpbind"
version = "0.1.0"
edition = "2024"

[dependencies]
nix = { version = "0.30.1", features = ["net"] }
anyhow = "*"

main.rs

use anyhow::Context;
use nix::sys::socket::{self, AddressFamily, SockFlag, SockProtocol, SockType, SockaddrStorage};
use std::{
    net::{IpAddr, Ipv4Addr, SocketAddr, UdpSocket},
    os::fd::AsRawFd,
    thread::{self, JoinHandle},
    time::Duration,
};

const LISTEN_IPS: [IpAddr; 3] = [
    IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
    IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
    IpAddr::V4(Ipv4Addr::new(127, 1, 0, 254)),
];

const SEND_TO_IPS: [IpAddr; 2] = [
    IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
    IpAddr::V4(Ipv4Addr::new(127, 1, 0, 254)),
];

const LISTEN_PORT: u16 = 10053;

fn main() -> anyhow::Result<()> {
    // Create interface alias before running this executable:
    // sudo ifconfig lo0 alias 127.1.0.254 255.255.255.0
    // ifconfig lo0

    for reuse_addr_a in [true, false] {
        for addr_a in LISTEN_IPS.iter() {
            for addr_b in LISTEN_IPS.iter() {
                if addr_a == addr_b {
                    continue;
                }

                for reuse_addr_b in [true, false] {
                    let sockaddr_a = SocketAddr::new(*addr_a, LISTEN_PORT);
                    let sockaddr_b = SocketAddr::new(*addr_b, LISTEN_PORT);

                    println!("== Begin test case ==");
                    println!(
                        "{sockaddr_a} {}",
                        if reuse_addr_a {
                            "with SO_REUSEADDR"
                        } else {
                            "without SO_REUSEADDR"
                        }
                    );
                    println!(
                        "{sockaddr_b} {}",
                        if reuse_addr_b {
                            "with SO_REUSEADDR"
                        } else {
                            "without SO_REUSEADDR"
                        }
                    );
                    println!("====================");

                    match run_test_case(sockaddr_a, reuse_addr_a, sockaddr_b, reuse_addr_b) {
                        Ok(()) => println!("Test case passed"),
                        Err(e) => println!("Test case failed: {}", e),
                    }

                    println!("== End test case ==");
                    println!();
                }
            }
        }
    }

    Ok(())
}

fn run_test_case(
    sockaddr_a: SocketAddr,
    reuse_addr_a: bool,
    sockaddr_b: SocketAddr,
    reuse_addr_b: bool,
) -> anyhow::Result<()> {
    let socket_a = create_socket(sockaddr_a, reuse_addr_a)?;
    let socket_b = create_socket(sockaddr_b, reuse_addr_b)?;

    let sockaddr_a = if sockaddr_a.ip().is_unspecified() {
        SocketAddr::new(
            *SEND_TO_IPS
                .iter()
                .find(|ip| **ip != sockaddr_b.ip())
                .unwrap(),
            sockaddr_a.port(),
        )
    } else {
        sockaddr_a
    };
    let sockaddr_b = if sockaddr_b.ip().is_unspecified() {
        SocketAddr::new(
            *SEND_TO_IPS
                .iter()
                .find(|ip| **ip != sockaddr_a.ip())
                .unwrap(),
            sockaddr_b.port(),
        )
    } else {
        sockaddr_b
    };

    send_hello(sockaddr_a, true)
        .with_context(|| format!("Failed to send hello to {sockaddr_a}"))?;
    send_hello(sockaddr_b, false)
        .with_context(|| format!("Failed to send hello to {sockaddr_b}"))?;

    let handle_a = listen_socket(socket_a);
    let handle_b = listen_socket(socket_b);

    handle_a.join().expect("fail to join handle_a");
    handle_b.join().expect("fail to join handle_b");

    Ok(())
}

fn create_socket(listen_addr: SocketAddr, reuse_addr: bool) -> anyhow::Result<UdpSocket> {
    let sock = socket::socket(
        if listen_addr.is_ipv4() {
            AddressFamily::Inet
        } else {
            AddressFamily::Inet6
        },
        SockType::Datagram,
        SockFlag::empty(),
        SockProtocol::Udp,
    )
    .with_context(|| "Failed to open IPv4/UDP socket")?;

    if reuse_addr {
        if let Err(error) = socket::setsockopt(&sock, socket::sockopt::ReuseAddr, &true) {
            eprintln!("Failed to set SO_REUSEADDR on resolver socket: {error}");
        }
    }

    let sin = SockaddrStorage::from(listen_addr);
    socket::bind(sock.as_raw_fd(), &sin)
        .with_context(|| format!("Failed to bind DNS server to {}", listen_addr))?;

    Ok(UdpSocket::from(sock))
}

fn send_hello(ep: SocketAddr, is_a: bool) -> anyhow::Result<()> {
    if is_a {
        println!("Sending hello A to {ep}");
    } else {
        println!("Sending hello B to {ep}");
    }
    let any_addr = SocketAddr::V4("0.0.0.0:0".parse().unwrap());
    let socket =
        UdpSocket::bind(any_addr).with_context(|| format!("Failed to bind socket to {ep}"))?;
    let data: &[u8; 8] = if is_a { b"Hello A!" } else { b"Hello B!" };
    socket
        .send_to(data, ep)
        .with_context(|| format!("Failed to send data to {ep}"))?;
    Ok(())
}

fn listen_socket(socket: UdpSocket) -> JoinHandle<()> {
    let local_addr = socket.local_addr().unwrap();
    thread::spawn(move || {
        let mut buf = [0; 8];
        socket
            .set_read_timeout(Some(Duration::from_secs(2)))
            .expect("failed to set read timeout");
        match socket.recv_from(&mut buf) {
            Ok((amt, src)) => {
                let data = String::from_utf8_lossy(&buf[..amt]);
                println!("{local_addr}: received hello: {data} from {src}");
            }
            Err(error) => {
                eprintln!("{local_addr}: failed to receive data: {error}");
            }
        }
    })
}