internet 0.1.0

//! IPv4 Header Mapping.
//!
//! Provides zero-copy read and write access to IPv4 header fields.
//!
//! As defined in [RFC 791].
//!
//! [IETF RFC 791]: https://datatracker.ietf.org/doc/html/rfc791

use crate::ietf::ip::Protocol;
use crate::ietf::ipv4::{Address, Options};
use crate::{Buf, BufMut};

/// A zero-copy mapping for an IPv4 header.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HeaderMapping<T> {
    buffer: T,
}

impl<T> HeaderMapping<T> {
    /// Creates a new header mapping.
    pub fn new(buffer: T) -> Self {
        Self { buffer }
    }

    /// Consumes the mapping and returns the underlying buffer.
    pub fn into_inner(self) -> T {
        self.buffer
    }
}

impl<T: Buf> HeaderMapping<T> {
    /// Reads the IP version.
    pub fn read_version(&self) -> u8 {
        unsafe { (self.buffer.get_u8_unchecked(0) >> 4) & 0x0F }
    }

    /// Reads the Internet Header Length (in 32-bit words).
    pub fn read_ihl(&self) -> u8 {
        unsafe { self.buffer.get_u8_unchecked(0) & 0x0F }
    }

    /// Reads the Differentiated Services Code Point.
    pub fn read_dscp(&self) -> u8 {
        unsafe { (self.buffer.get_u8_unchecked(1) >> 2) & 0x3F }
    }

    /// Reads the Explicit Congestion Notification.
    pub fn read_ecn(&self) -> u8 {
        unsafe { self.buffer.get_u8_unchecked(1) & 0x03 }
    }

    /// Reads the total length of the datagram.
    pub fn read_total_length(&self) -> u16 {
        unsafe { self.buffer.get_u16_be_unchecked(2) }
    }

    /// Reads the identification field.
    pub fn read_identification(&self) -> u16 {
        unsafe { self.buffer.get_u16_be_unchecked(4) }
    }

    /// Reads the Don't Fragment flag.
    pub fn read_df(&self) -> bool {
        unsafe { (self.buffer.get_u16_be_unchecked(6) & 0x4000) != 0 }
    }

    /// Reads the More Fragments flag.
    pub fn read_mf(&self) -> bool {
        unsafe { (self.buffer.get_u16_be_unchecked(6) & 0x2000) != 0 }
    }

    /// Reads the fragment offset.
    pub fn read_fragment_offset(&self) -> u16 {
        unsafe { self.buffer.get_u16_be_unchecked(6) & 0x1FFF }
    }

    /// Reads the Time to Live.
    pub fn read_ttl(&self) -> u8 {
        unsafe { self.buffer.get_u8_unchecked(8) }
    }

    /// Reads the protocol.
    pub fn read_protocol(&self) -> Protocol {
        unsafe { Protocol::from(self.buffer.get_u8_unchecked(9)) }
    }

    /// Reads the header checksum.
    pub fn read_checksum(&self) -> u16 {
        unsafe { self.buffer.get_u16_be_unchecked(10) }
    }

    /// Reads the source address.
    pub fn read_source_address(&self) -> Address {
        unsafe {
            Address::from([
                self.buffer.get_u8_unchecked(12),
                self.buffer.get_u8_unchecked(13),
                self.buffer.get_u8_unchecked(14),
                self.buffer.get_u8_unchecked(15),
            ])
        }
    }

    /// Reads the destination address.
    pub fn read_destination_address(&self) -> Address {
        unsafe {
            Address::from([
                self.buffer.get_u8_unchecked(16),
                self.buffer.get_u8_unchecked(17),
                self.buffer.get_u8_unchecked(18),
                self.buffer.get_u8_unchecked(19),
            ])
        }
    }

    /// Reads the IPv4 options.
    pub fn read_options(&self) -> Options {
        let ihl = self.read_ihl() as usize;
        let len = (ihl * 4).saturating_sub(20);
        let mut data = [0u8; 40];
        unsafe { self.buffer.read_into_unchecked(20, &mut data[..len]) };
        // len никогда не превысит 40, так как максимальный IHL = 15
        Options::try_from(&data[..len]).unwrap()
    }
}

impl<T: BufMut> HeaderMapping<T> {
    /// Writes the version and IHL.
    pub fn write_version_and_ihl(&mut self, version: u8, ihl: u8) {
        let val = ((version & 0x0F) << 4) | (ihl & 0x0F);
        unsafe { self.buffer.set_u8_unchecked(0, val) }
    }

