internet 0.1.0

//! UDP header encoding following [RFC 768].
//!
//! Encoding is supported for the following structures:
//!
//!  - [`Port`]
//!  - [`Checksum`]
//!  - [`Ipv4PseudoHeader`]
//!  - [`Ipv6PseudoHeader`]
//!  - [`Header`]
//!
//! [RFC 768]: https://datatracker.ietf.org/doc/html/rfc768

use crate::{
    Buf, BufMut, BufResult, Codec, Cursor,
    ietf::{ipv4, transport},
};

/// A UDP header following [RFC 768].
///
/// [RFC 768]: https://datatracker.ietf.org/doc/html/rfc768
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Header {
    /// The source port.
    pub source_port: Port,
    /// The destination port.
    pub destination_port: Port,
    /// The length of the header and data octets.
    pub length: u16,
    /// The checksum.
    pub checksum: Checksum,
}

impl Codec for Header {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        self.source_port.encode(writer, ())?;
        self.destination_port.encode(writer, ())?;
        self.length.encode(writer, ())?;
        self.checksum.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        Ok(Self {
            source_port: Port::decode(reader, ())?,
            destination_port: Port::decode(reader, ())?,
            length: u16::decode(reader, ())?,
            checksum: Checksum::decode(reader, ())?,
        })
    }
}

/// A UDP port number following [IANA service-names-port-numbers].
///
/// A 16-bit number used to identify a UDP endpoint on a host.
///
/// [IANA service-names-port-numbers]: https://www.iana.org/assignments/service-names-port-numbers
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct Port(pub transport::Port);

impl Port {
    /// DNS.
    pub const DNS: Self = Self(transport::Port(53));
    /// DHCP Server.
    pub const DHCP_SERVER: Self = Self(transport::Port(67));
    /// DHCP Client.
    pub const DHCP_CLIENT: Self = Self(transport::Port(68));
    /// TFTP.
    pub const TFTP: Self = Self(transport::Port(69));
    /// HTTP.
    pub const HTTP: Self = Self(transport::Port(80));
    /// NTP.
    pub const NTP: Self = Self(transport::Port(123));
    /// HTTPS.
    pub const HTTPS: Self = Self(transport::Port(443));

    /// Returns true if this is a system port (0-1023).
    pub const fn is_system(self) -> bool {
        self.0.is_system()
    }

    /// Returns true if this is a user port (1024-49151).
    pub const fn is_user(self) -> bool {
        self.0.is_user()
    }

    /// Returns true if this is a dynamic port (49152-65535).
    pub const fn is_dynamic(self) -> bool {
        self.0.is_dynamic()
    }
}

impl Default for Port {
    fn default() -> Self {
        Self(transport::Port::default())
    }
}

impl From<u16> for Port {
    #[inline]
    fn from(val: u16) -> Self {
        Port(transport::Port::new(val))
    }
}

impl From<Port> for u16 {
    #[inline]
    fn from(port: Port) -> Self {
        port.0.as_u16()
    }
}

impl From<transport::Port> for Port {
    #[inline]
    fn from(port: transport::Port) -> Self {
        Port(port)
    }
}

impl From<Port> for transport::Port {
    #[inline]
    fn from(port: Port) -> Self {
        port.0
    }
}

impl Codec for Port {
    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(transport::Port::decode(reader, ())?))
    }
}

/// A checksum of the UDP header following [RFC 768].
///
/// A 16-bit number used to detect data corruption.
///
/// [RFC 768]: https://datatracker.ietf.org/doc/html/rfc768
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct Checksum(pub u16);

impl Checksum {
    /// Calculates the checksum for IPv4.
    ///
    /// The checksum covers an IPv4 pseudo header, the UDP header, and the UDP payload.
    pub fn calculate_ipv4(
        pseudo_header: &Ipv4PseudoHeader,
        header: &Header,
        payload: &[u8],
    ) -> Self {
        let mut sum: u32 = 0;

        let source_ip_bytes: [u8; 4] = pseudo_header.source_address.into();
        sum += u16::from_be_bytes([source_ip_bytes[0], source_ip_bytes[1]]) as u32;
        sum += u16::from_be_bytes([source_ip_bytes[2], source_ip_bytes[3]]) as u32;
        let destination_ip_bytes: [u8; 4] = pseudo_header.destination_address.into();
        sum += u16::from_be_bytes([destination_ip_bytes[0], destination_ip_bytes[1]]) as u32;
        sum += u16::from_be_bytes([destination_ip_bytes[2], destination_ip_bytes[3]]) as u32;
        sum += 0x0011;
        sum += pseudo_header.udp_length as u32;

        sum += u16::from(header.source_port) as u32;
        sum += u16::from(header.destination_port) as u32;
        sum += header.length as u32;

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

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

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

        let checksum = !(sum as u16);

        Checksum(if checksum == 0 { 0xFFFF } else { checksum })
    }

