kapiti 0.0.3

use anyhow::{bail, Context, Result};
use std::convert::TryFrom;
use std::net::IpAddr;
use std::str;

use bytes::{BufMut, BytesMut};
use packed_struct::prelude::*;

use crate::codec::{domain_name, rdata};
use crate::specs::enums_generated::{self, ResourceClass, ResourceType};
use crate::specs::message::*;

#[derive(Clone, Debug)]
/// Information extracted from a DNS request, used internally for checking filters and sending a filtered response.
pub struct RequestInfo {
    /// The domain name to query, without a trailing '.'
    pub name: String,
    /// The resource type being queried
    pub resource_type: ResourceType,
    /// Whether dnssec_ok was enabled in the request
    pub dnssec_ok: bool,
    /// Responses to upstream clients must have the matching request ID.
    pub received_request_id: u16,
    /// We use the upstream client's requested UDP size for our response.
    pub requested_udp_size: u16,
}

/// Wire protocol type for OPT records, for use in some match statements.
pub const OPT_RESOURCE_TYPE: u16 = 41;

// Header, from RFC2535 section 6.1 or RFC2929 section 2.
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      ID                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |QR|   Opcode  |AA|TC|RD|RA| Z|AD|CD|   RCODE   |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    QDCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ANCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    NSCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                    ARCOUNT                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct HeaderBits {
    #[packed_field(bits = "0:15")]
    pub id: u16,

    #[packed_field(bits = "16:16")]
    pub is_response: bool,
    #[packed_field(bits = "17:20")]
    pub op_code: Integer<u8, packed_bits::Bits<4>>,
    #[packed_field(bits = "21:21")]
    pub authoritative: bool,
    #[packed_field(bits = "22:22")]
    pub truncated: bool,
    #[packed_field(bits = "23:23")]
    pub recursion_desired: bool,
    #[packed_field(bits = "24:24")]
    pub recursion_available: bool,
    #[packed_field(bits = "25:25")]
    pub reserved_9: bool,
    #[packed_field(bits = "26:26")]
    pub authentic_data: bool,
    #[packed_field(bits = "27:27")]
    pub checking_disabled: bool,
    #[packed_field(bits = "28:31")]
    pub response_code: Integer<u8, packed_bits::Bits<4>>,

    #[packed_field(bits = "32:47")]
    pub question_count: u16,
    #[packed_field(bits = "48:63")]
    pub answer_count: u16,
    #[packed_field(bits = "64:79")]
    pub authority_count: u16,
    #[packed_field(bits = "80:95")]
    pub additional_count: u16,
}

#[derive(Clone, Debug, PartialEq)]
pub struct RecordCounts {
    pub question: u16,
    pub answer: u16,
    pub authority: u16,
    pub additional: u16,
}

impl RecordCounts {
    pub fn new() -> RecordCounts {
        RecordCounts {
            question: 0,
            answer: 0,
            authority: 0,
            additional: 0,
        }
    }
}

/// Writes the message header to the provided buffer.
/// Allows specifying the transaction id since this is frequently overridden.
pub fn write_header_id(message: &Message, id_override: u16, buf: &mut BytesMut) -> Result<()> {
    // Ensure that the OPT record, if any, is included in additional_count
    let opt_len = match &message.opt {
        Some(_opt) => 1,
        None => 0,
    };
    write_header_bits(
        HeaderBits {
            id: id_override,

            is_response: message.header.is_response,
            op_code: match message.header.op_code {
                IntEnum::Enum(e) => Integer::from(e as u8),
                IntEnum::Unknown(i) => Integer::from(i),
            },
            authoritative: message.header.authoritative,
            truncated: message.header.truncated,
            recursion_desired: message.header.recursion_desired,
            recursion_available: message.header.recursion_available,
            reserved_9: message.header.reserved_9,
            authentic_data: message.header.authentic_data,
            checking_disabled: message.header.checking_disabled,
            response_code: match message.header.response_code {
                IntEnum::Enum(e) => Integer::from(e as u8),
                IntEnum::Unknown(i) => Integer::from(i),
            },

            question_count: u16::try_from(message.question.len())
                .with_context(|| "question length doesn't fit")?,
            answer_count: u16::try_from(message.answer.len())
                .with_context(|| "answer length doesn't fit")?,
            authority_count: u16::try_from(message.authority.len())
                .with_context(|| "authority length doesn't fit")?,
            additional_count: u16::try_from(message.additional.len() + opt_len)
                .with_context(|| "additional length doesn't fit")?,
        },
        buf,
    )
}

