dnscrypt 0.2.0

//! Asynchronous UDP/TCP transport for raw DNS packets.
//!
//! `DNSCrypt` packets are transported over plain UDP or TCP (port 8443 for the
//! bundled Quad9 resolvers).  TCP uses the standard DNS-over-TCP length-prefix
//! framing (`u16` big-endian length followed by the payload).
//!
//! UDP is tried first.  If the response is truncated (TC bit set) **or** the
//! UDP send/receive fails, the implementation transparently falls back to TCP.

use std::net::SocketAddr;
use std::time::Duration;

/// Number of UDP send/receive attempts before giving up.
const UDP_MAX_ATTEMPTS: u32 = 2;
/// Delay between UDP retry attempts.
const UDP_RETRY_DELAY: Duration = Duration::from_millis(200);
/// Timeout for individual UDP send and receive operations.
const UDP_TIMEOUT: Duration = Duration::from_secs(3);
/// Maximum size of a single UDP DNS response (bytes).
const UDP_BUF_SIZE: usize = 4096;
/// Timeout for the TCP connect call.
const TCP_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
/// Timeout for TCP read and write operations.
const TCP_IO_TIMEOUT: Duration = Duration::from_secs(5);

/// The 8-byte magic that `DNSCrypt` resolvers prepend to every response.
/// Used to distinguish a `DNSCrypt` response from a plain DNS response so we
/// can correctly detect truncation.
const DNSCRYPT_RESPONSE_MAGIC: &[u8; 8] = b"r6fnvWj8";

/// Send a DNS/`DNSCrypt` query over UDP with automatic retry.
///
/// Returns the raw response bytes on success.
///
/// # Errors
///
/// Returns `Err` if the socket cannot be bound, or if every send/receive
/// attempt fails or times out.
pub async fn send_dns_query_udp(resolver_ip: SocketAddr, query: &[u8]) -> Result<Vec<u8>, String> {
    use tokio::net::UdpSocket;
    use tokio::time::timeout;

    let socket = UdpSocket::bind("0.0.0.0:0")
        .await
        .map_err(|e| format!("Failed to bind UDP socket: {e}"))?;

    let mut buf = vec![0u8; UDP_BUF_SIZE];
    let mut last_err = String::new();

    for attempt in 1..=UDP_MAX_ATTEMPTS {
        if let Err(e) = socket.send_to(query, resolver_ip).await {
            last_err = format!("UDP send failed (attempt {attempt}): {e}");
            if attempt < UDP_MAX_ATTEMPTS {
                tokio::time::sleep(UDP_RETRY_DELAY).await;
            }
            continue;
        }

        // Loop on recv_from until we get a packet actually sent by the
        // resolver we queried: on unconnected UDP sockets any host that can
        // guess (or, on a shared network, observe) the ephemeral source port
        // can otherwise inject a spoofed response and win the race against
        // the real one. Bound the total wait so a flood of spoofed packets
        // can't extend the attempt past the configured timeout.
        let recv_matching = async {
            loop {
                match socket.recv_from(&mut buf).await {
                    Ok((n, src)) if src == resolver_ip => return Ok(n),
                    Ok(_) => {}
                    Err(e) => return Err(e),
                }
            }
        };

        match timeout(UDP_TIMEOUT, recv_matching).await {
            Ok(Ok(n)) => return Ok(buf.get(..n).unwrap_or_default().to_vec()),
            Ok(Err(e)) => {
                last_err = format!("UDP recv failed (attempt {attempt}): {e}");
                if attempt < UDP_MAX_ATTEMPTS {
                    tokio::time::sleep(UDP_RETRY_DELAY).await;
                }
            }
            Err(_) => {
                last_err = format!("UDP query timed out (attempt {attempt})");
                if attempt < UDP_MAX_ATTEMPTS {
                    tokio::time::sleep(UDP_RETRY_DELAY).await;
                }
            }
        }
    }

    Err(last_err)
}

/// Send a DNS/`DNSCrypt` query over TCP (length-prefixed framing).
///
/// Returns the raw response payload (without the 2-byte length prefix).
///
/// # Errors
///
/// Returns `Err` if `query` exceeds the DNS-over-TCP length-prefix limit
/// (65535 bytes), or if the connect/write/read sequence fails or times out.
pub async fn send_dns_query_tcp(resolver_ip: SocketAddr, query: &[u8]) -> Result<Vec<u8>, String> {
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpStream;
    use tokio::time::timeout;

    let mut stream = timeout(TCP_CONNECT_TIMEOUT, TcpStream::connect(resolver_ip))
        .await
        .map_err(|_| "TCP connection timeout".to_string())?
        .map_err(|e| format!("TCP connect failed: {e}"))?;

    let len: u16 = query
        .len()
        .try_into()
        .map_err(|_| "query too large for DNS-over-TCP framing (max 65535 bytes)".to_string())?;
    let mut payload = Vec::with_capacity(2usize.saturating_add(query.len()));
    payload.extend_from_slice(&len.to_be_bytes());
    payload.extend_from_slice(query);

    timeout(TCP_IO_TIMEOUT, stream.write_all(&payload))
        .await
        .map_err(|_| "TCP write payload timeout".to_string())?
        .map_err(|e| format!("TCP write payload failed: {e}"))?;

    let mut len_buf = [0u8; 2];
    timeout(TCP_IO_TIMEOUT, stream.read_exact(&mut len_buf))
        .await
        .map_err(|_| "TCP read len timeout".to_string())?
        .map_err(|e| format!("TCP read len failed: {e}"))?;
    let resp_len = usize::from(u16::from_be_bytes(len_buf));

    let mut resp_buf = vec![0u8; resp_len];
    timeout(TCP_IO_TIMEOUT, stream.read_exact(&mut resp_buf))
        .await
        .map_err(|_| "TCP read data timeout".to_string())?
        .map_err(|e| format!("TCP read data failed: {e}"))?;

    Ok(resp_buf)
}

