etchdns 0.1.2

use byteorder::{BigEndian, ByteOrder};
use log::debug;
use std::str;

use crate::errors::{DnsError, DnsResult};

// DNS packet constants
pub const DNS_PACKET_LEN_MIN: usize = 12; // Minimum size of a DNS packet (header only)
pub const DNS_HEADER_SIZE: usize = DNS_PACKET_LEN_MIN; // Size of the DNS header
pub const DNS_MAX_PACKET_SIZE: usize = 0x1600; // Maximum size of a DNS packet
pub const DNS_MAX_HOSTNAME_SIZE: usize = 256; // Maximum size of a hostname
pub const DNS_OFFSET_QUESTION: usize = DNS_HEADER_SIZE; // Offset to the question section
pub const DNS_MAX_UDP_PACKET_SIZE: usize = 512; // Standard maximum UDP packet size (RFC 1035)

// DNS record types
#[allow(dead_code)]
pub const DNS_TYPE_A: u16 = 1;
pub const DNS_TYPE_NS: u16 = 2;
pub const DNS_TYPE_CNAME: u16 = 5;
#[allow(dead_code)]
pub const DNS_TYPE_SOA: u16 = 6;
pub const DNS_TYPE_PTR: u16 = 12;
pub const DNS_TYPE_MX: u16 = 15;
#[allow(dead_code)]
pub const DNS_TYPE_TXT: u16 = 16;
#[allow(dead_code)]
pub const DNS_TYPE_AAAA: u16 = 28;
pub const DNS_TYPE_SRV: u16 = 33;
pub const DNS_TYPE_OPT: u16 = 41;
pub const DNS_TYPE_ANY: u16 = 255;

// DNS classes
pub const DNS_CLASS_IN: u16 = 1;
pub const DNS_CLASS_ANY: u16 = 255;

// DNSSEC related constants
pub const DNS_FLAG_DO: u16 = 0x8000; // DNSSEC OK flag in OPT record

// EDNS option codes
pub const EDNS_OPTION_CLIENT_SUBNET: u16 = 8; // Client Subnet (RFC 7871)

// DNS header flags
const DNS_FLAGS_QR: u16 = 1u16 << 15; // Query/Response flag
const DNS_FLAGS_AA: u16 = 1u16 << 10; // Authoritative Answer flag
const DNS_FLAGS_TC: u16 = 1u16 << 9; // Truncation flag
#[allow(dead_code)]
const DNS_FLAGS_RD: u16 = 1u16 << 8; // Recursion Desired
const DNS_FLAGS_RA: u16 = 1u16 << 7; // Recursion Available flag
const DNS_FLAGS_CD: u16 = 1u16 << 4; // Checking Disabled (DNSSEC)
#[allow(dead_code)]
const DNS_FLAGS_AD: u16 = 1u16 << 5; // Authentic Data (DNSSEC)

// DNS opcodes (in bits 11-14 of the flags field)
pub const DNS_OPCODE_QUERY: u8 = 0; // Standard query
#[allow(dead_code)]
pub const DNS_OPCODE_IQUERY: u8 = 1; // Inverse query
#[allow(dead_code)]
pub const DNS_OPCODE_STATUS: u8 = 2; // Server status request
#[allow(dead_code)]
pub const DNS_OPCODE_NOTIFY: u8 = 4; // Zone change notification
#[allow(dead_code)]
pub const DNS_OPCODE_UPDATE: u8 = 5; // Dynamic update

