internet 0.0.5

//! TCP Header encoding.
//!
//! [RFC 9293]: https://datatracker.ietf.org/doc/html/rfc9293

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

/// TCP header following [Section 3.1].
///
/// [Section 3.1]: https://datatracker.ietf.org/doc/html/rfc9293#section-3.1
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Header {
    /// Source port.
    pub source_port: Port,
    /// Destination port.
    pub destination_port: Port,
    /// Sequence number.
    pub sequence_number: u32,
    /// Acknowledgment number.
    pub acknowledgment_number: u32,
    /// Data offset in 32-bit words.
    pub data_offset: u8,
    /// ECN-nonce concealment protection flag.
    pub ns: bool,
    /// Congestion Window Reduced flag.
    pub cwr: bool,
    /// ECN-Echo flag.
    pub ece: bool,
    /// Urgent pointer flag.
    pub urg: bool,
    /// Acknowledgment flag.
    pub ack: bool,
    /// Push function flag.
    pub psh: bool,
    /// Reset connection flag.
    pub rst: bool,
    /// Synchronize sequence numbers flag.
    pub syn: bool,
    /// No more data from sender flag.
    pub fin: bool,
    /// Window size.
    pub window: u16,
    /// Checksum.
    pub checksum: Checksum,
    /// Urgent pointer.
    pub urgent_pointer: u16,
}

impl Header {
    /// Packs the data offset, reserved, NS and control flags into the 16-bit field at bytes 12-13.
    fn pack_flags(&self) -> u16 {
        let mut flags: u16 = 0;
        flags |= ((self.data_offset as u16) & 0x0F) << 12;
        if self.ns {
            flags |= 0x0100;
        }
        if self.cwr {
            flags |= 0x0080;
        }
        if self.ece {
            flags |= 0x0040;
        }
        if self.urg {
            flags |= 0x0020;
        }
        if self.ack {
            flags |= 0x0010;
        }
        if self.psh {
            flags |= 0x0008;
        }
        if self.rst {
            flags |= 0x0004;
        }
        if self.syn {
            flags |= 0x0002;
        }
        if self.fin {
            flags |= 0x0001;
        }
        flags
    }

    /// Unpacks the 16-bit field at bytes 12-13 into data offset, NS and control flags.
    fn unpack_flags(val: u16) -> (u8, bool, bool, bool, bool, bool, bool, bool, bool, bool) {
        let data_offset = ((val >> 12) & 0x0F) as u8;
        let ns = (val & 0x0100) != 0;
        let cwr = (val & 0x0080) != 0;
        let ece = (val & 0x0040) != 0;
        let urg = (val & 0x0020) != 0;
        let ack = (val & 0x0010) != 0;
        let psh = (val & 0x0008) != 0;
        let rst = (val & 0x0004) != 0;
        let syn = (val & 0x0002) != 0;
        let fin = (val & 0x0001) != 0;
        (data_offset, ns, cwr, ece, urg, ack, psh, rst, syn, fin)
    }
}

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.sequence_number.encode(writer, ())?;
        self.acknowledgment_number.encode(writer, ())?;
        self.pack_flags().encode(writer, ())?;
        self.window.encode(writer, ())?;
        self.checksum.encode(writer, ())?;
        self.urgent_pointer.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let source_port = Port::decode(reader, ())?;
        let destination_port = Port::decode(reader, ())?;
        let sequence_number = u32::decode(reader, ())?;
        let acknowledgment_number = u32::decode(reader, ())?;
        let flags = u16::decode(reader, ())?;
        let (data_offset, ns, cwr, ece, urg, ack, psh, rst, syn, fin) = Self::unpack_flags(flags);
        let window = u16::decode(reader, ())?;
        let checksum = Checksum::decode(reader, ())?;
        let urgent_pointer = u16::decode(reader, ())?;

        Ok(Self {
            source_port,
            destination_port,
            sequence_number,
            acknowledgment_number,
            data_offset,
            ns,
            cwr,
            ece,
            urg,
            ack,
            psh,
            rst,
            syn,
            fin,
            window,
            checksum,
            urgent_pointer,
        })
    }
}

use std::fmt;
use std::num::ParseIntError;
use std::str::FromStr;

use crate::transport::Port as BasePort;

/// A TCP port number following [IANA service-names-port-numbers].
///
/// A 16-bit number used to identify a TCP 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 BasePort);

impl Port {
    /// Creates a new TCP port.
    pub const fn new(port: u16) -> Self {
        Port(BasePort::new(port))
    }

    /// Returns the port number as u16.
    pub const fn as_u16(self) -> u16 {
        self.0.as_u16()
    }

