dnscrypt 0.1.1

//! 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 = if current.is_multiple_of(64) {
        current
    } else {
        (current / 64 + 1) * 64
    };

    // Add 0, 1, or 2 extra 64-byte blocks of random padding.
    let rb: [u8; 4] = crate::rng::random_bytes();
    let extra_blocks = (u32::from_ne_bytes(rb) % 3) as usize;
    let target = base_target + (extra_blocks * 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> {
    let mut end = decrypted.len();
    while end > 0 && decrypted[end - 1] == 0x00 {
        end -= 1;
    }
    if end == 0 || decrypted[end - 1] != 0x80 {
        return Err("Invalid ISO/IEC 7816-4 padding".to_string());
    }
    Ok(decrypted[..end - 1].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 += 1 + 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);
        // Safe: `min(63)` bounds `len` to the `u8` range.
        #[allow(clippy::cast_possible_truncation)]
        query.push(len as u8);
        query.extend_from_slice(&label.as_bytes()[..len]);
    }
    query.push(0x00); // root label
    query.extend_from_slice(&qtype.to_be_bytes());
    query.extend_from_slice(&[0x00, 0x01]); // QCLASS = IN
    query
}

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

/// 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 ancount = u16::from_be_bytes([data[6], data[7]]) as usize;
    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);

        if offset + 10 > data.len() {
            break;
        }
        let rtype = u16::from_be_bytes([data[offset], data[offset + 1]]);
        let rdlength = u16::from_be_bytes([data[offset + 8], data[offset + 9]]) as usize;
        offset += 10;

        if offset + rdlength > data.len() {
            break;
        }

        if rtype == 16 {
            // TXT RDATA is a sequence of length-prefixed strings.
            let mut txt_data = Vec::new();
            let rdata = &data[offset..offset + rdlength];
            let mut i = 0;
            while i < rdata.len() {
                let chunk_len = rdata[i] as usize;
                i += 1;
                if i + chunk_len > rdata.len() {
                    break;
                }
                txt_data.extend_from_slice(&rdata[i..i + chunk_len]);
                i += chunk_len;
            }
            if !txt_data.is_empty() {
                results.push(txt_data);
            }
        }
        offset += 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 ancount = u16::from_be_bytes([data[6], data[7]]) as usize;
    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);

        if offset + 10 > data.len() {
            break;
        }
        let atype = u16::from_be_bytes([data[offset], data[offset + 1]]);
        let rdlength = u16::from_be_bytes([data[offset + 8], data[offset + 9]]) as usize;
        offset += 10;

        if atype == 1 && rdlength == 4 && offset + 4 <= data.len() {
            ips.push(IpAddr::V4(Ipv4Addr::new(
                data[offset],
                data[offset + 1],
                data[offset + 2],
                data[offset + 3],
            )));
        } else if atype == 28 && rdlength == 16 && offset + 16 <= data.len() {
            let mut addr = [0u8; 16];
            addr.copy_from_slice(&data[offset..offset + 16]);
            ips.push(IpAddr::V6(Ipv6Addr::from(addr)));
        }

        if offset
            .checked_add(rdlength)
            .is_none_or(|end| end > data.len())
        {
            break;
        }
        offset += rdlength;
    }
    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 {
        if offset >= data.len() {
            return offset;
        }
        if (data[offset] & 0xC0) == 0xC0 {
            // Compression pointer: 2 bytes total.
            return offset + 2;
        }
        let label_len = data[offset] as usize;
        if label_len == 0 {
            return offset + 1;
        }
        offset += 1 + label_len;
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::cast_possible_truncation)]
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);
    }
}