pub fn add_ecs_to_packet(
    packet: &mut Vec<u8>,
    client_ip: IpAddr,
    prefix_v4: u8,
    prefix_v6: u8,
) -> Result<(), Error> {
    use crate::dns;

    let packet_len = packet.len();
    ensure!(packet_len >= 12, "DNS packet too short");

    // Check if we already have EDNS
    let arcount = dns::arcount(packet);
    if arcount == 0 {
        // No additional records, need to add EDNS first
        dns::set_edns_max_payload_size(packet, 4096)?;
    }

    // Now find the OPT record and add ECS option
    let qdcount = dns::qdcount(packet);
    let ancount = dns::ancount(packet);
    let nscount = BigEndian::read_u16(&packet[8..10]);

    // Skip to additional section
    let mut offset = 12;

    // Skip question
    for _ in 0..qdcount {
        // Skip name
        while offset < packet_len && packet[offset] != 0 {
            if packet[offset] & 0xc0 == 0xc0 {
                offset += 2;
                break;
            }
            let len = packet[offset] as usize;
            offset += 1 + len;
        }
        if offset < packet_len && packet[offset] == 0 {
            offset += 1;
        }
        offset += 4; // type + class
    }

    // Skip answer and authority sections
    for _ in 0..(ancount + nscount) {
        offset = skip_rr(packet, offset)?;
    }

    // Find OPT record in additional section
    let mut opt_found = false;
    let _additional_start = offset;

    for _ in 0..arcount {
        let rr_start = offset;

        // Check if this is root domain (OPT record)
        if offset < packet_len && packet[offset] == 0 {
            // Check type
            if offset + 3 < packet_len {
                let rtype = BigEndian::read_u16(&packet[offset + 1..offset + 3]);
                if rtype == dns::DNS_TYPE_OPT {
                    opt_found = true;
                    // Found OPT record
                    offset += 1; // skip root domain
                    offset += 8; // skip type, class, ttl

                    if offset + 2 <= packet_len {
                        let rdlength = BigEndian::read_u16(&packet[offset..offset + 2]) as usize;
                        offset += 2;

                        // Build new ECS option
                        let ecs_data = build_ecs_option(client_ip, prefix_v4, prefix_v6);
                        let _ecs_option_len = 4 + ecs_data.len(); // 2 bytes code + 2 bytes length + data

                        // Check if we already have ECS option and remove it
                        let mut new_rdata = Vec::new();
                        let mut rdata_offset = 0;
                        while rdata_offset < rdlength {
                            if offset + rdata_offset + 4 <= packet_len {
                                let opt_code =
                                    BigEndian::read_u16(&packet[offset + rdata_offset..]);
                                let opt_len =
                                    BigEndian::read_u16(&packet[offset + rdata_offset + 2..])
                                        as usize;

                                if opt_code != EDNS_CLIENT_SUBNET {
                                    // Keep this option
                                    new_rdata.extend_from_slice(
                                        &packet[offset + rdata_offset
                                            ..offset + rdata_offset + 4 + opt_len],
                                    );
                                }
                                rdata_offset += 4 + opt_len;
                            } else {
                                break;
                            }
                        }

                        // Add our ECS option
                        new_rdata.extend_from_slice(&EDNS_CLIENT_SUBNET.to_be_bytes());
                        new_rdata.extend_from_slice(&(ecs_data.len() as u16).to_be_bytes());
                        new_rdata.extend_from_slice(&ecs_data);

                        // Update packet
                        let new_rdlength = new_rdata.len();
                        BigEndian::write_u16(&mut packet[offset - 2..], new_rdlength as u16);

                        // Save the data after the OPT record before modifying
                        let remaining_start = offset + rdlength;
                        let remaining_data = if remaining_start < packet_len {
                            packet[remaining_start..].to_vec()
                        } else {
                            Vec::new()
                        };

                        // Replace old rdata with new
                        packet.truncate(offset);
                        packet.extend_from_slice(&new_rdata);

                        // Add any remaining data after OPT record
                        if !remaining_data.is_empty() {
                            packet.extend_from_slice(&remaining_data);
                        }

                        return Ok(());
                    }
                }
            }
        }

        // Skip this RR
        offset = skip_rr(packet, rr_start)?;
    }

    if !opt_found {
        // Should have been added by set_edns_max_payload_size
        return Err(anyhow::anyhow!("Failed to find OPT record"));
    }

    Ok(())
}
Add or update EDNS Client Subnet in a DNS packet
src/libdoh/src/edns_ecs.rs