pub fn write_header_bits(bits: HeaderBits, buf: &mut BytesMut) -> Result<()> {
    let bits_packed = bits.pack()?;
    buf.reserve(bits_packed.len());
    buf.put_slice(&bits_packed);
    Ok(())
}

/// Returns the header content, along with whether it declares that the rest of the Message is truncated.
/// When this is encountered the message should be abandoned and a different connection method should be used.
/// (For example if the original request was UDP then try TCP)
/// Alternatively returns Ok(None) when there isn't enough data available to parse the header yet.
pub fn read_header(buf: &[u8], offset: &mut usize) -> Result<Option<(Header, RecordCounts, bool)>> {
    let headerbits_size = 12; // size of HeaderBits in bytes
    if buf.len() < *offset + headerbits_size {
        // Not enough bytes for the DNS header, wait for more bytes.
        return Ok(None);
    }
    let bits = HeaderBits::unpack_from_slice(&buf[*offset..*offset + headerbits_size])
        .with_context(|| "couldn't unpack header bits")?;

    let header = Header {
        id: bits.id,
        is_response: bits.is_response,
        op_code: match enums_generated::opcode_int(*bits.op_code as usize) {
            Some(e) => IntEnum::Enum(e),
            None => IntEnum::Unknown(*bits.op_code),
        },
        authoritative: bits.authoritative,
        truncated: bits.truncated,
        recursion_desired: bits.recursion_desired,
        recursion_available: bits.recursion_available,
        reserved_9: bits.reserved_9,
        authentic_data: bits.authentic_data,
        checking_disabled: bits.checking_disabled,
        response_code: match enums_generated::responsecode_int(*bits.response_code as usize) {
            Some(e) => IntEnum::Enum(e),
            None => IntEnum::Unknown(*bits.response_code),
        },
    };

    let record_counts = RecordCounts {
        question: bits.question_count,
        answer: bits.answer_count,
        authority: bits.authority_count,
        additional: bits.additional_count,
    };

    *offset += headerbits_size;
    Ok(Some((header, record_counts, bits.truncated)))
}

#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct HeaderIDBits {
    #[packed_field(bits = "0:15")]
    id: u16,
}

/// Directly updates the request ID in a serialized Message.
/// This is a shortcut to avoid needing to reserialize a new message in the event of a retry.
pub fn update_message_id(id: u16, buf: &mut BytesMut, message_offset: usize) -> Result<()> {
    let bits = HeaderIDBits { id }.pack()?;
    if buf.len() < message_offset + bits.len() {
        bail!(
            "Buffer is too small to update request ID of length {} at message_offset={}: buf=0x{:X}",
            buf.len(),
            message_offset,
            buf
        );
    }
    // The ID is (conveniently) located in the first two bytes of the message, so we can edit in-place relative to the offset.
    for i in 0..bits.len() {
        buf[message_offset + i] = bits[i];
    }
    Ok(())
}

/// Finds the minimum TTL across all resources in a Message, or None if there are no resources.
/// This is a utility for caching where we want to detect how much the response TTLs should be reduced.
/// For example, if the cache says that the remaining TTL is 20s, but the minimum resource TTL value is 30s,
/// then 10s has passed and all resources should be have TTLs reduced by 10s.
pub fn get_min_ttl_secs(message: &Message) -> Option<u32> {
    let mut min_ttl: Option<u32> = None;

    // Skip header and question(s): No TTLs

    for a in &message.answer {
        match min_ttl {
            Some(min_val) => {
                if a.ttl < min_val {
                    min_ttl = Some(a.ttl);
                }
            }
            None => {
                min_ttl = Some(a.ttl);
            }
        }
    }

    for a in &message.authority {
        match min_ttl {
            Some(min_val) => {
                if a.ttl < min_val {
                    min_ttl = Some(a.ttl);
                }
            }
            None => {
                min_ttl = Some(a.ttl);
            }
        }
    }

    // Note: any OPT resource is in a separate message.opt() field. OPT resources don't have TTLs.
    for a in &message.additional {
        match min_ttl {
            Some(min_val) => {
                if a.ttl < min_val {
                    min_ttl = Some(a.ttl);
                }
            }
            None => {
                min_ttl = Some(a.ttl);
            }
        }
    }

    return min_ttl;
}