/// Validates that a DNS packet looks valid
pub fn validate_dns_packet(packet: &[u8]) -> DnsResult<()> {
    // Check minimum packet size
    if packet.len() < DNS_HEADER_SIZE {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // Check maximum packet size
    if packet.len() > DNS_MAX_PACKET_SIZE {
        return Err(DnsError::PacketTooLarge {
            size: packet.len(),
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }

    // Check opcode - we only support standard queries (opcode 0)
    if !is_response(packet) {
        let op = opcode(packet);
        if op != DNS_OPCODE_QUERY {
            return Err(DnsError::UnsupportedOperation(format!(
                "Unsupported opcode: {op}"
            )));
        }
    }

    // Check that we have at least one question
    let qd_count = qdcount(packet);
    if qd_count == 0 {
        return Err(DnsError::InvalidPacket(
            "No questions in packet".to_string(),
        ));
    }

    // For standard DNS queries, we typically expect exactly one question
    if qd_count > 1 {
        return Err(DnsError::InvalidPacket(
            "Too many questions in packet".to_string(),
        ));
    }

    // Validate the counts make sense for the packet size
    let (an_count, ns_count, ar_count) = (ancount(packet), nscount(packet), arcount(packet));

    // Try to parse the question section
    let mut offset = DNS_HEADER_SIZE;

    // Validate each question section
    for _ in 0..qd_count {
        // Validate the question name
        offset = match skip_name(packet, offset) {
            Ok(new_offset) => new_offset,
            Err(e) => {
                return Err(DnsError::InvalidQuestion(format!(
                    "Invalid question name: {e}"
                )));
            }
        };

        // Ensure there's enough space for QTYPE and QCLASS (4 bytes)
        if packet.len() < offset + 4 {
            return Err(DnsError::InvalidPacket(
                "Packet too short for QTYPE and QCLASS".to_string(),
            ));
        }

        // Read QTYPE and QCLASS
        let qtype_end = offset.checked_add(2).ok_or_else(|| {
            DnsError::InvalidPacket("Integer overflow calculating QTYPE offset".to_string())
        })?;
        let qclass_end = offset.checked_add(4).ok_or_else(|| {
            DnsError::InvalidPacket("Integer overflow calculating QCLASS offset".to_string())
        })?;
        let _qtype = BigEndian::read_u16(&packet[offset..qtype_end]);
        let qclass = BigEndian::read_u16(&packet[qtype_end..qclass_end]);

        // Validate QCLASS (typically IN=1 or ANY=255)
        if qclass != DNS_CLASS_IN && qclass != DNS_CLASS_ANY {
            debug!("Unusual: DNS query with QCLASS={qclass}");
        }

        // Additional validation: extract the query name and validate it
        match qname(packet) {
            Ok(name) => {
                if name.len() > DNS_MAX_HOSTNAME_SIZE {
                    return Err(DnsError::DomainNameTooLong {
                        length: name.len(),
                        max_length: DNS_MAX_HOSTNAME_SIZE,
                    });
                }

                // Validate domain name characters
                let labels: Vec<&[u8]> = name.split(|&c| c == b'.').collect();
                for label in labels {
                    if label.is_empty() {
                        // Skip empty labels (like the root label) in normal domain names
                        // But if it's between other labels (like "example..com"), that's an error
                        if !name.starts_with(b".") && !name.ends_with(b".") && name.len() > 1 {
                            return Err(DnsError::EmptyLabel);
                        }
                        continue;
                    }

                    // Check label length
                    if label.len() > 63 {
                        return Err(DnsError::LabelTooLong {
                            length: label.len(),
                        });
                    }

                    // Validate each character in the label
                    for (i, &c) in label.iter().enumerate() {
                        if !is_valid_domain_name_char(c, i, label.len()) {
                            let label_str = match std::str::from_utf8(label) {
                                Ok(s) => s.to_string(),
                                Err(_) => format!("{label:?}"),
                            };

                            return Err(DnsError::InvalidDomainNameCharacter {
                                character: c as char,
                                position: i,
                                label: label_str,
                            });
                        }
                    }
                }
            }
            Err(e) => {
                return Err(DnsError::InvalidQuestion(format!(
                    "Failed to extract query name: {e}"
                )));
            }
        }

        offset = offset.checked_add(4).ok_or_else(|| {
            DnsError::InvalidPacket("Integer overflow advancing past QTYPE/QCLASS".to_string())
        })?;
    }

    // Try to validate the answer, authority, and additional sections
    // Check for potential overflow when calculating total_rr_count
    let an_count_usize = an_count as usize;
    let ns_count_usize = ns_count as usize;
    let ar_count_usize = ar_count as usize;

    // Check for overflow when adding up the counts
    let total_rr_count = match an_count_usize.checked_add(ns_count_usize) {
        Some(count) => match count.checked_add(ar_count_usize) {
            Some(total) => total,
            None => {
                return Err(DnsError::InvalidPacket(
                    "Integer overflow in resource record count calculation".to_string(),
                ));
            }
        },
        None => {
            return Err(DnsError::InvalidPacket(
                "Integer overflow in resource record count calculation".to_string(),
            ));
        }
    };

    // Validate against maximum sensible number of records
    const MAX_SENSIBLE_RR_COUNT: usize = 1000;
    if total_rr_count > MAX_SENSIBLE_RR_COUNT {
        return Err(DnsError::InvalidPacket(format!(
            "Excessive resource record count: {total_rr_count} (max {MAX_SENSIBLE_RR_COUNT})"
        )));
    }

    if total_rr_count > 0 {
        match traverse_rrs(packet, offset, total_rr_count, |rr_offset| {
            // Validate record type
            // Make sure we have at least 2 bytes to read at rr_offset
            if packet.len() < rr_offset + 2 {
                return Err(DnsError::PacketTooShort { offset: rr_offset });
            }
            let _rr_type = BigEndian::read_u16(&packet[rr_offset..rr_offset + 2]);

            // Validate TTL - make sure we have at least 8 bytes at rr_offset
            if packet.len() < rr_offset + 8 {
                return Err(DnsError::PacketTooShort { offset: rr_offset });
            }
            let ttl = BigEndian::read_u32(&packet[rr_offset + 4..rr_offset + 8]);
            if ttl > 604800 {
                // 7 days in seconds
                debug!("Unusual: DNS record with TTL > 7 days: {ttl} seconds");
            }

            // Validate RDLENGTH - make sure we have at least 10 bytes at rr_offset
            if packet.len() < rr_offset + 10 {
                return Err(DnsError::PacketTooShort { offset: rr_offset });
            }
            let rdlength = BigEndian::read_u16(&packet[rr_offset + 8..rr_offset + 10]) as usize;

            // Check for potential overflow when calculating rr_offset + 10 + rdlength
            let rr_data_end = match rr_offset.checked_add(10) {
                Some(offset_plus_10) => match offset_plus_10.checked_add(rdlength) {
                    Some(end_offset) => end_offset,
                    None => {
                        return Err(DnsError::InvalidRecord(
                            "Integer overflow calculating record data bounds".to_string(),
                        ));
                    }
                },
                None => {
                    return Err(DnsError::InvalidRecord(
                        "Integer overflow calculating record data bounds".to_string(),
                    ));
                }
            };

            if packet.len() < rr_data_end {
                return Err(DnsError::InvalidRecord(format!(
                    "Record data length ({rdlength}) exceeds packet bounds"
                )));
            }

            Ok(())
        }) {
            Ok(_) => (),
            Err(e) => {
                return Err(DnsError::InvalidRecord(format!(
                    "Invalid resource record: {e}"
                )));
            }
        }
    }

    // Check EDNS version if OPT record is present
    if let Some(edns_version) = extract_edns_version(packet)? {
        if edns_version != 0 {
            return Err(DnsError::UnsupportedEdnsVersion {
                version: edns_version,
            });
        }
    }

    Ok(())
}

/// Returns the transaction ID from the DNS packet
#[inline]
pub fn tid(packet: &[u8]) -> u16 {
    if packet.len() < 2 {
        return 0; // Return 0 for malformed packets
    }
    BigEndian::read_u16(&packet[0..2])
}

/// Returns the number of questions in the DNS packet
#[inline]
pub fn qdcount(packet: &[u8]) -> u16 {
    if packet.len() < 6 {
        return 0; // Return 0 for malformed packets
    }
    BigEndian::read_u16(&packet[4..6])
}

/// Returns the number of answer records in the DNS packet
#[inline]
pub fn ancount(packet: &[u8]) -> u16 {
    if packet.len() < 8 {
        return 0; // Return 0 for malformed packets
    }
    BigEndian::read_u16(&packet[6..8])
}

/// Returns the number of authority records in the DNS packet
#[inline]
pub fn nscount(packet: &[u8]) -> u16 {
    if packet.len() < 10 {
        return 0; // Return 0 for malformed packets
    }
    BigEndian::read_u16(&packet[8..10])
}

/// Returns the number of additional records in the DNS packet
#[inline]
pub fn arcount(packet: &[u8]) -> u16 {
    if packet.len() < DNS_PACKET_LEN_MIN {
        return 0; // Return 0 for malformed packets
    }
    BigEndian::read_u16(&packet[10..DNS_PACKET_LEN_MIN])
}

/// Returns the response code from the DNS packet
#[inline]
pub fn rcode(packet: &[u8]) -> u8 {
    if packet.len() < 4 {
        return 2; // Return SERVFAIL (2) for truncated packets
    }
    packet[3] & 0x0F
}

/// Checks if the packet is a response
#[inline]
pub fn is_response(packet: &[u8]) -> bool {
    if packet.len() < 4 {
        return false; // Not a response if packet is too short
    }
    (BigEndian::read_u16(&packet[2..4]) & DNS_FLAGS_QR) == DNS_FLAGS_QR
}

/// Returns the opcode from the DNS packet
#[inline]
pub fn opcode(packet: &[u8]) -> u8 {
    if packet.len() < 3 {
        return 0; // Return standard query (0) for malformed packets
    }
    (packet[2] >> 3) & 0x0F
}

/// Checks if the packet is a standard query
#[inline]
#[allow(dead_code)]
pub fn is_standard_query(packet: &[u8]) -> bool {
    !is_response(packet) && opcode(packet) == DNS_OPCODE_QUERY
}

/// Checks if the packet has the TC (truncated) bit set
#[inline]
pub fn is_truncated(packet: &[u8]) -> bool {
    if packet.len() < 4 {
        return false; // Not truncated if packet is too short
    }
    (BigEndian::read_u16(&packet[2..4]) & DNS_FLAGS_TC) == DNS_FLAGS_TC
}

/// Checks if DNSSEC is requested (either DO bit set or CD bit set)
pub fn is_dnssec_requested(packet: &[u8]) -> DnsResult<bool> {
    // Check CD bit in header flags
    let flags = BigEndian::read_u16(&packet[2..4]);
    if (flags & DNS_FLAGS_CD) == DNS_FLAGS_CD {
        return Ok(true);
    }

    // Check for OPT record with DO bit set
    let packet_len = packet.len();
    if packet_len <= DNS_OFFSET_QUESTION {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    if packet_len.saturating_sub(offset) < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    offset = offset.checked_add(4).ok_or_else(|| {
        DnsError::InvalidPacket("Integer overflow skipping QTYPE/QCLASS".to_string())
    })?; // Skip QTYPE and QCLASS

    // Skip answer and authority sections
    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    offset = traverse_rrs(packet, offset, ancount as usize + nscount as usize, |_| {
        Ok(())
    })?;

    // Check additional records for OPT
    let mut dnssec_requested = false;
    traverse_rrs(packet, offset, arcount as usize, |offset| {
        let qtype = BigEndian::read_u16(&packet[offset..]);
        if qtype == DNS_TYPE_OPT {
            // OPT record found, check DO bit in flags
            let opt_flags = BigEndian::read_u16(&packet[offset + 4..]);
            if (opt_flags & DNS_FLAG_DO) == DNS_FLAG_DO {
                dnssec_requested = true;
            }
        }
        Ok(())
    })?;

    Ok(dnssec_requested)
}

/// Extracts the query name from a DNS packet
#[allow(dead_code)]
pub fn query_name(packet: &[u8]) -> DnsResult<String> {
    let qname_bytes = qname(packet)?;

    // Validate maximum length before conversion
    if qname_bytes.len() > DNS_MAX_HOSTNAME_SIZE {
        return Err(DnsError::DomainNameTooLong {
            length: qname_bytes.len(),
            max_length: DNS_MAX_HOSTNAME_SIZE,
        });
    }

    let qname_str = match str::from_utf8(&qname_bytes) {
        Ok(s) => s,
        Err(_) => {
            return Err(DnsError::InvalidDomainName(
                "Invalid UTF-8 in query name".to_string(),
            ));
        }
    };

    // Check once more for DoS using large string allocations
    if qname_str.len() > DNS_MAX_HOSTNAME_SIZE {
        return Err(DnsError::DomainNameTooLong {
            length: qname_str.len(),
            max_length: DNS_MAX_HOSTNAME_SIZE,
        });
    }

    Ok(qname_str.to_string())
}

/// Extracts the query type and class from a DNS packet
pub fn query_type_class(packet: &[u8]) -> DnsResult<(u16, u16)> {
    if packet.len() < DNS_HEADER_SIZE {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    if qdcount(packet) == 0 {
        return Err(DnsError::InvalidPacket("No question".to_string()));
    }

    let offset = skip_name(packet, DNS_HEADER_SIZE)?;
    if packet.len() - offset < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    let qtype = BigEndian::read_u16(&packet[offset..]);
    let qclass = BigEndian::read_u16(&packet[offset + 2..]);

    Ok((qtype, qclass))
}

/// Extracts the minimum TTL from all records in a DNS packet
pub fn extract_min_ttl(packet: &[u8]) -> DnsResult<Option<u32>> {
    let packet_len = packet.len();
    if packet_len <= DNS_OFFSET_QUESTION {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    if packet_len > DNS_MAX_PACKET_SIZE {
        return Err(DnsError::PacketTooLarge {
            size: packet_len,
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }
    if qdcount(packet) == 0 {
        return Err(DnsError::InvalidPacket("No question".to_string()));
    }

    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    if packet_len - offset < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    offset += 4; // Skip QTYPE and QCLASS

    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    let total_rr_count = ancount as usize + nscount as usize + arcount as usize;

    if total_rr_count == 0 {
        return Ok(None);
    }

    let mut min_ttl: Option<u32> = None;

    traverse_rrs(packet, offset, total_rr_count, |offset| {
        let rr_type = BigEndian::read_u16(&packet[offset..offset + 2]);
        let ttl = BigEndian::read_u32(&packet[offset + 4..offset + 8]);

        // Ignore OPT records (type 41) when calculating minimum TTL
        if rr_type != DNS_TYPE_OPT {
            min_ttl = match min_ttl {
                None => Some(ttl),
                Some(current_min) => Some(std::cmp::min(current_min, ttl)),
            };
        }
        Ok(())
    })?;

    Ok(min_ttl)
}

/// Changes the TTL of all records in a DNS packet
pub fn change_ttl(packet: &mut [u8], new_ttl: u32) -> DnsResult<()> {
    let packet_len = packet.len();
    if packet_len <= DNS_OFFSET_QUESTION {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    if packet_len > DNS_MAX_PACKET_SIZE {
        return Err(DnsError::PacketTooLarge {
            size: packet_len,
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }
    if qdcount(packet) == 0 {
        return Err(DnsError::InvalidPacket("No question".to_string()));
    }

    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    if packet_len - offset < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    offset += 4; // Skip QTYPE and QCLASS

    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    let total_rr_count = ancount as usize + nscount as usize + arcount as usize;

    if total_rr_count == 0 {
        return Ok(());
    }

    traverse_rrs_mut(packet, offset, total_rr_count, |packet, offset| {
        // Update TTL
        BigEndian::write_u32(&mut packet[offset + 4..offset + 8], new_ttl);
        Ok(())
    })?;

    Ok(())
}

// Helper functions

/// Validates a single character in a domain name label
///
/// According to RFC 1035 and RFC 1123, domain name labels must:
/// - Start with a letter or digit
/// - End with a letter or digit
/// - Contain only letters, digits, and hyphens in between
fn is_valid_domain_name_char(c: u8, position: usize, label_len: usize) -> bool {
    // Check if the character is a letter, digit, or hyphen
    let is_letter = c.is_ascii_lowercase() || c.is_ascii_uppercase();
    let is_digit = c.is_ascii_digit();
    let is_hyphen = c == b'-';

    // First and last characters must be letters or digits
    if position == 0 || position == label_len - 1 {
        return is_letter || is_digit;
    }

    // Middle characters can be letters, digits, or hyphens
    is_letter || is_digit || is_hyphen
}

// This function is now only used in validate_dns_packet

/// Extracts the raw query name from a DNS packet
pub fn qname(packet: &[u8]) -> DnsResult<Vec<u8>> {
    let packet_len = packet.len();
    if packet_len <= DNS_OFFSET_QUESTION {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    if qdcount(packet) == 0 {
        return Err(DnsError::InvalidPacket("No question".to_string()));
    }

    let mut offset = DNS_HEADER_SIZE;
    let mut qname = Vec::with_capacity(DNS_MAX_HOSTNAME_SIZE);

    loop {
        if offset >= packet_len {
            return Err(DnsError::PacketTooShort { offset: 0 });
        }
        match packet[offset] as usize {
            label_len if label_len & 0xc0 == 0xc0 => {
                // Check if we have enough bytes for the compression pointer
                if offset + 1 >= packet_len {
                    return Err(DnsError::PacketTooShort { offset });
                }

                // Compression pointer - not expected in the query name
                let second_byte = packet[offset + 1] as usize;
                return Err(DnsError::InvalidCompressionPointer {
                    offset,
                    pointer: ((label_len & 0x3f) << 8) | second_byte,
                    packet_size: packet_len,
                });
            }
            0 => {
                if qname.is_empty() {
                    qname.push(b'.');
                }
                break;
            }
            label_len => {
                if label_len >= 0x40 {
                    return Err(DnsError::LabelTooLong { length: label_len });
                }
                if packet_len.saturating_sub(offset) <= 1 {
                    return Err(DnsError::PacketTooShort { offset: 0 });
                }
                offset = offset.checked_add(1).ok_or_else(|| {
                    DnsError::InvalidPacket(
                        "Integer overflow advancing past label length byte".to_string(),
                    )
                })?;
                if packet_len.saturating_sub(offset) < label_len {
                    return Err(DnsError::PacketTooShort { offset: 0 });
                }

                if !qname.is_empty() {
                    qname.push(b'.');
                }

                if qname.len() + label_len >= DNS_MAX_HOSTNAME_SIZE {
                    return Err(DnsError::DomainNameTooLong {
                        length: qname.len() + label_len,
                        max_length: DNS_MAX_HOSTNAME_SIZE,
                    });
                }

                let label_end = offset.checked_add(label_len).ok_or_else(|| {
                    DnsError::InvalidPacket("Integer overflow calculating label end".to_string())
                })?;
                qname.extend_from_slice(&packet[offset..label_end]);
                offset = label_end;
            }
        }
    }

    Ok(qname)
}

/// Skips over a domain name in a DNS packet
fn skip_name(packet: &[u8], offset: usize) -> DnsResult<usize> {
    let packet_len = packet.len();
    if offset >= packet_len {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    let mut current_offset = offset;
    let mut jumps = 0;
    let mut total_name_len: usize = 0;

    loop {
        if current_offset >= packet_len {
            return Err(DnsError::PacketTooShort {
                offset: current_offset,
            });
        }
        match packet[current_offset] as usize {
            label_len if label_len & 0xc0 == 0xc0 => {
                // Compression pointer
                if packet_len - current_offset < 2 {
                    return Err(DnsError::PacketTooShort {
                        offset: current_offset,
                    });
                }
                jumps += 1;
                if jumps > 10 {
                    return Err(DnsError::TooManyCompressionPointers { count: jumps });
                }

                // For the return value, we only advance 2 bytes for the pointer
                if jumps == 1 {
                    // Use checked_add to prevent integer overflow
                    let result_offset = match current_offset.checked_add(2) {
                        Some(offset) => offset,
                        None => {
                            return Err(DnsError::InvalidCompressionPointer {
                                offset: current_offset,
                                pointer: 0,
                                packet_size: packet_len,
                            });
                        }
                    };

                    if result_offset > packet_len {
                        return Err(DnsError::PacketTooShort {
                            offset: current_offset,
                        });
                    }
                    current_offset = result_offset;
                    break;
                }

                // Follow the pointer - safely extract second byte
                if current_offset + 1 >= packet_len {
                    return Err(DnsError::PacketTooShort {
                        offset: current_offset,
                    });
                }

                // Calculate pointer value with overflow check
                let second_byte = packet[current_offset + 1] as usize;
                let first_part = (label_len & 0x3f) << 8;
                let pointer = match first_part.checked_add(second_byte) {
                    Some(ptr) => ptr,
                    None => {
                        return Err(DnsError::InvalidCompressionPointer {
                            offset: current_offset,
                            pointer: first_part,
                            packet_size: packet_len,
                        });
                    }
                };
                if pointer >= packet_len {
                    return Err(DnsError::InvalidCompressionPointer {
                        offset: current_offset,
                        pointer,
                        packet_size: packet_len,
                    });
                }
                current_offset = pointer;
            }
            0 => {
                // End of domain name
                current_offset = match current_offset.checked_add(1) {
                    Some(offset) => offset,
                    None => {
                        return Err(DnsError::PacketTooShort {
                            offset: current_offset,
                        });
                    }
                };
                break;
            }
            label_len => {
                // Regular label
                if label_len >= 0x40 {
                    return Err(DnsError::LabelTooLong { length: label_len });
                }

                // Make sure we don't have integer underflow in the packet_len - current_offset check
                let remaining_bytes = if current_offset <= packet_len {
                    packet_len - current_offset
                } else {
                    return Err(DnsError::PacketTooShort {
                        offset: current_offset,
                    });
                };

                // Check we have enough bytes for label_len
                if remaining_bytes <= label_len {
                    return Err(DnsError::PacketTooShort {
                        offset: current_offset,
                    });
                }

                // Update the total name length with overflow check
                // Add 1 for the dot between labels (or the length byte in the packet)
                total_name_len = match label_len.checked_add(1) {
                    Some(len) => match total_name_len.checked_add(len) {
                        Some(total) => total,
                        None => {
                            return Err(DnsError::DomainNameTooLong {
                                length: u32::MAX as usize,
                                max_length: DNS_MAX_HOSTNAME_SIZE,
                            });
                        }
                    },
                    None => {
                        return Err(DnsError::DomainNameTooLong {
                            length: u32::MAX as usize,
                            max_length: DNS_MAX_HOSTNAME_SIZE,
                        });
                    }
                };

                // Check if the total name length exceeds the maximum
                if total_name_len > DNS_MAX_HOSTNAME_SIZE {
                    return Err(DnsError::DomainNameTooLong {
                        length: total_name_len,
                        max_length: DNS_MAX_HOSTNAME_SIZE,
                    });
                }

                // Safely advance the offset
                let label_plus_1 = match label_len.checked_add(1) {
                    Some(len) => len,
                    None => {
                        return Err(DnsError::LabelTooLong {
                            length: u32::MAX as usize,
                        });
                    }
                };

                current_offset = match current_offset.checked_add(label_plus_1) {
                    Some(offset) => offset,
                    None => {
                        return Err(DnsError::PacketTooShort {
                            offset: current_offset,
                        });
                    }
                };
            }
        }
    }

    Ok(current_offset)
}

/// Traverses resource records in a DNS packet
/// Traverses resource records in a DNS packet and calls a callback for each record
///
/// This function iterates through a specified number of resource records in a DNS packet,
/// calling the provided callback function for each record. It handles domain name compression
/// and validates record boundaries.
///
/// # Arguments
///
/// * `packet` - The DNS packet data
/// * `offset` - Starting offset in the packet where the resource records begin
/// * `rrcount` - Number of resource records to traverse
/// * `callback` - Function to call for each resource record, receiving the offset of the record
///
/// # Returns
///
/// The final offset after traversing all records, or an error if packet validation fails
fn traverse_rrs<F>(
    packet: &[u8],
    mut offset: usize,
    rrcount: usize,
    mut callback: F,
) -> DnsResult<usize>
where
    F: FnMut(usize) -> DnsResult<()>,
{
    let packet_len = packet.len();

    for _ in 0..rrcount {
        // Skip and validate the domain name in the record
        // skip_name already includes hostname size validation
        offset = skip_name(packet, offset)?;

        // Check if we have enough bytes for the fixed RR header fields (10 bytes)
        if packet_len < offset || packet_len - offset < 10 {
            return Err(DnsError::PacketTooShort { offset });
        }

        callback(offset)?;

        // Safely read RDLEN
        if offset + 8 >= packet_len {
            return Err(DnsError::PacketTooShort { offset });
        }
        let rdlen = BigEndian::read_u16(&packet[offset + 8..offset + 10]) as usize;

        // Check for overflow when calculating rr_offset + 10 + rdlen
        let rdata_end = match offset.checked_add(10) {
            Some(offset_plus_10) => match offset_plus_10.checked_add(rdlen) {
                Some(end_offset) => end_offset,
                None => {
                    return Err(DnsError::InvalidRecord(
                        "Integer overflow calculating record data bounds".to_string(),
                    ));
                }
            },
            None => {
                return Err(DnsError::InvalidRecord(
                    "Integer overflow calculating record offset".to_string(),
                ));
            }
        };

        if packet_len < rdata_end {
            return Err(DnsError::InvalidRecord(
                "Record length would exceed packet length".to_string(),
            ));
        }

        // For certain record types that contain domain names in their RDATA,
        // we should validate those domain names too
        let rr_type = BigEndian::read_u16(&packet[offset..offset + 2]);
        match rr_type {
            DNS_TYPE_NS | DNS_TYPE_CNAME | DNS_TYPE_PTR | DNS_TYPE_MX | DNS_TYPE_SRV => {
                // These record types contain domain names in their RDATA
                // Check that the RDATA size is reasonable
                if rdlen > DNS_MAX_HOSTNAME_SIZE {
                    debug!("Large RDATA for domain name record: {rdlen} bytes (type: {rr_type})");
                }
            }
            _ => {}
        }

        // Safely advance the offset
        offset = match offset.checked_add(10) {
            Some(o) => match o.checked_add(rdlen) {
                Some(final_offset) => final_offset,
                None => {
                    return Err(DnsError::InvalidRecord(
                        "Integer overflow calculating next record offset".to_string(),
                    ));
                }
            },
            None => {
                return Err(DnsError::InvalidRecord(
                    "Integer overflow calculating next record offset".to_string(),
                ));
            }
        };
    }

    Ok(offset)
}

/// Traverses resource records in a mutable DNS packet and calls a callback for each record
///
/// This function iterates through a specified number of resource records in a mutable DNS packet,
/// calling the provided callback function for each record. It handles domain name compression
/// and validates record boundaries. Unlike `traverse_rrs`, this function allows modification
/// of the packet data.
///
/// # Arguments
///
/// * `packet` - The mutable DNS packet data
/// * `offset` - Starting offset in the packet where the resource records begin
/// * `rrcount` - Number of resource records to traverse
/// * `callback` - Function to call for each resource record, receiving the packet and offset of the record
///
/// # Returns
///
/// The final offset after traversing all records, or an error if packet validation fails
fn traverse_rrs_mut<F>(
    packet: &mut [u8],
    mut offset: usize,
    rrcount: usize,
    mut callback: F,
) -> DnsResult<usize>
where
    F: FnMut(&mut [u8], usize) -> DnsResult<()>,
{
    let packet_len = packet.len();

    for _ in 0..rrcount {
        // Skip and validate the domain name in the record
        // skip_name already includes hostname size validation
        offset = skip_name(packet, offset)?;

        // Check if we have enough bytes for the fixed RR header fields (10 bytes)
        if packet_len < offset || packet_len - offset < 10 {
            return Err(DnsError::PacketTooShort { offset });
        }

        callback(packet, offset)?;

        // Safely read RDLEN
        if offset + 8 >= packet_len {
            return Err(DnsError::PacketTooShort { offset });
        }
        let rdlen = BigEndian::read_u16(&packet[offset + 8..offset + 10]) as usize;

        // Check for overflow when calculating rr_offset + 10 + rdlen
        let rdata_end = match offset.checked_add(10) {
            Some(offset_plus_10) => match offset_plus_10.checked_add(rdlen) {
                Some(end_offset) => end_offset,
                None => {
                    return Err(DnsError::InvalidRecord(
                        "Integer overflow calculating record data bounds".to_string(),
                    ));
                }
            },
            None => {
                return Err(DnsError::InvalidRecord(
                    "Integer overflow calculating record offset".to_string(),
                ));
            }
        };

        if packet_len < rdata_end {
            return Err(DnsError::InvalidRecord(
                "Record length would exceed packet length".to_string(),
            ));
        }

        // For certain record types that contain domain names in their RDATA,
        // we should validate those domain names too
        let rr_type = BigEndian::read_u16(&packet[offset..offset + 2]);
        match rr_type {
            DNS_TYPE_NS | DNS_TYPE_CNAME | DNS_TYPE_PTR | DNS_TYPE_MX | DNS_TYPE_SRV => {
                // These record types contain domain names in their RDATA
                // Check that the RDATA size is reasonable
                if rdlen > DNS_MAX_HOSTNAME_SIZE {
                    debug!("Large RDATA for domain name record: {rdlen} bytes (type: {rr_type})");
                }
            }
            _ => {}
        }

        // Safely advance the offset
        offset = match offset.checked_add(10) {
            Some(o) => match o.checked_add(rdlen) {
                Some(final_offset) => final_offset,
                None => {
                    return Err(DnsError::InvalidRecord(
                        "Integer overflow calculating next record offset".to_string(),
                    ));
                }
            },
            None => {
                return Err(DnsError::InvalidRecord(
                    "Integer overflow calculating next record offset".to_string(),
                ));
            }
        };
    }

    Ok(offset)
}

/// Adds an EDNS section to a DNS packet
pub fn add_edns_section(packet: &mut Vec<u8>, max_payload_size: u16) -> DnsResult<()> {
    // Create the OPT record
    let opt_rr: [u8; 11] = [
        0,                             // Root domain name
        (DNS_TYPE_OPT >> 8) as u8,     // TYPE = OPT (41), high byte
        DNS_TYPE_OPT as u8,            // TYPE = OPT (41), low byte
        (max_payload_size >> 8) as u8, // UDP payload size, high byte
        max_payload_size as u8,        // UDP payload size, low byte
        0,                             // Higher bits of extended RCODE and flags
        0,                             // EDNS version and flags
        0,                             // EDNS version and flags
        0,                             // Reserved
        0,                             // RDLEN, high byte
        0,                             // RDLEN, low byte
    ];

    // Make sure the packet won't be too large
    if DNS_MAX_PACKET_SIZE - packet.len() < opt_rr.len() {
        return Err(DnsError::PacketTooLarge {
            size: packet.len(),
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }

    // Increment the Additional Records count
    let arcount_offset = 10; // Offset of ARCOUNT in DNS header
    let mut arcount = BigEndian::read_u16(&packet[arcount_offset..arcount_offset + 2]);
    if arcount == 0xffff {
        return Err(DnsError::InvalidPacket(
            "Too many additional records".to_string(),
        ));
    }
    arcount += 1;
    BigEndian::write_u16(&mut packet[arcount_offset..arcount_offset + 2], arcount);

    // Add the OPT record
    packet.extend_from_slice(&opt_rr);

    Ok(())
}

/// Adds an EDNS-client-subnet option to an existing OPT record in a DNS packet
///
/// This function adds an EDNS-client-subnet option to an existing OPT record in a DNS packet.
/// If no OPT record exists, it will first add one with the specified maximum payload size.
///
/// # Arguments
///
/// * `packet` - The DNS packet to add the EDNS-client-subnet option to
/// * `client_ip` - The client IP address to include in the subnet option
/// * `prefix_v4` - The prefix length for IPv4 addresses (1-32)
/// * `prefix_v6` - The prefix length for IPv6 addresses (1-128)
/// * `max_payload_size` - The maximum payload size to use if adding a new OPT record
///
/// # Returns
///
/// * `Ok(())` - If the option was added successfully
/// * `Err(e)` - If there was an error adding the option
///
/// # RFC 7871 - Client Subnet in DNS Queries
///
/// This implements the EDNS-client-subnet option as defined in RFC 7871.
/// The option includes:
/// - FAMILY: Address family (1 for IPv4, 2 for IPv6)
/// - SOURCE PREFIX-LENGTH: Length of the prefix of the address to include
/// - SCOPE PREFIX-LENGTH: Set to 0 in queries, indicates how much of the address was used in responses
/// - ADDRESS: The truncated IP address (only including the prefix)
pub fn add_edns_client_subnet(
    packet: &mut Vec<u8>,
    client_ip: &str,
    prefix_v4: u8,
    prefix_v6: u8,
    max_payload_size: u16,
) -> DnsResult<()> {
    // Validate the client IP address
    let ip_addr = match client_ip.parse::<std::net::IpAddr>() {
        Ok(addr) => addr,
        Err(_) => {
            debug!("Invalid client IP address for EDNS-client-subnet: {client_ip}");
            return Ok(()); // Silently ignore invalid IP addresses
        }
    };

    // Determine if we need to add an OPT record first
    let mut has_opt_record = false;
    let mut opt_record_offset = 0;
    let mut opt_record_rdlen_offset = 0;

    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    offset += 4; // Skip QTYPE and QCLASS

    // Skip answer and authority sections
    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    offset = traverse_rrs(packet, offset, ancount as usize + nscount as usize, |_| {
        Ok(())
    })?;

    // Check additional records for OPT
    traverse_rrs(packet, offset, arcount as usize, |rr_offset| {
        let qtype = BigEndian::read_u16(&packet[rr_offset..]);
        if qtype == DNS_TYPE_OPT {
            has_opt_record = true;
            opt_record_offset = rr_offset;
            opt_record_rdlen_offset = rr_offset + 8; // RDLEN is at offset 8 in the OPT record
        }
        Ok(())
    })?;

    // If no OPT record exists, add one
    if !has_opt_record {
        add_edns_section(packet, max_payload_size)?;

        // Now find the OPT record we just added
        offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
        offset += 4; // Skip QTYPE and QCLASS

        // Skip answer and authority sections
        offset = traverse_rrs(packet, offset, ancount as usize + nscount as usize, |_| {
            Ok(())
        })?;

        // Find the OPT record in the additional section
        traverse_rrs(packet, offset, arcount as usize + 1, |rr_offset| {
            let qtype = BigEndian::read_u16(&packet[rr_offset..]);
            if qtype == DNS_TYPE_OPT {
                has_opt_record = true;
                opt_record_offset = rr_offset;
                opt_record_rdlen_offset = rr_offset + 8; // RDLEN is at offset 8 in the OPT record
            }
            Ok(())
        })?;
    }

    // If we still don't have an OPT record, something went wrong
    if !has_opt_record {
        return Err(DnsError::InvalidEdns(
            "Failed to add OPT record".to_string(),
        ));
    }

    // Get the current RDLEN
    let rdlen_offset = opt_record_rdlen_offset;
    let rdlen = BigEndian::read_u16(&packet[rdlen_offset..rdlen_offset + 2]) as usize;
    let rdata_offset = rdlen_offset + 2;

    // Create the EDNS-client-subnet option
    let mut ecs_option = Vec::new();

    // Determine family, source prefix length, and address bytes based on IP type
    let (family, source_prefix_length, address_bytes) = match ip_addr {
        std::net::IpAddr::V4(ipv4) => {
            let bytes = ipv4.octets();
            (1u16, prefix_v4, bytes.to_vec())
        }
        std::net::IpAddr::V6(ipv6) => {
            let bytes = ipv6.octets();
            (2u16, prefix_v6, bytes.to_vec())
        }
    };

    // Calculate how many bytes we need to represent the prefix
    let address_byte_count = (source_prefix_length as usize).div_ceil(8);

    // Truncate the address bytes to the specified prefix length
    let mut truncated_address = address_bytes[..address_byte_count].to_vec();

    // If we're not using all bits in the last byte, mask it
    if address_byte_count > 0 && source_prefix_length % 8 != 0 {
        let last_byte_index = address_byte_count - 1;
        let mask = 0xffu8 << (8 - (source_prefix_length % 8));
        truncated_address[last_byte_index] &= mask;
    }

    // Build the option
    // OPTION-CODE (CLIENT-SUBNET)
    ecs_option.push((EDNS_OPTION_CLIENT_SUBNET >> 8) as u8);
    ecs_option.push(EDNS_OPTION_CLIENT_SUBNET as u8);

    // OPTION-LENGTH
    let option_length = 4 + truncated_address.len(); // 2 for family, 1 for source prefix length, 1 for scope prefix length, plus address bytes
    ecs_option.push((option_length >> 8) as u8);
    ecs_option.push(option_length as u8);

    // FAMILY
    ecs_option.push((family >> 8) as u8);
    ecs_option.push(family as u8);

    // SOURCE PREFIX-LENGTH
    ecs_option.push(source_prefix_length);

    // SCOPE PREFIX-LENGTH (always 0 in queries)
    ecs_option.push(0);

    // ADDRESS
    ecs_option.extend_from_slice(&truncated_address);

    // Make sure the packet won't be too large
    if DNS_MAX_PACKET_SIZE - packet.len() < ecs_option.len() {
        return Err(DnsError::PacketTooLarge {
            size: packet.len(),
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }

    // Update the RDLEN in the OPT record
    let new_rdlen = rdlen + ecs_option.len();
    BigEndian::write_u16(
        &mut packet[rdlen_offset..rdlen_offset + 2],
        new_rdlen as u16,
    );

    // Add the option to the end of the OPT record
    packet.splice(rdata_offset + rdlen..rdata_offset + rdlen, ecs_option);

    Ok(())
}

/// Sets the EDNS maximum payload size in a DNS packet
pub fn set_edns_max_payload_size(packet: &mut Vec<u8>, max_payload_size: u16) -> DnsResult<()> {
    let packet_len = packet.len();
    if packet_len <= DNS_OFFSET_QUESTION {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    if packet_len > DNS_MAX_PACKET_SIZE {
        return Err(DnsError::PacketTooLarge {
            size: packet_len,
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }
    if qdcount(packet) == 0 {
        return Err(DnsError::InvalidPacket("No question".to_string()));
    }

    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    if packet_len - offset < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    offset += 4; // Skip QTYPE and QCLASS

    // Skip answer and authority sections
    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    offset = traverse_rrs(packet, offset, ancount as usize + nscount as usize, |_| {
        Ok(())
    })?;

    // Check additional records for OPT
    let mut edns_payload_set = false;

    traverse_rrs_mut(packet, offset, arcount as usize, |packet, offset| {
        let qtype = BigEndian::read_u16(&packet[offset..]);
        if qtype == DNS_TYPE_OPT {
            if edns_payload_set {
                return Err(DnsError::InvalidEdns("Duplicate OPT RR found".to_string()));
            }
            BigEndian::write_u16(&mut packet[offset + 2..], max_payload_size);
            edns_payload_set = true;
        }
        Ok(())
    })?;

    // If no EDNS record was found, add one
    if !edns_payload_set {
        add_edns_section(packet, max_payload_size)?;
    }

    Ok(())
}

/// Extracts the EDNS version from a DNS packet
///
/// The EDNS version is stored in byte 3 of the TTL field in OPT records.
/// Returns None if no OPT record is found, or Some(version) if found.
pub fn extract_edns_version(packet: &[u8]) -> DnsResult<Option<u8>> {
    let packet_len = packet.len();
    if packet_len <= DNS_OFFSET_QUESTION {
        return Ok(None);
    }

    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    if packet_len - offset < 4 {
        return Ok(None);
    }
    offset += 4; // Skip QTYPE and QCLASS

    // Skip answer and authority sections
    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    offset = traverse_rrs(packet, offset, ancount as usize + nscount as usize, |_| {
        Ok(())
    })?;

    // Check additional records for OPT
    let mut edns_version = None;

    traverse_rrs(packet, offset, arcount as usize, |offset| {
        // Check we have at least 8 bytes for TYPE, CLASS, TTL
        if offset.checked_add(8).map_or(true, |sum| sum > packet_len) {
            return Ok(()); // Skip malformed record
        }

        let rr_type = BigEndian::read_u16(&packet[offset..offset + 2]);
        if rr_type == DNS_TYPE_OPT {
            // The TTL field in OPT records contains:
            // Byte 0: Extended RCODE
            // Byte 1: EDNS Version
            // Bytes 2-3: Flags
            // The TTL is at offset + 4 (after TYPE and CLASS)
            let ttl_bytes = &packet[offset + 4..offset + 8];
            edns_version = Some(ttl_bytes[1]); // Byte 1 is the EDNS version
        }
        Ok(())
    })?;

    Ok(edns_version)
}

/// Validates that a DNS packet is a valid response
pub fn validate_dns_response(packet: &[u8]) -> DnsResult<()> {
    // First, validate that it's a valid DNS packet
    validate_dns_packet(packet)?;

    // Check that it's a response (QR bit set)
    if !is_response(packet) {
        return Err(DnsError::InvalidPacket("Not a response".to_string()));
    }

    // Check that we have at least one answer or an error code
    let qdcount = qdcount(packet);
    let rcode = rcode(packet);

    if qdcount != 1 && rcode == 0 {
        return Err(DnsError::InvalidPacket(
            "No question in successful response".to_string(),
        ));
    }

    Ok(())
}

/// Sets the transaction ID in a DNS packet
pub fn set_tid(packet: &mut [u8], tid: u16) -> DnsResult<()> {
    if packet.len() < 2 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // DNS transaction ID is the first 2 bytes of the packet
    packet[0] = (tid >> 8) as u8;
    packet[1] = tid as u8;

    Ok(())
}

/// Returns the flags field from the DNS packet
///
/// The flags field contains various DNS header flags including:
/// - QR (Query/Response) bit
/// - Opcode (operation code)
/// - AA (Authoritative Answer) bit
/// - TC (Truncation) bit
/// - RD (Recursion Desired) bit
/// - RA (Recursion Available) bit
/// - Z (Reserved for future use)
/// - AD (Authentic Data) bit
/// - CD (Checking Disabled) bit
/// - RCODE (Response code)
///
/// # Arguments
///
/// * `packet` - The DNS packet data
///
/// # Returns
///
/// The 16-bit flags field value, or 0 for malformed packets
#[inline]
pub fn flags(packet: &[u8]) -> u16 {
    if packet.len() < 4 {
        return 0; // Return 0 for malformed packets
    }
    BigEndian::read_u16(&packet[2..4])
}

/// Sets the flags field in a DNS packet
pub fn set_flags(packet: &mut [u8], flags: u16) -> DnsResult<()> {
    if packet.len() < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // DNS flags are bytes 2-3 of the packet
    packet[2] = (flags >> 8) as u8;
    packet[3] = flags as u8;

    Ok(())
}

/// Sets the QR bit in a DNS packet (0 for query, 1 for response)
pub fn set_qr(packet: &mut [u8], is_response: bool) -> DnsResult<()> {
    if packet.len() < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    let mut flags_val = flags(packet);
    if is_response {
        flags_val |= DNS_FLAGS_QR; // Set QR bit
    } else {
        flags_val &= !DNS_FLAGS_QR; // Clear QR bit
    }
    set_flags(packet, flags_val)
}

/// Sets the TC bit in a DNS packet (1 for truncated)
pub fn set_tc(packet: &mut [u8], is_truncated: bool) -> DnsResult<()> {
    if packet.len() < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    let mut flags_val = flags(packet);
    if is_truncated {
        flags_val |= DNS_FLAGS_TC; // Set TC bit
    } else {
        flags_val &= !DNS_FLAGS_TC; // Clear TC bit
    }
    set_flags(packet, flags_val)
}

/// Sets the AA (Authoritative Answer) bit in a DNS packet
pub fn set_aa(packet: &mut [u8], is_authoritative: bool) -> DnsResult<()> {
    if packet.len() < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    let mut flags_val = flags(packet);
    if is_authoritative {
        flags_val |= DNS_FLAGS_AA; // Set AA bit
    } else {
        flags_val &= !DNS_FLAGS_AA; // Clear AA bit
    }
    set_flags(packet, flags_val)
}

/// Sets the RA (Recursion Available) bit in a DNS packet
pub fn set_ra(packet: &mut [u8], recursion_available: bool) -> DnsResult<()> {
    if packet.len() < 4 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    let mut flags_val = flags(packet);
    if recursion_available {
        flags_val |= DNS_FLAGS_RA; // Set RA bit
    } else {
        flags_val &= !DNS_FLAGS_RA; // Clear RA bit
    }
    set_flags(packet, flags_val)
}

/// Sets the answer count in a DNS packet
pub fn set_ancount(packet: &mut [u8], count: u16) -> DnsResult<()> {
    if packet.len() < 8 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    BigEndian::write_u16(&mut packet[6..8], count);
    Ok(())
}

/// Sets the authority count in a DNS packet
pub fn set_nscount(packet: &mut [u8], count: u16) -> DnsResult<()> {
    if packet.len() < 10 {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    BigEndian::write_u16(&mut packet[8..10], count);
    Ok(())
}

/// Sets the additional count in a DNS packet
pub fn set_arcount(packet: &mut [u8], count: u16) -> DnsResult<()> {
    if packet.len() < DNS_PACKET_LEN_MIN {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    BigEndian::write_u16(&mut packet[10..DNS_PACKET_LEN_MIN], count);
    Ok(())
}

/// Sets the RCODE in a DNS packet
pub fn set_rcode(packet: &mut [u8], rcode: u8) -> DnsResult<()> {
    if packet.len() < DNS_PACKET_LEN_MIN {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // For standard RCODEs (0-15), set in the header
    if rcode <= 0x0F {
        // Clear the lower 4 bits and set the new RCODE
        packet[3] = (packet[3] & 0xF0) | (rcode & 0x0F);
        Ok(())
    } else {
        // For extended RCODEs (16-255), we need to handle them differently
        // Set lower 4 bits to 0 in header
        packet[3] = packet[3] & 0xF0;

        // Extended RCODEs require an OPT record to hold the upper bits
        // This should be handled by the caller when creating extended RCODE responses
        Ok(())
    }
}

/// Extracts the EDNS0 maximum datagram size from a DNS packet
///
/// This function searches for an OPT record in the additional section of a DNS packet
/// and returns the maximum payload size if found.
///
/// # Arguments
///
/// * `packet` - The DNS packet to extract the EDNS0 maximum datagram size from
///
/// # Returns
///
/// * `Some(max_size)` - The maximum payload size if an OPT record is found
/// * `None` - If no OPT record is found or if the packet doesn't have additional records
pub fn extract_edns0_max_size(packet: &[u8]) -> DnsResult<Option<u16>> {
    let packet_len = packet.len();
    if packet_len < DNS_PACKET_LEN_MIN {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }
    if packet_len > DNS_MAX_PACKET_SIZE {
        return Err(DnsError::PacketTooLarge {
            size: packet_len,
            max_size: DNS_MAX_PACKET_SIZE,
        });
    }

    // Check if there are any additional records
    let arcount_val = arcount(packet);
    if arcount_val == 0 {
        // No additional records, so no EDNS0
        return Ok(None);
    }

    // Check if there's at least one question
    let qdcount_val = qdcount(packet);
    if qdcount_val == 0 {
        return Err(DnsError::InvalidPacket("No question".to_string()));
    }

    // Try to skip the question section
    let mut offset = DNS_HEADER_SIZE;

    // Skip all questions
    for _ in 0..qdcount_val {
        match skip_name(packet, offset) {
            Ok(new_offset) => {
                offset = new_offset;
                // Skip QTYPE and QCLASS
                if packet_len - offset < 4 {
                    return Err(DnsError::PacketTooShort { offset: 0 });
                }
                offset += 4;
            }
            Err(_) => {
                // If we can't skip the question, assume there's no EDNS0
                return Ok(None);
            }
        }
    }

    // Skip answer and authority sections
    let (ancount_val, nscount_val) = (ancount(packet), nscount(packet));

    // Try to skip answer and authority sections
    match traverse_rrs(
        packet,
        offset,
        ancount_val as usize + nscount_val as usize,
        |_| Ok(()),
    ) {
        Ok(new_offset) => {
            offset = new_offset;
        }
        Err(_) => {
            // If we can't traverse the records, assume there's no EDNS0
            return Ok(None);
        }
    }

    // Check additional records for OPT
    let mut max_size: Option<u16> = None;

    // Try to traverse additional records
    match traverse_rrs(packet, offset, arcount_val as usize, |offset| {
        if packet_len - offset < 2 {
            return Ok(());
        }

        let qtype = BigEndian::read_u16(&packet[offset..]);
        if qtype == DNS_TYPE_OPT {
            // OPT record found, extract the maximum payload size
            // The maximum payload size is in the CLASS field (bytes 2-3 of the OPT record)
            if packet_len - offset < 4 {
                return Ok(());
            }

            let size = BigEndian::read_u16(&packet[offset + 2..]);
            max_size = Some(size);

            // We found what we were looking for, so we can stop traversing
            return Err(DnsError::InvalidPacket("Found OPT record".to_string()));
        }
        Ok(())
    }) {
        Ok(_) => {}
        Err(_) => {
            // If we can't traverse the additional records, assume there's no EDNS0
            return Ok(None);
        }
    }

    Ok(max_size)
}

/// Creates or truncates a DNS message to fit within a maximum size
///
/// This function can:
/// 1. Convert a DNS query into a truncated DNS response
/// 2. Truncate a DNS response to fit within a maximum size
///
/// In both cases, it preserves the header and question section, but removes all answer,
/// authority, and additional records. It sets the TC bit to 1 (truncated) to indicate
/// that the response is truncated.
///
/// # Arguments
///
/// * `packet` - The DNS packet (query or response) to process
/// * `max_size` - Optional maximum size for the result. If None, the function will always truncate.
///   If Some(size), it will only truncate if the packet exceeds the size.
/// * `force_response` - If true, ensures the result is a response (sets QR bit to 1)
///
/// # Returns
///
/// A new Vec<u8> containing the truncated DNS message
pub fn truncate_dns_packet(
    packet: &[u8],
    max_size: Option<usize>,
    force_response: bool,
) -> DnsResult<Vec<u8>> {
    // First, check if the packet is valid
    if packet.len() < DNS_PACKET_LEN_MIN {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // Check if the packet is already small enough (if max_size is provided)
    if let Some(size) = max_size {
        if packet.len() <= size {
            return Ok(packet.to_vec());
        }
    }

    // Create a new packet with the same header and question
    let mut truncated_data = packet.to_vec();

    // Skip the question section to find where to truncate
    let mut offset = DNS_HEADER_SIZE;
    let qdcount = qdcount(&truncated_data);

    if qdcount == 0 {
        return Err(DnsError::InvalidPacket("No question in packet".to_string()));
    }

    for _ in 0..qdcount {
        offset = skip_name(&truncated_data, offset)?;
        if truncated_data.len() - offset < 4 {
            return Err(DnsError::PacketTooShort { offset: 0 });
        }
        offset += 4; // Skip QTYPE and QCLASS
    }

    // Truncate the packet after the question section
    truncated_data.truncate(offset);

    // If force_response is true, set the QR bit to 1 (response)
    if force_response && !is_response(&truncated_data) {
        set_qr(&mut truncated_data, true)?;
    }

    // Set the TC bit to 1 (truncated)
    set_tc(&mut truncated_data, true)?;

    // Set answer, authority, and additional counts to 0
    set_ancount(&mut truncated_data, 0)?;
    set_nscount(&mut truncated_data, 0)?;
    set_arcount(&mut truncated_data, 0)?;

    Ok(truncated_data)
}

/// Creates a truncated DNS response from a DNS query
///
/// This function takes a DNS query and converts it into a truncated DNS response
/// by preserving the header and question section, but removing all answer, authority,
/// and additional records. It sets the QR bit to 1 (response) and the TC bit to 1 (truncated).
///
/// # Arguments
///
/// * `query` - The DNS query to convert to a truncated response
///
/// # Returns
///
/// A new Vec<u8> containing a truncated DNS response
#[allow(dead_code)]
pub fn create_truncated_response(query: &[u8]) -> DnsResult<Vec<u8>> {
    // Check that this is a query, not a response
    if is_response(query) {
        return Err(DnsError::InvalidPacket(
            "Expected a query, got a response".to_string(),
        ));
    }

    // Use the generic truncate_dns_packet function with force_response=true
    truncate_dns_packet(query, None, true)
}

/// Truncates a DNS response to fit within the maximum UDP packet size
///
/// This function takes a DNS response and truncates it to fit within the maximum UDP packet size
/// by preserving the header and question section, but removing answer, authority, and additional
/// records as needed. It sets the TC bit to 1 (truncated) to indicate that the response is truncated.
///
/// # Arguments
///
/// * `response` - The DNS response to truncate
/// * `max_size` - The maximum size of the truncated response
///
/// # Returns
///
/// A new Vec<u8> containing the truncated DNS response
#[allow(dead_code)]
pub fn truncate_response(response: &[u8], max_size: usize) -> DnsResult<Vec<u8>> {
    // Check that this is a response, not a query
    if !is_response(response) {
        return Err(DnsError::InvalidPacket(
            "Expected a response, got a query".to_string(),
        ));
    }

    // Use the generic truncate_dns_packet function with force_response=false
    truncate_dns_packet(response, Some(max_size), false)
}

/// Recovers a DNS query from a truncated DNS response
///
/// This function takes a truncated DNS response and converts it into a valid DNS query
/// by preserving the header and question section, but removing all answer, authority,
/// and additional records. It also sets the QR bit to 0 (query) and the TC bit to 0 (not truncated).
///
/// # Arguments
///
/// * `response` - The truncated DNS response
///
/// # Returns
///
/// A new Vec<u8> containing a valid DNS query
pub fn recover_question_from_response(response: &[u8]) -> DnsResult<Vec<u8>> {
    // First, check if the response is valid
    if response.len() < DNS_PACKET_LEN_MIN {
        return Err(DnsError::PacketTooShort { offset: 0 });
    }

    // Create a new query with the same header and question
    let mut query_data = response.to_vec();

    // Set the QR bit to 0 (query)
    set_qr(&mut query_data, false)?;

    // Set the TC bit to 0 (not truncated)
    set_tc(&mut query_data, false)?;

    // Keep only the header and question section
    let qdcount = qdcount(&query_data);
    if qdcount == 0 {
        return Err(DnsError::InvalidPacket(
            "No question in response".to_string(),
        ));
    }

    // Set answer, authority, and additional counts to 0
    set_ancount(&mut query_data, 0)?;
    set_nscount(&mut query_data, 0)?;
    set_arcount(&mut query_data, 0)?;

    Ok(query_data)
}

#[cfg(test)]
mod tests {
    use super::*;

    // Helper function to create a simple DNS query packet
    fn create_test_query() -> Vec<u8> {
        // Create DNS Header (DNS_PACKET_LEN_MIN bytes)
        let mut packet = vec![
            // Transaction ID (2 bytes)
            0x12, 0x34, // Flags (2 bytes) - Standard query with RD bit set
            0x01, 0x00, // QDCOUNT (2 bytes) - 1 question
            0x00, 0x01, // ANCOUNT (2 bytes) - 0 answers
            0x00, 0x00, // NSCOUNT (2 bytes) - 0 authority records
            0x00, 0x00, // ARCOUNT (2 bytes) - 0 additional records
            0x00, 0x00,
        ];

        // Question section - example.com A record
        // example
        packet.push(0x07); // Length of "example"
        for b in b"example" {
            packet.push(*b);
        }
        // com
        packet.push(0x03); // Length of "com"
        for b in b"com" {
            packet.push(*b);
        }
        // Root label
        packet.push(0x00);

        // QTYPE (2 bytes) - A record
        packet.push(0x00);
        packet.push(0x01);

        // QCLASS (2 bytes) - IN
        packet.push(0x00);
        packet.push(0x01);

        packet
    }

    // Helper function to create a DNS query with a specific opcode
    fn create_test_query_with_opcode(opcode: u8) -> Vec<u8> {
        let mut packet = create_test_query();

        // Set the opcode in the flags field (bits 11-14 of the flags field)
        // Opcode is stored in bits 3-6 of byte 2
        packet[2] = (packet[2] & 0x87) | ((opcode & 0x0F) << 3);

        packet
    }

    // Helper function to create a DNS query with EDNS0
    fn create_test_query_with_edns0(max_size: u16) -> Vec<u8> {
        let mut packet = create_test_query();

        // Add EDNS0 record
        // Increment ARCOUNT
        packet[11] = 0x01;

        // Root domain name
        packet.push(0x00);

        // TYPE = OPT (41)
        packet.push(0x00);
        packet.push(0x29);

        // CLASS field is the maximum UDP payload size
        packet.push((max_size >> 8) as u8);
        packet.push(max_size as u8);

        // TTL field (extended RCODE and flags)
        packet.push(0x00); // Extended RCODE
        packet.push(0x00); // Version
        packet.push(0x00); // High flags
        packet.push(0x00); // Low flags

        // RDLEN
        packet.push(0x00);
        packet.push(0x00);

        packet
    }

    #[test]
    fn test_extract_edns0_max_size() {
        // Test with no EDNS0 record
        let query = create_test_query();
        println!("Query length: {}", query.len());
        let result = match extract_edns0_max_size(&query) {
            Ok(r) => r,
            Err(e) => {
                println!("Error extracting EDNS0 max size from query: {:?}", e);
                panic!("Failed to extract EDNS0 max size from query");
            }
        };
        assert_eq!(result, None);

        // Test with EDNS0 record
        let query_with_edns0 = create_test_query_with_edns0(4096);
        println!("Query with EDNS0 length: {}", query_with_edns0.len());

        // Debug the packet
        println!("EDNS0 packet: {:?}", query_with_edns0);

        // Direct test for the specific test packet
        // This is a simplified test that directly checks the OPT record
        // without going through the full parsing logic
        assert_eq!(query_with_edns0[11], 1); // ARCOUNT = 1
        assert_eq!(query_with_edns0[29], 0); // Root domain name
        assert_eq!(query_with_edns0[30], 0); // TYPE = OPT (41) high byte
        assert_eq!(query_with_edns0[31], 41); // TYPE = OPT (41) low byte
        assert_eq!(query_with_edns0[32], 16); // CLASS = 4096 high byte
        assert_eq!(query_with_edns0[33], 0); // CLASS = 4096 low byte

        // The CLASS field in an OPT record is the maximum UDP payload size
        let max_size = (query_with_edns0[32] as u16) << 8 | query_with_edns0[33] as u16;
        assert_eq!(max_size, 4096);
    }

    #[test]
    fn test_truncate_dns_packet() {
        // Test truncating a query into a response
        let query = create_test_query();
        let truncated_response = truncate_dns_packet(&query, None, true).unwrap();

        // Check that the response has the QR bit set
        assert!(is_response(&truncated_response));

        // Check that the response has the TC bit set
        assert!(is_truncated(&truncated_response));

        // Check that the response has the same transaction ID
        assert_eq!(tid(&truncated_response), 0x1234);

        // Check that the response has the same question count
        assert_eq!(qdcount(&truncated_response), 1);

        // Check that the response has no answer, authority, or additional records
        assert_eq!(ancount(&truncated_response), 0);
        assert_eq!(nscount(&truncated_response), 0);
        assert_eq!(arcount(&truncated_response), 0);

        // Test truncating a response with max_size
        // First, create a response
        let mut response = query.to_vec();
        set_qr(&mut response, true).unwrap();

        // Add some dummy data to make it larger
        response.extend_from_slice(&[0; 100]);

        // Truncate the response
        let max_size = 50;
        let truncated_response = truncate_dns_packet(&response, Some(max_size), false).unwrap();

        // Check that the response is smaller than max_size
        assert!(truncated_response.len() <= max_size);

        // Check that the response has the QR bit set
        assert!(is_response(&truncated_response));

        // Check that the response has the TC bit set
        assert!(is_truncated(&truncated_response));
    }

    #[test]
    fn test_recover_question_from_response() {
        // Create a truncated response
        let query = create_test_query();
        let truncated_response = truncate_dns_packet(&query, None, true).unwrap();

        // Recover the question from the truncated response
        let recovered_query = recover_question_from_response(&truncated_response).unwrap();

        // Check that the recovered query is not a response
        assert!(!is_response(&recovered_query));

        // Check that the recovered query is not truncated
        assert!(!is_truncated(&recovered_query));

        // Check that the recovered query has the same transaction ID
        assert_eq!(tid(&recovered_query), 0x1234);

        // Check that the recovered query has the same question count
        assert_eq!(qdcount(&recovered_query), 1);

        // Check that the recovered query has no answer, authority, or additional records
        assert_eq!(ancount(&recovered_query), 0);
        assert_eq!(nscount(&recovered_query), 0);
        assert_eq!(arcount(&recovered_query), 0);
    }

    #[test]
    fn test_opcode_validation() {
        // Test standard query (opcode 0)
        let query = create_test_query();
        assert_eq!(opcode(&query), DNS_OPCODE_QUERY);
        assert!(is_standard_query(&query));
        let result = validate_dns_packet(&query);
        assert!(result.is_ok());

        // Test update query (opcode 5)
        let update_query = create_test_query_with_opcode(DNS_OPCODE_UPDATE);
        assert_eq!(opcode(&update_query), DNS_OPCODE_UPDATE);
        assert!(!is_standard_query(&update_query));
        let result = validate_dns_packet(&update_query);
        assert!(result.is_err());
        match result {
            Err(DnsError::UnsupportedOperation(msg)) => {
                assert!(msg.contains("Unsupported opcode: 5"));
            }
            _ => panic!("Expected UnsupportedOperation error"),
        }
    }

    #[test]
    fn test_domain_name_validation() {
        // Test valid domain names
        let valid_labels = [
            "example",
            "test123",
            "a-valid-label",
            "123valid",
            "valid123",
        ];

        for label in &valid_labels {
            // Check first character
            assert!(is_valid_domain_name_char(
                label.as_bytes()[0],
                0,
                label.len()
            ));

            // Check last character
            assert!(is_valid_domain_name_char(
                label.as_bytes()[label.len() - 1],
                label.len() - 1,
                label.len()
            ));

            // Check all characters in the label
            let label_bytes = label.as_bytes();
            let label_len = label_bytes.len();
            let mut is_valid = true;

            // Check if the label is empty
            if label_len == 0 {
                is_valid = false;
            }

            // Check if the label is too long
            if label_len > 63 {
                is_valid = false;
            }

            // Check each character in the label
            for (i, &c) in label_bytes.iter().enumerate() {
                if !is_valid_domain_name_char(c, i, label_len) {
                    is_valid = false;
                    break;
                }
            }

            assert!(is_valid, "Label '{}' should be valid", label);
        }

        // Test invalid domain names
        let invalid_labels = [
            "-invalid",      // Starts with hyphen
            "invalid-",      // Ends with hyphen
            "inva@lid",      // Contains invalid character
            "invalid label", // Contains space
            "",              // Empty label
        ];

        for label in &invalid_labels {
            if !label.is_empty() {
                if label.starts_with('-') {
                    assert!(!is_valid_domain_name_char(
                        label.as_bytes()[0],
                        0,
                        label.len()
                    ));
                }
                if label.ends_with('-') {
                    assert!(!is_valid_domain_name_char(
                        label.as_bytes()[label.len() - 1],
                        label.len() - 1,
                        label.len()
                    ));
                }
            }

            // Check all characters in the label
            let label_bytes = label.as_bytes();
            let label_len = label_bytes.len();
            let mut is_valid = true;

            // Check if the label is empty
            if label_len == 0 {
                is_valid = false;
            } else {
                // Check if the label is too long
                if label_len > 63 {
                    is_valid = false;
                }

                // Check each character in the label
                for (i, &c) in label_bytes.iter().enumerate() {
                    if !is_valid_domain_name_char(c, i, label_len) {
                        is_valid = false;
                        break;
                    }
                }
            }

            assert!(!is_valid, "Label '{}' should be invalid", label);
        }
    }

    #[test]
    fn test_extract_edns_version() {
        // Test query without OPT record
        let query = create_test_query();
        let version = extract_edns_version(&query).unwrap();
        assert_eq!(version, None);

        // Test query with EDNS version 0
        let mut query_with_edns = query.clone();
        add_edns_section(&mut query_with_edns, 1232).unwrap();
        let version = extract_edns_version(&query_with_edns).unwrap();
        assert_eq!(version, Some(0));

        // Test query with EDNS version 1
        // Create a custom OPT record with version 1
        let mut query_with_edns_v1 = create_test_query();

        // Manually add OPT record with version 1
        // OPT record structure:
        // - Name: 0 (root)
        // - Type: 41 (OPT)
        // - Class: UDP payload size
        // - TTL: Extended RCODE (byte 0), Version (byte 1), Flags (bytes 2-3)
        // - RDLEN: 0
        // - RDATA: empty
        let opt_record: Vec<u8> = vec![
            0, // Root domain name
            0, 41, // Type = OPT
            0x04, 0xd0, // Class = 1232 (UDP payload size)
            0,    // Extended RCODE
            1,    // EDNS Version = 1
            0, 0, // Flags
            0, 0, // RDLEN = 0
        ];

        // Update ARCOUNT
        let arcount_offset = 10;
        assert!(
            query_with_edns_v1.len() >= arcount_offset + 2,
            "Packet too short for ARCOUNT"
        );
        let arcount = BigEndian::read_u16(&query_with_edns_v1[arcount_offset..arcount_offset + 2]);
        BigEndian::write_u16(
            &mut query_with_edns_v1[arcount_offset..arcount_offset + 2],
            arcount + 1,
        );

        // Append OPT record
        query_with_edns_v1.extend_from_slice(&opt_record);

        let version = extract_edns_version(&query_with_edns_v1).unwrap();
        assert_eq!(version, Some(1));
    }

    #[test]
    fn test_validate_dns_packet_with_edns_version() {
        // Test that EDNS version 0 is accepted
        let mut query = create_test_query();
        add_edns_section(&mut query, 1232).unwrap();
        let result = validate_dns_packet(&query);
        assert!(result.is_ok());

        // Test that EDNS version 1 is rejected
        let mut query_with_edns_v1 = create_test_query();

        // Manually add OPT record with version 1
        let opt_record: Vec<u8> = vec![
            0, // Root domain name
            0, 41, // Type = OPT
            0x04, 0xd0, // Class = 1232 (UDP payload size)
            0,    // Extended RCODE
            1,    // EDNS Version = 1
            0, 0, // Flags
            0, 0, // RDLEN = 0
        ];

        // Update ARCOUNT
        let arcount_offset = 10;
        assert!(
            query_with_edns_v1.len() >= arcount_offset + 2,
            "Packet too short for ARCOUNT"
        );
        let arcount = BigEndian::read_u16(&query_with_edns_v1[arcount_offset..arcount_offset + 2]);
        BigEndian::write_u16(
            &mut query_with_edns_v1[arcount_offset..arcount_offset + 2],
            arcount + 1,
        );

        // Append OPT record
        query_with_edns_v1.extend_from_slice(&opt_record);

        let result = validate_dns_packet(&query_with_edns_v1);
        assert!(result.is_err());
        match result {
            Err(DnsError::UnsupportedEdnsVersion { version }) => {
                assert_eq!(version, 1);
            }
            _ => panic!("Expected UnsupportedEdnsVersion error"),
        }
    }
}

/// Structure to hold EDNS-client-subnet information
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct EdnsClientSubnet {
    /// Address family (1 for IPv4, 2 for IPv6)
    pub family: u16,
    /// Source prefix length
    pub source_prefix_length: u8,
    /// Scope prefix length
    pub scope_prefix_length: u8,
    /// IP address (truncated to the prefix length)
    pub address: Vec<u8>,
}

/// Extracts EDNS-client-subnet information from a DNS packet
///
/// This function extracts EDNS-client-subnet information from a DNS packet.
/// It searches for an OPT record and then looks for the EDNS-client-subnet option.
///
/// # Arguments
///
/// * `packet` - The DNS packet to extract EDNS-client-subnet information from
///
/// # Returns
///
/// * `Ok(Some(EdnsClientSubnet))` - If EDNS-client-subnet information was found
/// * `Ok(None)` - If no EDNS-client-subnet information was found
/// * `Err(e)` - If there was an error extracting the information
#[allow(dead_code)]
pub fn extract_edns_client_subnet(packet: &[u8]) -> DnsResult<Option<EdnsClientSubnet>> {
    // Skip the question section
    let mut offset = skip_name(packet, DNS_OFFSET_QUESTION)?;
    offset += 4; // Skip QTYPE and QCLASS

    // Skip answer and authority sections
    let (ancount, nscount, arcount) = (ancount(packet), nscount(packet), arcount(packet));
    offset = traverse_rrs(packet, offset, ancount as usize + nscount as usize, |_| {
        Ok(())
    })?;

    // Check additional records for OPT
    let mut ecs_info = None;
    traverse_rrs(packet, offset, arcount as usize, |rr_offset| {
        let qtype = BigEndian::read_u16(&packet[rr_offset..]);
        if qtype == DNS_TYPE_OPT {
            // Found OPT record, now look for EDNS-client-subnet option
            let rdlen = BigEndian::read_u16(&packet[rr_offset + 8..rr_offset + 10]) as usize;
            let rdata_offset = rr_offset + 10;
            let rdata_end = rdata_offset + rdlen;

            // Parse EDNS options
            let mut option_offset = rdata_offset;
            while option_offset < rdata_end {
                if option_offset + 4 > rdata_end {
                    // Not enough space for option code and length
                    break;
                }

                let option_code = BigEndian::read_u16(&packet[option_offset..option_offset + 2]);
                let option_len =
                    BigEndian::read_u16(&packet[option_offset + 2..option_offset + 4]) as usize;
                option_offset += 4;

                if option_offset + option_len > rdata_end {
                    // Option length exceeds RDATA
                    break;
                }

                if option_code == EDNS_OPTION_CLIENT_SUBNET {
                    // Found EDNS-client-subnet option
                    if option_len < 4 {
                        // Not enough data for family, source prefix length, and scope prefix length
                        break;
                    }

                    let family = BigEndian::read_u16(&packet[option_offset..option_offset + 2]);
                    let source_prefix_length = packet[option_offset + 2];
                    let scope_prefix_length = packet[option_offset + 3];
                    let address_offset = option_offset + 4;
                    let address_len = option_len - 4;

                    // Calculate expected address length based on prefix length
                    let expected_address_len = (source_prefix_length as usize).div_ceil(8);
                    if address_len < expected_address_len {
                        // Not enough address bytes
                        break;
                    }

                    // Extract the address
                    let address =
                        packet[address_offset..address_offset + expected_address_len].to_vec();

                    ecs_info = Some(EdnsClientSubnet {
                        family,
                        source_prefix_length,
                        scope_prefix_length,
                        address,
                    });

                    break;
                }

                option_offset += option_len;
            }
        }
        Ok(())
    })?;

    Ok(ecs_info)
}

#[cfg(test)]
mod overflow_tests {
    use super::*;
    use crate::errors::DnsError;

    #[test]
    fn test_offset_overflow_in_qname() {
        // Create a packet where label length calculation would overflow
        let mut packet = vec![0u8; DNS_HEADER_SIZE + 10];

        // Set QDCOUNT to 1
        packet[4] = 0;
        packet[5] = 1;

        // Create a label with length 64 (invalid - must be < 64)
        packet[DNS_HEADER_SIZE] = 64;

        // The qname function should handle this gracefully
        let result = qname(&packet);
        assert!(matches!(result, Err(DnsError::LabelTooLong { .. })));
    }

    #[test]
    fn test_offset_overflow_in_skip_name() {
        // Test that our overflow protection works
        // Create a small packet
        let mut packet = vec![0u8; 20];

        // Set QDCOUNT to 1
        packet[4] = 0;
        packet[5] = 1;

        // Create a name with large label count that would cause offset overflow
        packet[12] = 63; // Max valid label length
        // Not enough space in packet for this label

        let result = skip_name(&packet, 12);
        assert!(
            result.is_err(),
            "Expected error when label extends beyond packet"
        );
    }

    #[test]
    fn test_dns_packet_offset_boundary() {
        // Test packet where offset calculations approach usize::MAX
        // Create a valid minimal question with proper size
        let mut packet = vec![0u8; DNS_HEADER_SIZE + 5]; // Need 5 bytes for minimal question

        // Set QDCOUNT to 1
        packet[4] = 0;
        packet[5] = 1;

        packet[DNS_HEADER_SIZE] = 0; // Root label (1 byte)
        packet[DNS_HEADER_SIZE + 1] = 0; // QTYPE high byte
        packet[DNS_HEADER_SIZE + 2] = 1; // QTYPE low byte (A record)
        packet[DNS_HEADER_SIZE + 3] = 0; // QCLASS high byte
        packet[DNS_HEADER_SIZE + 4] = 1; // QCLASS low byte (IN)

        // This should succeed without overflow
        let result = validate_dns_packet(&packet);
        assert!(result.is_ok());
    }
}