internet 0.1.0

//! IPv4 header types and codecs.
//!
//! As defined in [RFC 791].
//!
//! [IETF RFC 791]: https://datatracker.ietf.org/doc/html/rfc791

use crate::ietf::{
    ip::{Protocol, Version},
    ipv4::{Address, Options},
};
use crate::{Buf, BufError, BufMut, BufResult, Codec, Cursor};

/// Differentiated Services Code Point (DSCP).
///
/// A 6-bit field used for packet classification.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Dscp(pub u8);

impl Dscp {
    /// Creates a new DSCP value.
    ///
    /// # Errors
    ///
    /// Returns [`BufError::UnexpectedValue`] if the value exceeds 6 bits (> 63).
    pub fn new(value: u8) -> BufResult<Self> {
        if value <= 63 {
            Ok(Self(value))
        } else {
            Err(BufError::UnexpectedValue)
        }
    }
}

/// Explicit Congestion Notification (ECN).
///
/// A 2-bit field used for end-to-end notification of network congestion.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum Ecn {
    /// Non-ECN-Capable Transport.
    NonEct = 0,
    /// ECN-Capable Transport (0).
    Ect0 = 1,
    /// ECN-Capable Transport (1).
    Ect1 = 2,
    /// Congestion Encountered.
    Ce = 3,
}

impl Ecn {
    /// Decodes a 2-bit ECN value.
    fn from_bits(bits: u8) -> BufResult<Self> {
        match bits & 0x03 {
            0 => Ok(Self::NonEct),
            1 => Ok(Self::Ect0),
            2 => Ok(Self::Ect1),
            3 => Ok(Self::Ce),
            _ => Err(BufError::UnexpectedValue),
        }
    }
}

/// Internet Header Length (IHL).
///
/// The length of the IPv4 header in 32-bit (4-byte) words.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Ihl(pub u8);

impl Ihl {
    /// Creates a new IHL value.
    ///
    /// # Errors
    ///
    /// Returns [`BufError::UnexpectedValue`] if the value is not between 5 and 15.
    ///
    /// *Note: The IHL field is strictly 4 bits in the IPv4 header, making 15
    /// (binary 1111) the absolute maximum value, which corresponds to a 60-byte header.*
    pub fn new(value: u8) -> BufResult<Self> {
        if (5..=15).contains(&value) {
            Ok(Self(value))
        } else {
            Err(BufError::UnexpectedValue)
        }
    }
}

impl PartialEq<u8> for Ihl {
    fn eq(&self, other: &u8) -> bool {
        self.0 == *other
    }
}

impl PartialEq<usize> for Ihl {
    fn eq(&self, other: &usize) -> bool {
        (self.0 as usize) == *other
    }
}

impl PartialOrd<u8> for Ihl {
    fn partial_cmp(&self, other: &u8) -> Option<core::cmp::Ordering> {
        self.0.partial_cmp(other)
    }
}

impl PartialOrd<usize> for Ihl {
    fn partial_cmp(&self, other: &usize) -> Option<core::cmp::Ordering> {
        (self.0 as usize).partial_cmp(other)
    }
}

/// An IPv4 header checksum.
///
/// Used to detect data corruption in the header, as defined in [RFC 791].
///
/// [RFC 791]: https://datatracker.ietf.org/doc/html/rfc791
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Checksum(pub u16);

impl Checksum {
    /// Calculates the IPv4 Header Checksum.
    ///
    /// The checksum field must be zeroed out before calculation.
    ///
    /// # Examples
    ///
    /// ```
    /// use internet::ipv4::Checksum;
    ///
    /// let mut header = [
    ///     0x45, 0x00, 0x00, 0x34,
    ///     0x00, 0x00, 0x00, 0x00,
    ///     0x40, 0x11, 0x00, 0x00,
    ///     0xc0, 0xa8, 0x00, 0x01,
    ///     0xc0, 0xa8, 0x00, 0x02,
    /// ];
    ///
    /// let checksum = Checksum::calculate(&header);
    /// header[10..12].copy_from_slice(&checksum.0.to_be_bytes());
    /// ```
    pub fn calculate(data: &[u8]) -> Self {
        let mut sum: u32 = 0;
        let mut i = 0;

        while i + 1 < data.len() {
            sum += u16::from_be_bytes([data[i], data[i + 1]]) as u32;
            i += 2;
        }

        if i < data.len() {
            sum += (data[i] as u32) << 8;
        }

        while sum >> 16 != 0 {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }

        Self(!(sum as u16))
    }
}

impl Codec for Checksum {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        self.0.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        Ok(Self(u16::decode(reader, ())?))
    }
}

