dnscrypt 0.2.0

//! DNS packet building, parsing, and padding utilities.
//!
//! This module handles raw DNS wire-format serialization and deserialization,
//! limited to the subset needed by `DNSCrypt`:
//! - Building A-record and TXT-record queries.
//! - Parsing TXT-record responses (for certificate fetching).
//! - Parsing A-record responses (for domain resolution).
//! - ISO/IEC 7816-4 padding applied to encrypted DNS payloads.

use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

// ── Padding ───────────────────────────────────────────────────────────────────

/// Pad `query` to a multiple of 64 bytes using ISO/IEC 7816-4 padding.
///
/// The result is at least `min_len` bytes. 0 to 2 extra 64-byte blocks of
/// random padding are also appended, to obscure the exact query size from
/// traffic analysts.
#[must_use]
pub fn pad_query(query: &[u8], min_len: usize) -> Vec<u8> {
    let mut padded = query.to_vec();
    padded.push(0x80);
    let current = padded.len().max(min_len);
    // Round up to the next 64-byte boundary.
    let base_target = current.next_multiple_of(64);

    // Add 0, 1, or 2 extra 64-byte blocks of random padding.
    let rb: [u8; 4] = crate::random_bytes();
    let extra_blocks = usize::try_from(u32::from_ne_bytes(rb) % 3).unwrap_or(0);
    let target = base_target.saturating_add(extra_blocks.saturating_mul(64));

    padded.resize(target, 0x00);
    padded
}

/// Strip ISO/IEC 7816-4 padding from a decrypted DNS response.
///
/// Scans backwards over trailing `0x00` bytes and removes the mandatory `0x80`
/// marker.
///
/// # Errors
///
/// Returns `Err` if no `0x80` marker is found (malformed or missing padding).
pub fn unpad_response(decrypted: &[u8]) -> Result<Vec<u8>, String> {
    // The last non-zero byte, scanning from the end, must be the 0x80 marker;
    // everything before it is the unpadded payload.
    let Some(marker_pos) = decrypted.iter().rposition(|&b| b != 0x00) else {
        return Err("Invalid ISO/IEC 7816-4 padding".to_string());
    };
    if decrypted.get(marker_pos) != Some(&0x80) {
        return Err("Invalid ISO/IEC 7816-4 padding".to_string());
    }
    Ok(decrypted.get(..marker_pos).unwrap_or_default().to_vec())
}

// ── Validation ────────────────────────────────────────────────────────────────

/// Validate that `domain` can be safely encoded as a DNS QNAME.
///
/// Rejects empty labels between dots, labels over 63 bytes, and names whose
/// encoded wire length exceeds 255 bytes. Without this check `build_query`
/// would silently truncate an oversized label, causing a query to be sent
/// for a *different* (truncated) domain than the caller requested.
///
/// # Errors
///
/// Returns `Err` with a human-readable reason if `domain` is empty, has a
/// label over 63 bytes, or encodes to more than 255 bytes.
pub fn validate_domain(domain: &str) -> Result<(), String> {
    if domain.is_empty() {
        return Err("domain name is empty".to_string());
    }

    let mut encoded_len = 1usize; // trailing root label
    for label in domain.split('.') {
        if label.is_empty() {
            continue;
        }
        if label.len() > 63 {
            return Err(format!("domain label '{label}' exceeds 63 bytes"));
        }
        encoded_len = encoded_len.saturating_add(1).saturating_add(label.len());
    }

    if encoded_len > 255 {
        return Err("domain name exceeds 255 bytes when encoded".to_string());
    }

    Ok(())
}

// ── Query builders ────────────────────────────────────────────────────────────

/// Build a minimal DNS query for `domain` requesting TXT records (type 16).
///
/// The transaction ID should be randomized.
#[must_use]
pub fn build_txt_record_query(domain: &str, txid: [u8; 2]) -> Vec<u8> {
    build_query(domain, 0x0010, txid) // QTYPE = TXT
}

/// Build a minimal DNS query for `domain` requesting A records (type 1).
///
/// The transaction ID should be randomized.
#[must_use]
pub fn build_a_record_query(domain: &str, txid: [u8; 2]) -> Vec<u8> {
    build_query(domain, 0x0001, txid) // QTYPE = A
}

