dnscrypt 0.1.0

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

/// 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.
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;
        }

        match socket.recv_from(&mut buf) {
            Ok((n, _)) => return Ok(buf[..n].to_vec()),
            Err(e) => {
                last_err = format!("UDP recv failed (attempt {attempt}): {e}");
                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).
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 = (query.len() as u16).to_be_bytes();
    let mut payload = Vec::with_capacity(2 + query.len());
    payload.extend_from_slice(&len);
    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.
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);
    }

    match send_dns_query_udp(resolver_ip, query) {
        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 {
                send_dns_query_tcp(resolver_ip, query)
            } else {
                Ok(resp)
            }
        }
        Err(_) => 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.
#[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;
        }

        match timeout(UDP_TIMEOUT, socket.recv_from(&mut buf)).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.
#[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 = (query.len() as u16).to_be_bytes();
    let mut payload = Vec::with_capacity(2 + query.len());
    payload.extend_from_slice(&len);
    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.
#[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 {
                send_dns_query_tcp_async(resolver_ip, query).await
            } else {
                Ok(resp)
            }
        }
        Err(_) => send_dns_query_tcp_async(resolver_ip, query).await,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::{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(format!("127.0.0.1:{}", port)).unwrap();

        let resolver_ip = SocketAddr::new(
            std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
            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 = u16::from_be_bytes(len_buf) as usize;
            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 = (resp.len() as u16).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);
    }
}