/// An IPv4 header following [RFC 791].
///
/// [IETF RFC 791]: https://datatracker.ietf.org/doc/html/rfc791
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Header {
    /// Differentiated Services Code Point.
    pub dscp: Dscp,
    /// Explicit Congestion Notification.
    pub ecn: Ecn,
    /// Internet Header Length (in 32-bit words).
    pub ihl: Ihl,
    /// Total length of the datagram (header + data) in bytes.
    pub total_length: u16,
    /// Identification field for fragmentation.
    pub identification: u16,
    /// Don't Fragment flag.
    pub df: bool,
    /// More Fragments flag.
    pub mf: bool,
    /// Fragment offset (in 8-byte blocks).
    pub fragment_offset: u16,
    /// Time to Live.
    pub ttl: u8,
    /// The protocol of the payload.
    pub protocol: Protocol,
    /// The header checksum.
    pub checksum: Checksum,
    /// The source IPv4 address.
    pub source_address: Address,
    /// The destination IPv4 address.
    pub destination_address: Address,
    /// Optional fields and padding.
    pub options: Options,
}

impl Codec for Header {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        let version_ihl = ((Version::Ip as u8) << 4) | (self.ihl.0 & 0x0F);
        writer.write_u8(version_ihl)?;

        let dscp_ecn = (self.dscp.0 << 2) | (self.ecn as u8);
        writer.write_u8(dscp_ecn)?;

        writer.write_u16_be(self.total_length)?;
        writer.write_u16_be(self.identification)?;

        let mut flags_offset = self.fragment_offset & 0x1FFF;
        if self.df {
            flags_offset |= 0x4000;
        }
        if self.mf {
            flags_offset |= 0x2000;
        }
        writer.write_u16_be(flags_offset)?;

        writer.write_u8(self.ttl)?;
        self.protocol.encode(writer, ())?;
        self.checksum.encode(writer, ())?;

        self.source_address.encode(writer, ())?;
        self.destination_address.encode(writer, ())?;

        let options_len = ((self.ihl.0 as usize) * 4).saturating_sub(20);
        self.options.encode(writer, options_len)?;

        Ok(())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let version_ihl = reader.read_u8()?;
        let version = version_ihl >> 4;

        if version != (Version::Ip as u8) {
            return Err(BufError::UnexpectedValue);
        }

        let ihl = Ihl::new(version_ihl & 0x0F)?;

        let dscp_ecn = reader.read_u8()?;
        let dscp = Dscp::new(dscp_ecn >> 2)?;
        let ecn = Ecn::from_bits(dscp_ecn)?;

        let total_length = reader.read_u16_be()?;
        let identification = reader.read_u16_be()?;

        let flags_offset = reader.read_u16_be()?;
        let df = (flags_offset & 0x4000) != 0;
        let mf = (flags_offset & 0x2000) != 0;
        let fragment_offset = flags_offset & 0x1FFF;

        let ttl = reader.read_u8()?;
        let protocol = Protocol::decode(reader, ())?;
        let checksum = Checksum::decode(reader, ())?;

        let source_address = Address::decode(reader, ())?;
        let destination_address = Address::decode(reader, ())?;

        let options_len = ((ihl.0 as usize) * 4).saturating_sub(20);
        let options = Options::decode(reader, options_len)?;

        Ok(Self {
            dscp,
            ecn,
            ihl,
            total_length,
            identification,
            df,
            mf,
            fragment_offset,
            ttl,
            protocol,
            checksum,
            source_address,
            destination_address,
            options,
        })
    }
}

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

    #[test]
    fn ihl_comparisons() {
        let ihl = Ihl::new(5).unwrap();
        assert_eq!(ihl, 5u8);
        assert_eq!(ihl, 5usize);
        assert!(ihl < 6u8);
        assert!(ihl <= 5usize);
    }

    #[test]
    fn ihl_validation() {
        assert!(Ihl::new(4).is_err());
        assert!(Ihl::new(5).is_ok());
        assert!(Ihl::new(15).is_ok());
        assert!(Ihl::new(16).is_err());
    }

    #[test]
    fn header_roundtrip() {
        let etalon_bytes = &[
            0x45, 0x00, 0x00, 0x34, 0x00, 0x00, 0x40, 0x00, 0x40, 0x11, 0x00, 0x00, 0xc0, 0xa8,
            0x00, 0x01, 0xc0, 0xa8, 0x00, 0x02,
        ];

        let etalon_struct = Header {
            dscp: Dscp(0),
            ecn: Ecn::NonEct,
            ihl: Ihl(5),
            total_length: 52,
            identification: 0,
            df: true,
            mf: false,
            fragment_offset: 0,
            ttl: 64,
            protocol: Protocol::UDP,
            checksum: Checksum(0),
            source_address: Address::from([192, 168, 0, 1]),
            destination_address: Address::from([192, 168, 0, 2]),
            options: Options::try_from(&[][..]).unwrap(),
        };

        let mut encoded_bytes = vec![];
        {
            let writer = &mut Cursor::new(&mut encoded_bytes);
            etalon_struct.encode(writer, ()).unwrap();
        }
        assert_eq!(&etalon_bytes.to_vec(), &encoded_bytes);

        let decoded_struct = {
            let reader = &mut Cursor::new(&encoded_bytes);
            Header::decode(reader, ()).unwrap()
        };
        assert_eq!(etalon_struct, decoded_struct);
    }
}