dnscrypt 0.1.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, with a
/// minimum padded length of `min_len`. We also randomly append 0 to 2 extra
/// 64-byte blocks to obscure the exact size from traffic analysts.
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 mut rb = [0u8; 4];
    getrandom::fill(&mut rb).expect("getrandom failed");
    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())
}

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

/// Build a minimal DNS query for `domain` requesting TXT records (type 16).
///
/// The transaction ID should be randomized.
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.
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.
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(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.
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.
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;
    }
}