    /// Calculates the checksum for IPv6.
    ///
    /// The checksum covers a pseudo header, the UDP header, and the UDP payload.
    /// For IPv6, the checksum is mandatory (unlike IPv4 where it's optional).
    pub fn calculate_ipv6(
        pseudo_header: &Ipv6PseudoHeader,
        header: &Header,
        payload: &[u8],
    ) -> Self {
        let mut sum: u32 = 0;

        for i in 0..8 {
            sum += u16::from_be_bytes([
                pseudo_header.source_ip[i * 2],
                pseudo_header.source_ip[i * 2 + 1],
            ]) as u32;
        }
        for i in 0..8 {
            sum += u16::from_be_bytes([
                pseudo_header.destination_ip[i * 2],
                pseudo_header.destination_ip[i * 2 + 1],
            ]) as u32;
        }

        sum += ((pseudo_header.udp_length >> 16) & 0xFFFF) as u32;
        sum += (pseudo_header.udp_length & 0xFFFF) as u32;

        sum += 0x0011;

        sum += u16::from(header.source_port) as u32;
        sum += u16::from(header.destination_port) as u32;
        sum += header.length as u32;

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

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

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

        let checksum = !(sum as u16);

        Checksum(checksum)
    }
}

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, ())?))
    }
}

/// A pseudo header for IPv4 checksum calculation.
///
/// Used in checksum computation as defined in [RFC 768].
///
/// [RFC 768]: https://datatracker.ietf.org/doc/html/rfc768
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Ipv4PseudoHeader {
    /// The source IPv4 address.
    pub source_address: ipv4::Address,
    /// The destination IPv4 address.
    pub destination_address: ipv4::Address,
    /// The UDP length (header + data).
    pub udp_length: u16,
}

/// A pseudo header for IPv6 checksum calculation.
///
/// Used in checksum computation as defined in [RFC 2460].
///
/// [RFC 2460]: https://datatracker.ietf.org/doc/html/rfc2460
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Ipv6PseudoHeader {
    /// The source IPv6 address.
    pub source_ip: [u8; 16],
    /// The destination IPv6 address.
    pub destination_ip: [u8; 16],
    /// The UDP length (header + data).
    pub udp_length: u32,
}

#[cfg(test)]
mod tests {
    use core::fmt::Debug;

    use super::{Checksum, Header, Ipv4PseudoHeader, Port};
    use crate::{Codec, Cursor, ietf::ipv4};

    fn codec_roundtrip<T: Codec<C> + Debug + Eq, C: Copy>(
        etalon_struct: T,
        etalon_bytes: &[u8],
        context: C,
    ) {
        let mut encoded_bytes = vec![];
        {
            let writer = &mut Cursor::new(&mut encoded_bytes);
            etalon_struct.encode(writer, context).unwrap();
        }
        assert_eq!(etalon_bytes, &encoded_bytes);

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

        encoded_bytes.fill(0x00);
        {
            let writer = &mut Cursor::new(&mut encoded_bytes);
            decoded_struct.encode(writer, context).unwrap();
        }
        assert_eq!(etalon_bytes, &encoded_bytes);
    }

    #[test]
    fn port() {
        let etalon_bytes = &[0x00, 0x35]; // Port 53 (DNS)
        let etalon_struct = Port::DNS;

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn header() {
        let etalon_bytes = &[
            0x00, 0x35, // source port: 53 (DNS)
            0xC3, 0x50, // destination port: 50000
            0x00, 0x10, // length: 16
            0x00, 0x00, // checksum: 0
        ];
        let etalon_struct = Header {
            source_port: Port::DNS,
            destination_port: Port::from(50000u16),
            length: 16,
            checksum: Checksum(0),
        };

        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn checksum_calculation() {
        // Simple test case for checksum calculation
        let pseudo_header = Ipv4PseudoHeader {
            source_address: ipv4::Address::from([192, 168, 1, 1]),
            destination_address: ipv4::Address::from([192, 168, 1, 2]),
            udp_length: 16,
        };

        let header = Header {
            source_port: Port::DNS,
            destination_port: Port::from(50000u16),
            length: 16,
            checksum: Checksum(0),
        };

        let payload = &[];

        let checksum = Checksum::calculate_ipv4(&pseudo_header, &header, payload);
        assert_ne!(checksum.0, 0); // Should not be zero for this case
    }
}