/// Updates a cached DNS response to have the provided overrides.
pub fn update_cached_response(
    response: &mut Message,
    // received_request_id as orig_header_id, requested_udp_size as udp_size_override:
    request_info: &RequestInfo,
    cache_remaining_ttl_secs: u32,
) -> Result<()> {
    // Update TTL in response: Search all records for a minimum TTL value, then subtract that from all record TTLs.
    // For example, if the min TTL across all records is 30s and the cache response TTL is 20s, then 10s has passed and all TTLs should have 10s subtracted.
    let min_ttl_secs = get_min_ttl_secs(&response).with_context(|| {
        format!(
            "Missing resources in cached {:?} response for {}",
            request_info.resource_type, request_info.name
        )
    })?;
    if min_ttl_secs < cache_remaining_ttl_secs {
        // Shouldn't happen: Cache TTL should have started at min_ttl and gone down.
        bail!(
            "Redis had invalid TTL {}s with {:?} result for {}: {}",
            cache_remaining_ttl_secs,
            request_info.resource_type,
            request_info.name,
            response
        );
    }
    let ttl_subtract_secs = min_ttl_secs - cache_remaining_ttl_secs;

    // Do the updates
    response.header.id = request_info.received_request_id;
    for r in &mut response.answer {
        r.ttl -= ttl_subtract_secs;
    }
    for r in &mut response.authority {
        r.ttl -= ttl_subtract_secs;
    }
    if let Some(opt) = &mut response.opt {
        opt.udp_size = request_info.requested_udp_size;
    }
    for r in &mut response.additional {
        r.ttl -= ttl_subtract_secs;
    }
    Ok(())
}

// Question, from RFC1035 section 4.1.2
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                                               |
// /                     QNAME                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     QTYPE                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     QCLASS                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct QuestionBits {
    // name: variable length string
    #[packed_field(bits = "0:15")]
    resource_type: u16,
    #[packed_field(bits = "16:31")]
    resource_class: u16,
}

pub fn write_question(
    question: &Question,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    domain_name::write(&question.name, buf, ptr_offsets, "question.name")?;

    let bits = QuestionBits {
        resource_type: match question.resource_type {
            IntEnum::Enum(e) => e as u16,
            IntEnum::Unknown(i) => i,
        },
        resource_class: match question.resource_class {
            IntEnum::Enum(e) => e as u16,
            IntEnum::Unknown(i) => i,
        },
    }
    .pack()?;
    buf.reserve(bits.len());
    buf.put_slice(&bits);
    Ok(())
}

pub fn read_question(buf: &[u8], offset: &mut usize) -> Result<Option<Question>> {
    let (name_bytes_consumed, name_str) = domain_name::read(buf, *offset, "question.name")?;

    let bits_size = 4; // size of QuestionBits in bytes
    if buf.len() < *offset + name_bytes_consumed + bits_size {
        // Not enough bytes for the question header, wait for more bytes.
        return Ok(None);
    }
    let bits = QuestionBits::unpack_from_slice(
        &buf[*offset + name_bytes_consumed..*offset + name_bytes_consumed + bits_size],
    )
    .with_context(|| "couldn't unpack question bits")?;

    *offset += name_bytes_consumed + bits_size;

    Ok(Some(Question {
        name: name_str,
        resource_type: match enums_generated::resourcetype_int(bits.resource_type as usize) {
            Some(e) => IntEnum::Enum(e),
            None => IntEnum::Unknown(bits.resource_type),
        },
        resource_class: match enums_generated::resourceclass_int(bits.resource_class as usize) {
            Some(e) => IntEnum::Enum(e),
            None => IntEnum::Unknown(bits.resource_class),
        },
    }))
}

