dosh/src/udp.rs

//! UDP error classification shared by Dosh clients, servers, and embedders.

use std::net::SocketAddr;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;

pub fn is_transient_udp_error(err: &std::io::Error) -> bool {
    matches!(
        err.kind(),
        std::io::ErrorKind::Interrupted
            | std::io::ErrorKind::TimedOut
            | std::io::ErrorKind::WouldBlock
    ) || err.raw_os_error().is_some_and(is_transient_udp_os_error)
}

pub fn is_transient_udp_send_error(err: &std::io::Error) -> bool {
    is_transient_udp_error(err)
}

pub async fn send_udp_retrying_transient(
    socket: &UdpSocket,
    packet: &[u8],
    peer: SocketAddr,
    wait: Duration,
) -> std::io::Result<()> {
    let deadline = Instant::now() + wait.max(Duration::from_millis(1));
    loop {
        match socket.send_to(packet, peer).await {
            Ok(_) => return Ok(()),
            Err(err) if is_transient_udp_send_error(&err) && Instant::now() < deadline => {
                sleep_until_retry(deadline).await;
            }
            Err(err) => return Err(err),
        }
    }
}

pub async fn recv_udp_retrying_transient(
    socket: &UdpSocket,
    buf: &mut [u8],
    wait: Duration,
) -> std::io::Result<(usize, SocketAddr)> {
    let deadline = Instant::now() + wait.max(Duration::from_millis(1));
    loop {
        let now = Instant::now();
        if now >= deadline {
            return Err(std::io::Error::from(std::io::ErrorKind::TimedOut));
        }
        match tokio::time::timeout(deadline - now, socket.recv_from(buf)).await {
            Ok(Ok(value)) => return Ok(value),
            Ok(Err(err)) if is_transient_udp_error(&err) && Instant::now() < deadline => {
                sleep_until_retry(deadline).await;
            }
            Ok(Err(err)) => return Err(err),
            Err(_) => return Err(std::io::Error::from(std::io::ErrorKind::TimedOut)),
        }
    }
}

async fn sleep_until_retry(deadline: Instant) {
    let now = Instant::now();
    if now >= deadline {
        return;
    }
    tokio::time::sleep((deadline - now).min(Duration::from_millis(20))).await;
}

#[cfg(unix)]
fn is_transient_udp_os_error(code: i32) -> bool {
    matches!(
        code,
        libc::EADDRNOTAVAIL
            | libc::ECONNREFUSED
            | libc::ECONNRESET
            | libc::EHOSTDOWN
            | libc::EHOSTUNREACH
            | libc::ENETDOWN
            | libc::ENETRESET
            | libc::ENETUNREACH
            | libc::ETIMEDOUT
    )
}

#[cfg(windows)]
fn is_transient_udp_os_error(code: i32) -> bool {
    matches!(
        code,
        10049  // WSAEADDRNOTAVAIL
            | 10050 // WSAENETDOWN
            | 10051 // WSAENETUNREACH
            | 10052 // WSAENETRESET
            | 10054 // WSAECONNRESET
            | 10060 // WSAETIMEDOUT
            | 10061 // WSAECONNREFUSED
            | 10064 // WSAEHOSTDOWN
            | 10065 // WSAEHOSTUNREACH
    )
}

#[cfg(not(any(unix, windows)))]
fn is_transient_udp_os_error(_code: i32) -> bool {
    false
}

#[cfg(test)]
mod tests {
    use super::{
        is_transient_udp_error, is_transient_udp_send_error, recv_udp_retrying_transient,
        send_udp_retrying_transient,
    };
    use std::time::Duration;
    use tokio::net::UdpSocket;

    #[test]
    fn classifies_portable_transient_udp_errors() {
        for kind in [
            std::io::ErrorKind::Interrupted,
            std::io::ErrorKind::TimedOut,
            std::io::ErrorKind::WouldBlock,
        ] {
            let err = std::io::Error::from(kind);
            assert!(is_transient_udp_error(&err));
            assert!(is_transient_udp_send_error(&err));
        }
        let fatal = std::io::Error::from(std::io::ErrorKind::PermissionDenied);
        assert!(!is_transient_udp_error(&fatal));
        assert!(!is_transient_udp_send_error(&fatal));
    }

    #[tokio::test]
    async fn retrying_udp_send_and_receive_round_trip() {
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let sender = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let addr = receiver.local_addr().unwrap();

        send_udp_retrying_transient(&sender, b"hello", addr, Duration::from_millis(50))
            .await
            .unwrap();
        let mut buf = [0u8; 16];
        let (n, peer) = recv_udp_retrying_transient(&receiver, &mut buf, Duration::from_millis(50))
            .await
            .unwrap();

        assert_eq!(&buf[..n], b"hello");
        assert_eq!(peer, sender.local_addr().unwrap());
    }

    #[tokio::test]
    async fn retrying_udp_receive_times_out_cleanly() {
        let receiver = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let mut buf = [0u8; 16];
        let err = recv_udp_retrying_transient(&receiver, &mut buf, Duration::from_millis(1))
            .await
            .unwrap_err();

        assert_eq!(err.kind(), std::io::ErrorKind::TimedOut);
    }

    #[cfg(unix)]
    #[test]
    fn classifies_unix_network_churn_as_transient() {
        for code in [
            libc::EADDRNOTAVAIL,
            libc::ECONNREFUSED,
            libc::ECONNRESET,
            libc::EHOSTDOWN,
            libc::EHOSTUNREACH,
            libc::ENETDOWN,
            libc::ENETRESET,
            libc::ENETUNREACH,
            libc::ETIMEDOUT,
        ] {
            let err = std::io::Error::from_raw_os_error(code);
            assert!(is_transient_udp_error(&err));
            assert!(is_transient_udp_send_error(&err));
        }
    }

    #[cfg(windows)]
    #[test]
    fn classifies_windows_network_churn_as_transient() {
        for code in [
            10049, 10050, 10051, 10052, 10054, 10060, 10061, 10064, 10065,
        ] {
            let err = std::io::Error::from_raw_os_error(code);
            assert!(is_transient_udp_error(&err));
            assert!(is_transient_udp_send_error(&err));
        }
    }
}