zero-packet 0.1.0

use core::fmt;

/// The length of a Fragment extension header in bytes.
pub const FRAGMENT_HEADER_LENGTH: usize = 8;

/// Writes the Fragment header fields.
pub struct FragmentHeaderWriter<'a> {
    pub bytes: &'a mut [u8],
}

impl<'a> FragmentHeaderWriter<'a> {
    /// Creates a new `FragmentHeaderWriter` from the given data slice.
    #[inline]
    pub fn new(bytes: &'a mut [u8]) -> Result<Self, &'static str> {
        if bytes.len() < FRAGMENT_HEADER_LENGTH {
            panic!("Slice is too short to contain a Fragment header.");
        }

        Ok(Self { bytes })
    }

    /// Returns the length of the Fragment header.
    #[inline]
    pub fn header_len(&self) -> usize {
        FRAGMENT_HEADER_LENGTH
    }

    /// Sets the next header field.
    ///
    /// Identifies the type of the next header.
    #[inline]
    pub fn set_next_header(&mut self, next_header: u8) {
        self.bytes[0] = next_header;
    }

    /// Sets the reserved field.
    ///
    /// Should be all zeroes.
    #[inline]
    pub fn set_reserved(&mut self, reserved: u8) {
        self.bytes[1] = reserved;
    }

    /// Sets the fragment offset field.
    ///
    /// Offset in 8 bytes relative to the start of the fragmentable part of the original packet.
    #[inline]
    pub fn set_fragment_offset(&mut self, fragment_offset: u16) {
        // All 8 bits of the 2nd byte and the first 5 bits of the 3rd byte.
        let value = fragment_offset & 0x1FFF;
        self.bytes[2] = (value >> 5) as u8;
        self.bytes[3] = (self.bytes[3] & 0xE0) | ((value & 0x1F) as u8);
    }

    /// Sets the res field.
    ///
    /// Reserved and should be all zeroes.
    #[inline]
    pub fn set_res(&mut self, res: u8) {
        // Bits 5-6 of the 3rd byte.
        let value = res & 0b11;
        self.bytes[3] = (self.bytes[3] & 0x9F) | (value << 5);
    }

    /// Sets the M flag.
    ///
    /// 1 means more fragments follow, 0 means this is the last fragment.
    #[inline]
    pub fn set_m_flag(&mut self, m_flag: bool) {
        // Bit 7 of the 3rd byte.
        if m_flag {
            self.bytes[3] |= 0x80;
        } else {
            self.bytes[3] &= 0x7F;
        }
    }

    /// Sets the identification field.
    ///
    /// Identifies the packet, generated by the sender, to help the receiver reassemble the fragments.
    #[inline]
    pub fn set_identification(&mut self, identification: u32) {
        self.bytes[4] = (identification >> 24) as u8;
        self.bytes[5] = (identification >> 16) as u8;
        self.bytes[6] = (identification >> 8) as u8;
        self.bytes[7] = identification as u8;
    }
}

/// Reads the Fragment header fields.
pub struct FragmentHeaderReader<'a> {
    pub bytes: &'a [u8],
}

impl<'a> FragmentHeaderReader<'a> {
    #[inline]
    pub fn new(bytes: &'a [u8]) -> Result<Self, &'static str> {
        if bytes.len() < FRAGMENT_HEADER_LENGTH {
            return Err("Slice is too short to contain a Fragment header.");
        }

        Ok(Self { bytes })
    }

    /// Returns the next header field.
    ///
    /// Identifies the type of the next header.
    #[inline]
    pub fn next_header(&self) -> u8 {
        self.bytes[0]
    }

    /// Returns the reserved field.
    ///
    /// Should be all zeroes.
    #[inline]
    pub fn reserved(&self) -> u8 {
        self.bytes[1]
    }

    /// Returns the fragment offset field.
    ///
    /// Offset in 8 bytes relative to the start of the fragmentable part of the original packet.
    #[inline]
    pub fn fragment_offset(&self) -> u16 {
        ((self.bytes[2] as u16) << 5) | ((self.bytes[3] & 0x1F) as u16)
    }

    /// Returns the res field.
    ///
    /// Reserved and should be all zeroes.
    #[inline]
    pub fn res(&self) -> u8 {
        (self.bytes[3] >> 5) & 0b11
    }

    /// Returns the M flag.
    ///
    /// 1 means more fragments follow, 0 means this is the last fragment.
    #[inline]
    pub fn m_flag(&self) -> bool {
        (self.bytes[3] & 0x80) != 0
    }

    /// Returns the identification field.
    ///
    /// Identifies the packet, generated by the sender, to help the receiver reassemble the fragments.
    #[inline]
    pub fn identification(&self) -> u32 {
        ((self.bytes[4] as u32) << 24)
            | ((self.bytes[5] as u32) << 16)
            | ((self.bytes[6] as u32) << 8)
            | (self.bytes[7] as u32)
    }

    /// Returns the length of the Fragment header.
    #[inline]
    pub fn header_len(&self) -> usize {
        FRAGMENT_HEADER_LENGTH
    }

    /// Returns a reference to the header.
    #[inline]
    pub fn header(&self) -> &'a [u8] {
        &self.bytes[..FRAGMENT_HEADER_LENGTH]
    }

    /// Returns a reference to the payload.
    #[inline]
    pub fn payload(&self) -> &'a [u8] {
        &self.bytes[FRAGMENT_HEADER_LENGTH..]
    }
}

impl fmt::Debug for FragmentHeaderReader<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("FragmentHeader")
            .field("next_header", &self.next_header())
            .field("reserved", &self.reserved())
            .field("fragment_offset", &self.fragment_offset())
            .field("res", &self.res())
            .field("m_flag", &self.m_flag())
            .field("identification", &self.identification())
            .finish()
    }
}

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

    #[test]
    fn getters_and_setters() {
        // Raw packet.
        let mut bytes = [0u8; 16];

        // Random values.
        let next_header = 6;
        let reserved = 0;
        let fragment_offset = 255;
        let res = 0;
        let m_flag = true;
        let identification = 0x04050607;

        // Create a new Fragment header.
        let mut writer = FragmentHeaderWriter::new(&mut bytes).unwrap();

        // Set the fields.
        writer.set_next_header(next_header);
        writer.set_reserved(reserved);
        writer.set_fragment_offset(fragment_offset);
        writer.set_res(res);
        writer.set_m_flag(m_flag);
        writer.set_identification(identification);

        // Create a Fragment header reader.
        let reader = FragmentHeaderReader::new(&bytes).unwrap();

        // Check the fields.
        assert_eq!(reader.next_header(), next_header);
        assert_eq!(reader.reserved(), reserved);
        assert_eq!(reader.fragment_offset(), fragment_offset);
        assert_eq!(reader.res(), res);
        assert_eq!(reader.m_flag(), m_flag);
        assert_eq!(reader.identification(), identification);
    }
}