/// Build a minimal DNS query for `domain` requesting AAAA records (type 28).
///
/// The transaction ID should be randomized.
#[must_use]
pub fn build_aaaa_record_query(domain: &str, txid: [u8; 2]) -> Vec<u8> {
    build_query(domain, 0x001c, txid) // QTYPE = AAAA
}

/// Internal helper: build a single-question DNS query packet.
fn build_query(domain: &str, qtype: u16, txid: [u8; 2]) -> Vec<u8> {
    // Fixed 12-byte DNS header: QR=0 (query), QDCOUNT=1, rest zero.
    let mut query: Vec<u8> = vec![
        txid[0], txid[1], // Transaction ID
        0x01, 0x00, // Flags: RD=1
        0x00, 0x01, // QDCOUNT = 1
        0x00, 0x00, // ANCOUNT = 0
        0x00, 0x00, // NSCOUNT = 0
        0x00, 0x00, // ARCOUNT = 0
    ];
    // Encode the QNAME in DNS label format.
    for label in domain.split('.') {
        if label.is_empty() {
            continue;
        }
        let len = label.len().min(63);
        query.push(u8::try_from(len).unwrap_or(63));
        query.extend_from_slice(label.as_bytes().get(..len).unwrap_or(&[]));
    }
    query.push(0x00); // root label
    query.extend_from_slice(&qtype.to_be_bytes());
    query.extend_from_slice(&[0x00, 0x01]); // QCLASS = IN
    query
}

// ── Response parsers ──────────────────────────────────────────────────────────

/// Read a big-endian `u16` at `offset`, returning `None` if it doesn't fit
/// within `data`.
fn read_u16(data: &[u8], offset: usize) -> Option<u16> {
    let end = offset.checked_add(2)?;
    let bytes: [u8; 2] = data.get(offset..end)?.try_into().ok()?;
    Some(u16::from_be_bytes(bytes))
}

/// Parse all TXT record payloads from a raw DNS response.
///
/// Returns a list of concatenated TXT string data (each TXT record may span
/// multiple strings internally; they are joined here).  Empty records are
/// omitted.
#[must_use]
pub fn parse_txt_response(data: &[u8]) -> Vec<Vec<u8>> {
    let mut results = Vec::new();
    if data.len() < 12 {
        return results;
    }

    let Some(ancount) = read_u16(data, 6) else {
        return results;
    };
    let mut offset = 12;

    // Skip the question section QNAME.
    offset = skip_name(data, offset);
    // Skip QTYPE + QCLASS (4 bytes).
    offset = offset.saturating_add(4);

    for _ in 0..ancount {
        offset = skip_name(data, offset);

        let Some(rtype) = read_u16(data, offset) else {
            break;
        };
        let Some(rdlength) = read_u16(data, offset.saturating_add(8)) else {
            break;
        };
        let rdlength = usize::from(rdlength);
        offset = offset.saturating_add(10);

        let Some(rdata) = data.get(offset..offset.saturating_add(rdlength)) else {
            break;
        };

        if rtype == 16 {
            // TXT RDATA is a sequence of length-prefixed strings.
            let mut txt_data = Vec::new();
            let mut i = 0;
            while let Some(&chunk_len_byte) = rdata.get(i) {
                let chunk_len = usize::from(chunk_len_byte);
                i = i.saturating_add(1);
                let Some(chunk) = rdata.get(i..i.saturating_add(chunk_len)) else {
                    break;
                };
                txt_data.extend_from_slice(chunk);
                i = i.saturating_add(chunk_len);
            }
            if !txt_data.is_empty() {
                results.push(txt_data);
            }
        }
        offset = offset.saturating_add(rdlength);
    }
    results
}