// Resource, from RFC1035 section 4.1.3
//
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                                               |
// /                      NAME                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      TYPE                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                     CLASS                     |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                      TTL                      |
// |                                               |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                   RDLENGTH                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--|
// /                     RDATA                     /
// /                                               /
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+

// Just the initial type field, for reading out and deciding if we're OPT or not.
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceTypeBits {
    // name: variable length string
    #[packed_field(bits = "0:15")]
    resource_type: u16,
}

// The regular Class/TTL fields used by non-OPT records
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceClassTTLBits {
    #[packed_field(bits = "0:15")]
    class: u16,
    #[packed_field(bits = "16:47")]
    ttl: u32,
}

// OPT-specific Class/TTL packing, from RFC6891 section 6.1.3:
//                                 1  1  1  1  1  1
//   0  1  2  3  4  5  6  7  8  9  0  1  2  3  4  5
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |                   UDP-SIZE                    |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |     EXTENDED-RCODE    |        VERSION        |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// |DO|                    Z                       |
// +--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
// Mirrors ResourceClassTTLBits with OPT packing
// Members are public to allow access by rdata.rs
#[derive(Default, PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceClassTTLOPTBits {
    // Class field
    #[packed_field(bits = "0:15")]
    pub udp_size: u16,
    // TTL field
    #[packed_field(bits = "16:23")]
    pub response_code: u8,
    #[packed_field(bits = "24:31")]
    pub version: u8,
    #[packed_field(bits = "32:32")]
    pub dnssec_ok: bool,
    #[packed_field(bits = "33:47")]
    _reserved: ReservedZero<packed_bits::Bits<15>>,
}

// The RDLength field, initialized to zero then updated after writing the RData
#[derive(PackedStruct)]
#[packed_struct(endian = "msb", bit_numbering = "msb0")]
pub struct ResourceRdataLengthBits {
    #[packed_field(bits = "0:15")]
    rdata_len: u16,
}

