dnscrypt 0.1.1

//! Synchronous and 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::io::{Read, Write};
use std::net::{SocketAddr, TcpStream, UdpSocket};
use std::time::{Duration, Instant};

/// 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";

// ── Synchronous transport ────────────────────────────────────────────────────

/// 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 fn send_dns_query_udp(resolver_ip: SocketAddr, query: &[u8]) -> Result<Vec<u8>, String> {
    let socket =
        UdpSocket::bind("0.0.0.0:0").map_err(|e| format!("Failed to bind UDP socket: {e}"))?;
    socket.set_read_timeout(Some(UDP_TIMEOUT)).ok();
    socket.set_write_timeout(Some(UDP_TIMEOUT)).ok();

    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) {
            last_err = format!("UDP send failed (attempt {attempt}): {e}");
            if attempt < UDP_MAX_ATTEMPTS {
                std::thread::sleep(UDP_RETRY_DELAY);
            }
            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 deadline = Instant::now() + UDP_TIMEOUT;
        loop {
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                last_err = format!("UDP recv timed out (attempt {attempt})");
                break;
            }
            socket.set_read_timeout(Some(remaining)).ok();
            match socket.recv_from(&mut buf) {
                Ok((n, src)) if src == resolver_ip => return Ok(buf[..n].to_vec()),
                // Packet from an unexpected source; keep waiting.
                Ok(_) => {}
                Err(e) => {
                    last_err = format!("UDP recv failed (attempt {attempt}): {e}");
                    break;
                }
            }
        }
        if attempt < UDP_MAX_ATTEMPTS {
            std::thread::sleep(UDP_RETRY_DELAY);
        }
    }

    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.
pub fn send_dns_query_tcp(resolver_ip: SocketAddr, query: &[u8]) -> Result<Vec<u8>, String> {
    let mut stream = TcpStream::connect_timeout(&resolver_ip, TCP_CONNECT_TIMEOUT)
        .map_err(|e| format!("TCP connect failed: {e}"))?;
    stream.set_read_timeout(Some(TCP_IO_TIMEOUT)).ok();
    stream.set_write_timeout(Some(TCP_IO_TIMEOUT)).ok();

    // DNS-over-TCP: 2-byte big-endian length prefix followed by the query.
    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(2 + query.len());
    payload.extend_from_slice(&len.to_be_bytes());
    payload.extend_from_slice(query);
    stream
        .write_all(&payload)
        .map_err(|e| format!("TCP write payload failed: {e}"))?;

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

    let mut resp_buf = vec![0u8; resp_len];
    stream
        .read_exact(&mut resp_buf)
        .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 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);
    }

    send_dns_query_udp(resolver_ip, query).map_or_else(
        |_| send_dns_query_tcp(resolver_ip, query),
        |resp| {
            let is_dnscrypt = resp.len() >= 8 && resp[0..8] == *DNSCRYPT_RESPONSE_MAGIC;
            let tc_bit_set = resp.len() >= 3 && (resp[2] & 0x02 != 0);

            if is_dnscrypt || !tc_bit_set {
                Ok(resp)
            } else {
                send_dns_query_tcp(resolver_ip, query)
            }
        },
    )
}

// ── Asynchronous transport ───────────────────────────────────────────────────

/// Send a DNS/`DNSCrypt` query over UDP asynchronously with automatic retry.
///
/// Only available when the `tokio` feature is enabled.
///
/// # Errors
///
/// Returns `Err` if the socket cannot be bound, or if every send/receive
/// attempt fails or times out.
#[cfg(feature = "tokio")]
pub async fn send_dns_query_udp_async(
    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;
        }

        // As in the sync path: reject packets from anyone but the resolver
        // we actually queried, since an unconnected UDP socket will happily
        // hand back a spoofed response from another sender.
        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[..n].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 asynchronously (length-prefixed framing).
///
/// Only available when the `tokio` feature is enabled.
///
/// # 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.
#[cfg(feature = "tokio")]
pub async fn send_dns_query_tcp_async(
    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(2 + 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 = u16::from_be_bytes(len_buf) as usize;

    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 asynchronously, falling back from UDP to TCP as needed.
///
/// Only available when the `tokio` feature is enabled.
///
/// # 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.
#[cfg(feature = "tokio")]
pub async fn send_dns_query_async(
    resolver_ip: SocketAddr,
    query: &[u8],
    force_tcp: bool,
) -> Result<Vec<u8>, String> {
    if force_tcp {
        return send_dns_query_tcp_async(resolver_ip, query).await;
    }

    match send_dns_query_udp_async(resolver_ip, query).await {
        Ok(resp) => {
            let is_dnscrypt = resp.len() >= 8 && resp[0..8] == *DNSCRYPT_RESPONSE_MAGIC;
            let tc_bit_set = resp.len() >= 3 && (resp[2] & 0x02 != 0);

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

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

    #[test]
    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();

            // Build a non-DNSCrypt response with TC bit set (byte 2 & 0x02 != 0)
            let mut resp = vec![0u8; 12];
            resp[0] = buf[0];
            resp[1] = buf[1];
            resp[2] = 0x02; // TC bit set

            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();

            // Return a dummy DNSCrypt response over TCP
            let mut resp = vec![0u8; 32];
            resp[0..8].copy_from_slice(DNSCRYPT_RESPONSE_MAGIC);

            let resp_len = u16::try_from(resp.len()).unwrap().to_be_bytes();
            stream.write_all(&resp_len).unwrap();
            stream.write_all(&resp).unwrap();
        });

        let query = b"test query";
        let result = send_dns_query(resolver_ip, query, false).unwrap();

        assert!(result.len() >= 8);
        assert_eq!(&result[0..8], DNSCRYPT_RESPONSE_MAGIC);
    }

    #[test]
    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![0xaa; 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).unwrap();
        assert_eq!(result, vec![0xaa; 16]);
    }

    #[test]
    fn test_send_dns_query_tcp_rejects_oversized_query() {
        // The size check happens after connecting, so a listener must be up
        // (it never needs to read anything before the client bails out).
        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).unwrap_err();
        assert!(err.contains("too large"));
    }

    #[test]
    fn test_send_dns_query_udp_rejects_unmatched_source() {
        // A packet from an address other than `resolver_ip` must be ignored
        // rather than accepted as the response.
        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();

            // Spoofed reply from an unrelated socket/port.
            let spoofer = UdpSocket::bind("127.0.0.1:0").unwrap();
            spoofer.send_to(b"spoofed", client_addr).unwrap();

            // Give the client a moment to (correctly) ignore the spoofed
            // packet, then reply for real from the expected source.
            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").unwrap();
        assert_eq!(result.len(), 12);
    }

    #[cfg(feature = "tokio")]
    mod async_tests {
        use super::*;

        #[tokio::test]
        async fn test_async_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_async(resolver_ip, b"test query", false)
                .await
                .unwrap();
            assert_eq!(&result[0..8], DNSCRYPT_RESPONSE_MAGIC);
        }

        #[tokio::test]
        async fn test_async_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_async(resolver_ip, b"forced tcp query", true)
                .await
                .unwrap();
            assert_eq!(result, vec![0xbb; 16]);
        }

        #[tokio::test]
        async fn test_async_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_async(resolver_ip, &oversized)
                .await
                .unwrap_err();
            assert!(err.contains("too large"));
        }

        #[tokio::test]
        async fn test_async_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_async(resolver_ip, b"query")
                .await
                .unwrap();
            assert_eq!(result.len(), 12);
        }
    }
}