/// Parse all A and AAAA records from a raw DNS response.
///
/// Returns an empty list if the response is malformed or contains no IP records.
#[must_use]
pub fn parse_dns_response(data: &[u8]) -> Vec<IpAddr> {
    let mut ips = Vec::new();
    if data.len() < 12 {
        return ips;
    }

    let Some(ancount) = read_u16(data, 6) else {
        return ips;
    };
    let mut offset = 12;

    // Skip question QNAME, QTYPE, QCLASS.
    offset = skip_name(data, offset);
    offset = offset.saturating_add(4);

    for _ in 0..ancount {
        offset = skip_name(data, offset);

        let Some(atype) = read_u16(data, offset) else {
            break;
        };
        let Some(rdlength) = read_u16(data, offset.saturating_add(8)) else {
            break;
        };
        let rdlength = usize::from(rdlength);
        offset = offset.saturating_add(10);

        if atype == 1
            && rdlength == 4
            && let Some(Ok(bytes)) = data
                .get(offset..offset.saturating_add(4))
                .map(<[u8; 4]>::try_from)
        {
            ips.push(IpAddr::V4(Ipv4Addr::from(bytes)));
        } else if atype == 28
            && rdlength == 16
            && let Some(Ok(bytes)) = data
                .get(offset..offset.saturating_add(16))
                .map(<[u8; 16]>::try_from)
        {
            ips.push(IpAddr::V6(Ipv6Addr::from(bytes)));
        }

        let Some(next_offset) = offset
            .checked_add(rdlength)
            .filter(|&end| end <= data.len())
        else {
            break;
        };
        offset = next_offset;
    }
    ips
}

/// Advance `offset` past a DNS name (handles both label sequences and
/// pointer compression).
///
/// Returns the new offset positioned **after** the name.
fn skip_name(data: &[u8], mut offset: usize) -> usize {
    loop {
        let Some(&byte) = data.get(offset) else {
            return offset;
        };
        if (byte & 0xC0) == 0xC0 {
            // Compression pointer: 2 bytes total.
            return offset.saturating_add(2);
        }
        let label_len = usize::from(byte);
        if label_len == 0 {
            return offset.saturating_add(1);
        }
        offset = offset.saturating_add(1).saturating_add(label_len);
    }
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::cast_possible_truncation,
    clippy::indexing_slicing,
    clippy::arithmetic_side_effects,
    clippy::as_conversions
)]
mod tests {
    use super::*;

    #[test]
    fn test_validate_domain_accepts_normal_names() {
        assert!(validate_domain("github.com").is_ok());
        assert!(validate_domain("a.b.c.example.").is_ok());
    }

    #[test]
    fn test_validate_domain_rejects_empty() {
        assert!(validate_domain("").is_err());
    }

    #[test]
    fn test_validate_domain_rejects_oversized_label() {
        let label = "a".repeat(64);
        assert!(validate_domain(&format!("{label}.com")).is_err());

        // Exactly 63 bytes is still fine.
        let ok_label = "a".repeat(63);
        assert!(validate_domain(&format!("{ok_label}.com")).is_ok());
    }

    #[test]
    fn test_validate_domain_rejects_oversized_name() {
        // 4 labels of 63 bytes each plus separators exceeds the 255-byte
        // wire-format limit.
        let label = "a".repeat(63);
        let domain = [label.as_str(); 5].join(".");
        assert!(validate_domain(&domain).is_err());
    }