pub enum RDataFields<'a> {
    RDATA(&'a ResourceData),
    IP(IpAddr),
    SOA(rdata::SOAFields<'a>),
    TXT(&'a String),
}

/// Values to be used when writing a Resource, instead of using what's in the cached version.
/// Used when we're constructing our own custom response to something.
pub struct ResourceFields<'a> {
    pub name: &'a str,
    pub resource_type: IntEnum<u16, ResourceType>,
    pub resource_class: IntEnum<u16, ResourceClass>,
    pub ttl: u32,
    pub rdata: RDataFields<'a>,
}

/// Encodes a DNS resource to the provided buf
pub fn write_resource(
    resource: &Resource,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    write_resource_fields(
        &ResourceFields {
            name: resource.name.as_str(),
            resource_type: resource.resource_type,
            resource_class: resource.resource_class,
            ttl: resource.ttl,
            rdata: RDataFields::RDATA(&resource.rdata),
        },
        buf,
        ptr_offsets,
    )
}

/// Encodes a resource fields to the provided buf, with support for overriding field values
pub fn write_resource_fields(
    resource: &ResourceFields,
    buf: &mut BytesMut,
    ptr_offsets: &mut domain_name::LabelOffsets,
) -> Result<()> {
    domain_name::write(resource.name, buf, ptr_offsets, "resource.name")?;

    let type_bits = ResourceTypeBits {
        resource_type: match resource.resource_type {
            IntEnum::Enum(e) => e as u16,
            IntEnum::Unknown(i) => i,
        },
    }
    .pack()?;

    let classttl_bits = ResourceClassTTLBits {
        class: match resource.resource_class {
            IntEnum::Enum(e) => e as u16,
            IntEnum::Unknown(i) => i,
        },
        ttl: resource.ttl,
    }
    .pack()?;

    let rdlen_zero_bits = ResourceRdataLengthBits {
        rdata_len: 0 as u16, // To be updated below
    }
    .pack()?;

    buf.reserve(type_bits.len() + classttl_bits.len() + rdlen_zero_bits.len());
    buf.put_slice(&type_bits);
    buf.put_slice(&classttl_bits);
    let rdata_len_offset = buf.len();
    buf.put_slice(&rdlen_zero_bits);
    let rdata_offset = buf.len();

    // We write directly to 'buf' to ensure that any domain-name offsets are relative to the start
    // of the full 'buf'. Afterwards we go back and update the rdata_len to reflect the size of
    // what was written.
    match &resource.rdata {
        RDataFields::RDATA(rdata) => {
            rdata::write_rdata(&resource.resource_type, &rdata, buf, ptr_offsets)?
        }
        RDataFields::IP(IpAddr::V4(ip)) => rdata::write_a_ip(&ip, buf)?,
        RDataFields::IP(IpAddr::V6(ip)) => rdata::write_aaaa_ip(&ip, buf)?,
        RDataFields::SOA(soa) => rdata::write_soa_fields(&soa, buf, ptr_offsets)?,
        RDataFields::TXT(txt) => rdata::write_txt_entry(txt.as_bytes(), buf)?,
    }

    if buf.len() > rdata_offset {
        // We wrote some rdata, so go back and update the rdata length
        let rdlen_updated_bits = ResourceRdataLengthBits {
            rdata_len: u16::try_from(buf.len() - rdata_offset)
                .with_context(|| "rdata.length doesn't fit")?,
        }
        .pack()?;
        let buf_len = buf.len();
        let bits = buf
            .get_mut(rdata_len_offset..rdata_offset)
            .with_context(|| {
                format!(
                    "failed to get bytes {}..{} from buffer with length {}",
                    rdata_len_offset, rdata_offset, buf_len
                )
            })?;
        // Overwrite the previously appended rdlen_zero_bits with rdlen_updated_bits
        for i in 0..(rdata_offset - rdata_len_offset) {
            bits[i] = rdlen_updated_bits[i];
        }
    }

    Ok(())
}

/// Reads a resource in the common case when OPT records are not expected.
pub fn read_resource_non_opt<'a>(buf: &[u8], offset: &mut usize) -> Result<Option<Resource>> {
    match read_resource_name_type(buf, offset).context("Failed to read non-OPT resource")? {
        Some((_name, OPT_RESOURCE_TYPE)) => bail!("Got OPT resource in unexpected part of message"),
        Some((name, resource_type)) => {
            match read_resource_remainder(buf, offset, name, resource_type)? {
                Some(resource) => Ok(Some(resource)),
                None => Ok(None),
            }
        }
        None => Ok(None),
    }
}

/// Reads the resource name and type from a raw DNS message, or None if not enough data is available.
pub fn read_resource_name_type(buf: &[u8], offset: &mut usize) -> Result<Option<(String, u16)>> {
    let (mut total_bytes_consumed, name_str) = domain_name::read(buf, *offset, "resource.name")?;

    // Get the type field first: Determine whether this is an OPT record
    let type_bits_size = 2; // size of ResourceTypeBits in bytes
    if buf.len() < *offset + total_bytes_consumed + type_bits_size {
        // Not enough bytes for the resource type field, wait for more bytes.
        return Ok(None);
    }
    let resource_type = ResourceTypeBits::unpack_from_slice(
        &buf[*offset + total_bytes_consumed..*offset + total_bytes_consumed + type_bits_size],
    )
    .with_context(|| "couldn't unpack resource prelude bits")?
    .resource_type;
    total_bytes_consumed += type_bits_size;

    *offset += total_bytes_consumed;

    Ok(Some((name_str, resource_type)))
}

/// Reads the rest of an OPT resource from a raw DNS message.
/// This should be called after read_resource_name_type(), which provides the type of resource being read.
pub fn read_resource_remainder_opt(buf: &[u8], offset: &mut usize) -> Result<Option<OPT>> {
    // OPT/EDNS(0) hax: Custom packing with OPT-specific fields
    let opt_class_ttl_bits_size = 6; // size of ResourceClassTTLOPTBits in bytes
    if buf.len() < *offset + opt_class_ttl_bits_size {
        // Not enough bytes for the OPT class/ttl fields, wait for more bytes.
        return Ok(None);
    }
    let opt_class_ttl_bits = ResourceClassTTLOPTBits::unpack_from_slice(
        &buf[*offset..*offset + opt_class_ttl_bits_size],
    )
    .with_context(|| "couldn't unpack OPT resource custom class/ttl bits")?;
    let mut total_bytes_consumed = opt_class_ttl_bits_size;

    let rdata_location = read_rdata_location(
        buf,
        *offset + total_bytes_consumed,
        &mut total_bytes_consumed,
    )
    .context("Failed to read OPT rdata location")?;
    match rdata_location {
        Some((rdata_offset, rdata_len)) => {
            // Pass the OPT-specific data replacing the class/TTL bits, to be stored in the OPT object
            let opt = rdata::read_opt(buf, rdata_offset, rdata_len, opt_class_ttl_bits)?;
            total_bytes_consumed += rdata_len;
            *offset += total_bytes_consumed;
            Ok(Some(opt))
        }
        None => {
            // Not enough bytes for the size info (or the expected rdata), wait for more bytes.
            Ok(None)
        }
    }
}