    /// 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()
    }

    /// FTP.
    pub const FTP: Self = Self(BasePort(21));
    /// SSH.
    pub const SSH: Self = Self(BasePort(22));
    /// Telnet.
    pub const TELNET: Self = Self(BasePort(23));
    /// SMTP.
    pub const SMTP: Self = Self(BasePort(25));
    /// DNS.
    pub const DNS: Self = Self(BasePort(53));
    /// HTTP.
    pub const HTTP: Self = Self(BasePort(80));
    /// POP3.
    pub const POP3: Self = Self(BasePort(110));
    /// NTP.
    pub const NTP: Self = Self(BasePort(123));
    /// IMAP.
    pub const IMAP: Self = Self(BasePort(143));
    /// BGP.
    pub const BGP: Self = Self(BasePort(179));
    /// HTTPS.
    pub const HTTPS: Self = Self(BasePort(443));
    /// SMB.
    pub const SMB: Self = Self(BasePort(445));
    /// SMTPS.
    pub const SMTPS: Self = Self(BasePort(465));
    /// SMTP Submission.
    pub const SMTP_SUBMISSION: Self = Self(BasePort(587));
    /// Syslog.
    pub const SYSLOG: Self = Self(BasePort(514));
    /// RTSP.
    pub const RTSP: Self = Self(BasePort(554));
    /// MySQL.
    pub const MYSQL: Self = Self(BasePort(3306));
    /// RDP.
    pub const RDP: Self = Self(BasePort(3389));
    /// PostgreSQL.
    pub const POSTGRES: Self = Self(BasePort(5432));
}

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

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

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

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

impl fmt::Display for Port {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.0.as_u16())
    }
}

impl FromStr for Port {
    type Err = ParseIntError;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        s.parse::<u16>().map(Port::new)
    }
}

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(BasePort::decode(reader, ())?))
    }
}

/// TCP checksum.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(transparent)]
pub struct Checksum(pub u16);

impl Checksum {
    /// Calculates the TCP checksum for IPv4.
    pub fn calculate_ipv4(pseudo_header: &Ipv4PseudoHeader, tcp_segment: &[u8]) -> Self {
        let mut sum: u32 = 0;

        sum += u16::from_be_bytes([pseudo_header.source_ip[0], pseudo_header.source_ip[1]]) as u32;
        sum += u16::from_be_bytes([pseudo_header.source_ip[2], pseudo_header.source_ip[3]]) as u32;
        sum += u16::from_be_bytes([
            pseudo_header.destination_ip[0],
            pseudo_header.destination_ip[1],
        ]) as u32;
        sum += u16::from_be_bytes([
            pseudo_header.destination_ip[2],
            pseudo_header.destination_ip[3],
        ]) as u32;
        sum += 0x0006; // TCP protocol number
        sum += pseudo_header.tcp_length as u32;

        let mut i = 0;
        while i + 1 < tcp_segment.len() {
            sum += u16::from_be_bytes([tcp_segment[i], tcp_segment[i + 1]]) as u32;
            i += 2;
        }
        if i < tcp_segment.len() {
            sum += (tcp_segment[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 TCP checksum for IPv6.
    pub fn calculate_ipv6(pseudo_header: &Ipv6PseudoHeader, tcp_segment: &[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.tcp_length >> 16) & 0xFFFF) as u32;
        sum += (pseudo_header.tcp_length & 0xFFFF) as u32;
        sum += 0x0006;

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

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

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

/// TCP pseudo header for IPv4 checksum calculation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Ipv4PseudoHeader {
    /// Source IPv4 address.
    pub source_ip: [u8; 4],
    /// Destination IPv4 address.
    pub destination_ip: [u8; 4],
    /// TCP segment length.
    pub tcp_length: u16,
}

/// TCP pseudo header for IPv6 checksum calculation.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Ipv6PseudoHeader {
    /// Source IPv6 address.
    pub source_ip: [u8; 16],
    /// Destination IPv6 address.
    pub destination_ip: [u8; 16],
    /// TCP segment length.
    pub tcp_length: u32,
}

#[cfg(test)]
mod tests {
    use super::{Checksum, Header, Port};
    use crate::{Codec, Cursor};
    use core::fmt::Debug;

    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);
    }

    #[test]
    fn header() {
        let etalon_bytes = [
            0x00, 0x50, // src port: 80
            0xC0, 0x00, // dst port: 49152
            0xDE, 0xAD, 0xBE, 0xEF, // seq
            0xCA, 0xFE, 0xBA, 0xBE, // ack
            0x50, 0x12, // data_offset=5, SYN+ACK
            0xFF, 0xFF, // window
            0x12, 0x34, // checksum
            0x00, 0x00, // urgent pointer
        ];
        let etalon_struct = Header {
            source_port: Port::HTTP,
            destination_port: Port::new(49152),
            sequence_number: 0xDEADBEEF,
            acknowledgment_number: 0xCAFEBABE,
            data_offset: 5,
            ns: false,
            cwr: false,
            ece: false,
            urg: false,
            ack: true,
            psh: false,
            rst: false,
            syn: true,
            fin: false,
            window: 0xFFFF,
            checksum: Checksum(0x1234),
            urgent_pointer: 0,
        };
        codec_roundtrip(etalon_struct, &etalon_bytes, ());
    }

    #[test]
    fn port() {
        let etalon_bytes = &[0x00, 0x50]; // 80
        codec_roundtrip(Port::HTTP, etalon_bytes, ());
    }
}