    #[test]
    fn test_parse_dns_response_aaaa_record() {
        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[6]);
        response.extend_from_slice(b"github");
        response.extend_from_slice(&[3]);
        response.extend_from_slice(b"com");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x1c, 0x00, 0x01]); // QTYPE = AAAA
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x1c, 0x00, 0x01]);
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        response.extend_from_slice(&[0x00, 0x10]); // RDLENGTH = 16
        let addr = Ipv6Addr::new(0x2606, 0x4700, 0x4700, 0, 0, 0, 0, 0x1111);
        response.extend_from_slice(&addr.octets());

        let parsed = parse_dns_response(&response);
        assert_eq!(parsed, vec![IpAddr::V6(addr)]);
    }

    #[test]
    fn test_build_txt_record_query_has_txt_qtype() {
        let query = build_txt_record_query("2.dnscrypt-cert.quad9.net", [0x11, 0x22]);
        // QTYPE (0x0010 = TXT) is the two bytes just before the trailing QCLASS.
        assert_eq!(&query[query.len() - 4..query.len() - 2], &[0x00, 0x10]);
    }

    #[test]
    fn test_parse_txt_response_extracts_chunks() {
        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[4]);
        response.extend_from_slice(b"test");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]); // QTYPE = TXT
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]); // TYPE = TXT
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        let txt_chunk = b"hello";
        response.extend_from_slice(&[0x00, (txt_chunk.len() + 1) as u8]); // RDLENGTH
        response.push(txt_chunk.len() as u8);
        response.extend_from_slice(txt_chunk);

        let parsed = parse_txt_response(&response);
        assert_eq!(parsed, vec![txt_chunk.to_vec()]);
    }

    #[test]
    fn test_build_query_skips_empty_labels() {
        // A leading/trailing/double dot must not produce a zero-length label
        // in the wire encoding (which would be indistinguishable from the
        // root-label terminator).
        let query = build_a_record_query(".github..com.", [0, 0]);
        assert_eq!(query[12], 6);
        assert_eq!(&query[13..19], b"github");
        assert_eq!(query[19], 3);
        assert_eq!(&query[20..23], b"com");
        assert_eq!(query[23], 0); // root label terminator directly follows
    }

    #[test]
    fn test_parse_txt_response_rejects_short_input() {
        assert!(parse_txt_response(&[]).is_empty());
        assert!(parse_txt_response(&[0; 10]).is_empty());
    }

    #[test]
    fn test_parse_txt_response_ignores_non_txt_records() {
        // An answer section with a non-TXT record (type A = 1) followed by a
        // real TXT record; only the TXT record should be extracted.
        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[4]);
        response.extend_from_slice(b"test");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        // Answer 1: A record (ignored by parse_txt_response).
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        response.extend_from_slice(&[0x00, 0x04]);
        response.extend_from_slice(&[1, 2, 3, 4]);
        // Answer 2: TXT record.
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        let txt_chunk = b"hi";
        response.extend_from_slice(&[0x00, (txt_chunk.len() + 1) as u8]);
        response.push(txt_chunk.len() as u8);
        response.extend_from_slice(txt_chunk);

        let parsed = parse_txt_response(&response);
        assert_eq!(parsed, vec![txt_chunk.to_vec()]);
    }

    #[test]
    fn test_parse_txt_response_truncated_answer_header() {
        // Valid question section, but the answer's fixed 10-byte header
        // (type/class/ttl/rdlength) is cut short.
        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[4]);
        response.extend_from_slice(b"test");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x10]); // truncated: only 2 of 10 header bytes
        assert!(parse_txt_response(&response).is_empty());
    }

    #[test]
    fn test_parse_txt_response_rdlength_exceeds_buffer() {
        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[4]);
        response.extend_from_slice(b"test");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        response.extend_from_slice(&[0xff, 0xff]); // RDLENGTH far exceeds remaining data
        assert!(parse_txt_response(&response).is_empty());
    }

    #[test]
    fn test_parse_txt_response_chunk_len_overflow() {
        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[4]);
        response.extend_from_slice(b"test");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x10, 0x00, 0x01]);
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        response.extend_from_slice(&[0x00, 0x02]); // RDLENGTH = 2
        response.extend_from_slice(&[0xff, 0x00]); // chunk_len=255 but only 1 byte follows
        assert!(parse_txt_response(&response).is_empty());
    }

    #[test]
    fn test_parse_dns_response_truncated_answer_and_rdlength_overflow() {
        assert!(parse_dns_response(&[]).is_empty());
        assert!(parse_dns_response(&[0; 10]).is_empty());

        let mut response = Vec::new();
        response.extend_from_slice(&[
            0xab, 0xcd, 0x81, 0x80, 0x00, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
        ]);
        response.extend_from_slice(&[6]);
        response.extend_from_slice(b"github");
        response.extend_from_slice(&[3]);
        response.extend_from_slice(b"com");
        response.extend_from_slice(&[0]);
        response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
        response.extend_from_slice(&[0xc0, 0x0c]);
        response.extend_from_slice(&[0x00, 0x01, 0x00, 0x01]);
        response.extend_from_slice(&[0x00, 0x00, 0x00, 0x3c]);
        response.extend_from_slice(&[0xff, 0xff]); // RDLENGTH far exceeds remaining data
        assert!(parse_dns_response(&response).is_empty());
    }

    #[test]
    fn test_skip_name_handles_truncated_label() {
        // A label length byte with no following data: offset runs past the
        // end of the buffer inside the loop rather than at entry.
        let data = [0u8; 5];
        let mut truncated = data.to_vec();
        truncated[0] = 100; // claims a 100-byte label in a 5-byte buffer (top bits != 0xC0, so not a pointer)
        assert_eq!(skip_name(&truncated, 0), 101);
        assert_eq!(skip_name(&truncated, 10), 10);
    }
}