/// Reads the rest of a non-OPT resource.
/// This should be called after read_resource_name_type(), which provides the type of resource being read.
pub fn read_resource_remainder(
    buf: &[u8],
    offset: &mut usize,
    name: String,
    resource_type: u16,
) -> Result<Option<Resource>> {
    // Non-OPT resource: regular class and TTL fields
    let class_ttl_bits_size = 6; // size of ResourceClassTTLBits in bytes
    if buf.len() < *offset + class_ttl_bits_size {
        // Not enough bytes for the class/ttl fields, wait for more bytes.
        return Ok(None);
    }
    let class_ttl_bits =
        ResourceClassTTLBits::unpack_from_slice(&buf[*offset..*offset + class_ttl_bits_size])
            .with_context(|| "couldn't unpack resource class/ttl bits")?;
    let mut total_bytes_consumed = class_ttl_bits_size;

    let rdata_location = read_rdata_location(
        buf,
        *offset + total_bytes_consumed,
        &mut total_bytes_consumed,
    )
    .with_context(|| {
        format!(
            "Failed to read {:?} resource rdata location for {}",
            resource_type, name
        )
    })?;
    match rdata_location {
        Some((rdata_offset, rdata_len)) => {
            let rdata = rdata::read_rdata(buf, resource_type, rdata_offset, rdata_len)
                .with_context(|| {
                    format!("Failed to read {:?} resource for {}", resource_type, name)
                })?;
            total_bytes_consumed += rdata_len;
            *offset += total_bytes_consumed;
            Ok(Some(Resource {
                name,
                resource_type: match enums_generated::resourcetype_int(resource_type as usize) {
                    Some(r) => IntEnum::Enum(r),
                    None => IntEnum::Unknown(resource_type),
                },
                resource_class: match enums_generated::resourceclass_int(
                    class_ttl_bits.class as usize,
                ) {
                    Some(r) => IntEnum::Enum(r),
                    None => IntEnum::Unknown(resource_type),
                },
                ttl: class_ttl_bits.ttl,
                rdata,
            }))
        }
        None => {
            // Not enough bytes for the size info (or the expected rdata), wait for more bytes.
            Ok(None)
        }
    }
}

/// Returns the offset and length of the rdata available in the buffer,
/// or None if the buffer is incomplete
fn read_rdata_location(
    buf: &[u8],
    offset: usize,
    total_bytes_consumed: &mut usize,
) -> Result<Option<(usize, usize)>> {
    let rdata_length_bits_size = 2; // size of ResourceRdataLengthBits in bytes
    let rdata_offset = offset + rdata_length_bits_size;
    if buf.len() < rdata_offset {
        // Not enough bytes for the rdata length field, wait for more bytes.
        return Ok(None);
    }
    let rdata_length_bits =
        ResourceRdataLengthBits::unpack_from_slice(&buf[offset..offset + rdata_length_bits_size])
            .with_context(|| "couldn't unpack rdata length bits")?;

    let rdata_len = rdata_length_bits.rdata_len as usize;
    if buf.len() < rdata_offset + rdata_len {
        // Not enough bytes for the rdata content, wait for more bytes.
        return Ok(None);
    }

    *total_bytes_consumed += rdata_length_bits_size;
    Ok(Some((rdata_offset, rdata_len)))
}