    /// Writes the DSCP and ECN.
    pub fn write_dscp_and_ecn(&mut self, dscp: u8, ecn: u8) {
        let val = ((dscp & 0x3F) << 2) | (ecn & 0x03);
        unsafe { self.buffer.set_u8_unchecked(1, val) }
    }

    /// Writes the total length.
    pub fn write_total_length(&mut self, length: u16) {
        unsafe { self.buffer.set_u16_be_unchecked(2, length) }
    }

    /// Writes the identification.
    pub fn write_identification(&mut self, id: u16) {
        unsafe { self.buffer.set_u16_be_unchecked(4, id) }
    }

    /// Writes the flags and fragment offset.
    pub fn write_flags_and_offset(&mut self, df: bool, mf: bool, offset: u16) {
        let mut val = offset & 0x1FFF;
        if df {
            val |= 0x4000;
        }
        if mf {
            val |= 0x2000;
        }
        unsafe { self.buffer.set_u16_be_unchecked(6, val) }
    }

    /// Writes the TTL.
    pub fn write_ttl(&mut self, ttl: u8) {
        unsafe { self.buffer.set_u8_unchecked(8, ttl) }
    }

    /// Writes the protocol.
    pub fn write_protocol(&mut self, protocol: Protocol) {
        unsafe { self.buffer.set_u8_unchecked(9, u8::from(protocol)) }
    }

    /// Writes the checksum.
    pub fn write_checksum(&mut self, checksum: u16) {
        unsafe { self.buffer.set_u16_be_unchecked(10, checksum) }
    }

    /// Writes the source address.
    pub fn write_source_address(&mut self, addr: Address) {
        let octets = addr.octets();
        unsafe {
            self.buffer.set_u8_unchecked(12, octets[0]);
            self.buffer.set_u8_unchecked(13, octets[1]);
            self.buffer.set_u8_unchecked(14, octets[2]);
            self.buffer.set_u8_unchecked(15, octets[3]);
        }
    }

    /// Writes the destination address.
    pub fn write_destination_address(&mut self, addr: Address) {
        let octets = addr.octets();
        unsafe {
            self.buffer.set_u8_unchecked(16, octets[0]);
            self.buffer.set_u8_unchecked(17, octets[1]);
            self.buffer.set_u8_unchecked(18, octets[2]);
            self.buffer.set_u8_unchecked(19, octets[3]);
        }
    }

    /// Writes the IPv4 options from a byte slice.
    pub fn write_options(&mut self, options: &[u8]) {
        let ihl = self.read_ihl() as usize;
        let len = (ihl * 4).saturating_sub(20);
        unsafe {
            for i in 0..len {
                self.buffer.set_u8_unchecked(20 + i, options[i]);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{Address, HeaderMapping, Protocol};

    #[test]
    fn read_fields() {
        let buffer = [
            0x45, 0x00, 0x00, 0x34, 0x00, 0x00, 0x40, 0x00, 0x40, 0x11, 0x00, 0x00, 0xc0, 0xa8,
            0x00, 0x01, 0xc0, 0xa8, 0x00, 0x02,
        ];
        let mapping = HeaderMapping::new(&buffer[..]);
        assert_eq!(mapping.read_version(), 4);
        assert_eq!(mapping.read_ihl(), 5);
        assert_eq!(mapping.read_total_length(), 52);
        assert!(mapping.read_df());
        assert!(!mapping.read_mf());
        assert_eq!(mapping.read_ttl(), 64);
        assert_eq!(mapping.read_protocol(), Protocol::UDP);
        assert_eq!(mapping.read_source_address(), Address::new(192, 168, 0, 1));
    }

    #[test]
    fn write_fields() {
        let mut buffer = [0u8; 20];
        {
            let mut mapping = HeaderMapping::new(&mut buffer[..]);
            mapping.write_version_and_ihl(4, 5);
            mapping.write_total_length(52);
            mapping.write_flags_and_offset(true, false, 0);
            mapping.write_ttl(64);
            mapping.write_protocol(Protocol::UDP);
            mapping.write_source_address(Address::new(192, 168, 0, 1));
            mapping.write_destination_address(Address::new(192, 168, 0, 2));
        }
        assert_eq!(&buffer[0..2], &[0x45, 0x00]);
        assert_eq!(&buffer[12..16], &[0xc0, 0xa8, 0x00, 0x01]);
    }
}