/// Send a DNS/`DNSCrypt` query, falling back from UDP to TCP as needed.
///
/// When `force_tcp` is `false` (the default for most queries):
/// - Tries UDP first.
/// - Falls back to TCP if UDP fails or if the response has the TC (truncated)
///   bit set and is not already a `DNSCrypt` encrypted packet.
///
/// Set `force_tcp = true` only when you already know the payload will exceed
/// the UDP MTU.
///
/// # Errors
///
/// Returns `Err` if both the UDP attempt (or its TC-bit-triggered TCP
/// fallback) and, when `force_tcp` is set, the direct TCP attempt fail.
pub async fn send_dns_query(
    resolver_ip: SocketAddr,
    query: &[u8],
    force_tcp: bool,
) -> Result<Vec<u8>, String> {
    if force_tcp {
        return send_dns_query_tcp(resolver_ip, query).await;
    }

    match send_dns_query_udp(resolver_ip, query).await {
        Ok(resp) => {
            let is_dnscrypt = resp.get(0..8) == Some(DNSCRYPT_RESPONSE_MAGIC.as_slice());
            let tc_bit_set = resp.get(2).is_some_and(|b| b & 0x02 != 0);

            if is_dnscrypt || !tc_bit_set {
                Ok(resp)
            } else {
                send_dns_query_tcp(resolver_ip, query).await
            }
        }
        Err(_) => send_dns_query_tcp(resolver_ip, query).await,
    }
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::indexing_slicing,
    clippy::arithmetic_side_effects,
    clippy::as_conversions
)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::{IpAddr, Ipv4Addr, TcpListener, UdpSocket};
    use std::thread;

    #[tokio::test]
    async fn test_fallback_udp_to_tcp() {
        let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
        let port = udp_socket.local_addr().unwrap().port();
        let tcp_listener = TcpListener::bind(("127.0.0.1", port)).unwrap();
        let resolver_ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);

        thread::spawn(move || {
            let mut buf = [0u8; 1024];
            let (_, src) = udp_socket.recv_from(&mut buf).unwrap();
            let mut resp = vec![0u8; 12];
            resp[0] = buf[0];
            resp[1] = buf[1];
            resp[2] = 0x02; // TC bit set, not a DNSCrypt response
            udp_socket.send_to(&resp, src).unwrap();
        });

        thread::spawn(move || {
            let (mut stream, _) = tcp_listener.accept().unwrap();
            let mut len_buf = [0u8; 2];
            stream.read_exact(&mut len_buf).unwrap();
            let len = usize::from(u16::from_be_bytes(len_buf));
            let mut query = vec![0u8; len];
            stream.read_exact(&mut query).unwrap();

            let mut resp = vec![0u8; 32];
            resp[0..8].copy_from_slice(DNSCRYPT_RESPONSE_MAGIC);
            stream
                .write_all(&u16::try_from(resp.len()).unwrap().to_be_bytes())
                .unwrap();
            stream.write_all(&resp).unwrap();
        });

        let result = send_dns_query(resolver_ip, b"test query", false)
            .await
            .unwrap();
        assert_eq!(&result[0..8], DNSCRYPT_RESPONSE_MAGIC);
    }

    #[tokio::test]
    async fn test_send_dns_query_force_tcp() {
        let tcp_listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = tcp_listener.local_addr().unwrap().port();
        let resolver_ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);

        thread::spawn(move || {
            let (mut stream, _) = tcp_listener.accept().unwrap();
            let mut len_buf = [0u8; 2];
            stream.read_exact(&mut len_buf).unwrap();
            let len = usize::from(u16::from_be_bytes(len_buf));
            let mut query = vec![0u8; len];
            stream.read_exact(&mut query).unwrap();

            let resp = vec![0xbb; 16];
            stream
                .write_all(&u16::try_from(resp.len()).unwrap().to_be_bytes())
                .unwrap();
            stream.write_all(&resp).unwrap();
        });

        let result = send_dns_query(resolver_ip, b"forced tcp query", true)
            .await
            .unwrap();
        assert_eq!(result, vec![0xbb; 16]);
    }

    #[tokio::test]
    async fn test_send_dns_query_tcp_rejects_oversized_query() {
        let tcp_listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let port = tcp_listener.local_addr().unwrap().port();
        let resolver_ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);
        thread::spawn(move || {
            let _ = tcp_listener.accept();
        });

        let oversized = vec![0u8; 65536];
        let err = send_dns_query_tcp(resolver_ip, &oversized)
            .await
            .unwrap_err();
        assert!(err.contains("too large"));
    }

    #[tokio::test]
    async fn test_send_dns_query_udp_rejects_unmatched_source() {
        let udp_socket = UdpSocket::bind("127.0.0.1:0").unwrap();
        let port = udp_socket.local_addr().unwrap().port();
        let resolver_ip = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), port);

        thread::spawn(move || {
            let mut buf = [0u8; 1024];
            let (_, client_addr) = udp_socket.recv_from(&mut buf).unwrap();

            let spoofer = UdpSocket::bind("127.0.0.1:0").unwrap();
            spoofer.send_to(b"spoofed", client_addr).unwrap();

            thread::sleep(Duration::from_millis(100));
            let mut resp = vec![0u8; 12];
            resp[0] = buf[0];
            resp[1] = buf[1];
            udp_socket.send_to(&resp, client_addr).unwrap();
        });

        let result = send_dns_query_udp(resolver_ip, b"query").await.unwrap();
        assert_eq!(result.len(), 12);
    }
}