pub fn write_opt(opt: &OPT, udp_size_override: Option<u16>, buf: &mut BytesMut) -> Result<()> {
    domain_name::write_nopointer(".", buf, "opt.name")?;

    let type_bits = ResourceTypeBits {
        resource_type: ResourceType::OPT as u16,
    }
    .pack()?;

    let classttl_bits = ResourceClassTTLOPTBits {
        udp_size: match udp_size_override {
            Some(udp_size) => udp_size,
            None => opt.udp_size,
        },
        response_code: opt.response_code,
        version: opt.response_code,
        dnssec_ok: opt.dnssec_ok,
        ..ResourceClassTTLOPTBits::default() // for setting the _reserved field
    }
    .pack()?;

    let rdlen_zero_bits = ResourceRdataLengthBits {
        rdata_len: 0 as u16, // To be updated below
    }
    .pack()?;

    buf.reserve(type_bits.len() + classttl_bits.len() + rdlen_zero_bits.len());
    buf.put_slice(&type_bits);
    buf.put_slice(&classttl_bits);
    let rdata_len_offset = buf.len();
    buf.put_slice(&rdlen_zero_bits);
    let rdata_offset = buf.len();

    rdata::write_opt(opt, buf)?;

    if buf.len() > rdata_offset {
        // We wrote some rdata, so go back and update the rdata length
        let rdlen_updated_bits = ResourceRdataLengthBits {
            rdata_len: u16::try_from(buf.len() - rdata_offset)
                .with_context(|| "rdata.length doesn't fit")?,
        }
        .pack()?;
        let buf_len = buf.len();
        let bits = buf
            .get_mut(rdata_len_offset..rdata_offset)
            .with_context(|| {
                format!(
                    "failed to get bytes {}..{} from buffer with length {}",
                    rdata_len_offset, rdata_offset, buf_len
                )
            })?;
        // Overwrite the previously appended rdlen_zero_bits with rdlen_updated_bits
        for i in 0..(rdata_offset - rdata_len_offset) {
            bits[i] = rdlen_updated_bits[i];
        }
    }
    Ok(())
}

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

    #[test]
    fn test_update_message_id() {
        let mut buf = BytesMut::with_capacity(5);
        buf.resize(5, 0);
        assert_eq!(5, buf.len(), "{:?}", buf);

        // test value: upper=1 (256>>8), lower=255
        let id = 511;

        let mut res = update_message_id(id, &mut buf, 0);
        assert!(res.is_ok(), "{:?}", res);
        assert_eq!(1, buf[0]);
        assert_eq!(255, buf[1]);
        assert_eq!(0, buf[2]);
        assert_eq!(0, buf[3]);
        assert_eq!(0, buf[4]);

        buf.clear();
        buf.resize(5, 0);

        res = update_message_id(id, &mut buf, 1);
        assert!(res.is_ok(), "{:?}", res);
        assert_eq!(0, buf[0]);
        assert_eq!(1, buf[1]);
        assert_eq!(255, buf[2]);
        assert_eq!(0, buf[3]);
        assert_eq!(0, buf[4]);

        buf.clear();
        buf.resize(5, 0);

        res = update_message_id(id, &mut buf, 2);
        assert!(res.is_ok(), "{:?}", res);
        assert_eq!(0, buf[0]);
        assert_eq!(0, buf[1]);
        assert_eq!(1, buf[2]);
        assert_eq!(255, buf[3]);
        assert_eq!(0, buf[4]);

        buf.clear();
        buf.resize(5, 0);

        res = update_message_id(id, &mut buf, 3);
        assert!(res.is_ok(), "{:?}", res);
        assert_eq!(0, buf[0]);
        assert_eq!(0, buf[1]);
        assert_eq!(0, buf[2]);
        assert_eq!(1, buf[3]);
        assert_eq!(255, buf[4]);

        for offset in 4..10 {
            buf.clear();
            buf.resize(5, 0);

            res = update_message_id(id, &mut buf, offset);
            assert!(res.is_err(), "offset={} {:?}", offset, res);
            assert_eq!(0, buf[0]);
            assert_eq!(0, buf[1]);
            assert_eq!(0, buf[2]);
            assert_eq!(0, buf[3]);
            assert_eq!(0, buf[4]);
        }
    }
}