yggdrasil/lib/abi/src/net/protocols.rs

// TODO don't really want to document this now
#![allow(missing_docs)]

use core::{fmt, mem::size_of};

use crate::bitflags;

use super::{
    types::{net_value::WrappedValue, NetValue, NetValueImpl},
    MacAddress,
};

/////////// Layer 2 protocols ////////////

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct EthernetFrame {
    pub destination_mac: MacAddress,
    pub source_mac: MacAddress,
    pub ethertype: NetValue<EtherType>,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(transparent)]
pub struct EtherType(pub u16);

impl EtherType {
    pub const ARP: Self = Self(0x0806);
    pub const IPV4: Self = Self(0x0800);
}

impl WrappedValue for EtherType {
    type Inner = u16;

    fn into_inner(self) -> u16 {
        self.0
    }
    fn from_inner(value: u16) -> Self {
        Self(value)
    }
}

/////////// Layer 3 protocols ////////////

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct ArpFrame {
    pub hardware_type: NetValue<u16>,
    pub protocol: NetValue<EtherType>,
    pub hardware_size: u8,
    pub protocol_size: u8,
    pub opcode: NetValue<u16>,
    pub sender_mac: MacAddress,
    pub sender_ip: NetValue<u32>,
    pub target_mac: MacAddress,
    // TODO handle IPv6
    pub target_ip: NetValue<u32>,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct Ipv4Frame {
    pub version_length: u8,
    pub dscp_flags: u8,
    pub total_length: NetValue<u16>,
    pub id: NetValue<u16>,
    pub flags_frag: NetValue<u16>,
    pub ttl: u8,
    pub protocol: IpProtocol,
    pub header_checksum: NetValue<u16>,
    pub source_address: NetValue<u32>,
    pub destination_address: NetValue<u32>,
}

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(transparent)]
pub struct IpProtocol(pub u8);

impl IpProtocol {
    pub const ICMP: Self = Self(1);
    pub const TCP: Self = Self(6);
    pub const UDP: Self = Self(17);
}

/////////// Layer 4 protocols ////////////

#[derive(Clone, Copy)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct UdpFrame {
    pub source_port: NetValue<u16>,
    pub destination_port: NetValue<u16>,
    pub length: NetValue<u16>,
    pub checksum: NetValue<u16>,
}

bitflags! {
    #[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
    pub struct TcpFlags: u8 {
        const FIN: bit 0;
        const SYN: bit 1;
        const RST: bit 2;
        const PSH: bit 3;
        const ACK: bit 4;
        const URG: bit 5;
        const ECE: bit 6;
        const CWR: bit 7;
    }
}

#[derive(Clone, Copy)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct TcpFrame {
    pub source_port: NetValue<u16>,
    pub destination_port: NetValue<u16>,
    pub sequence_number: NetValue<u32>,
    pub acknowledge_number: NetValue<u32>,
    pub data_offset: u8,
    pub flags: TcpFlags,
    pub window_size: NetValue<u16>,
    pub checksum: NetValue<u16>,
    pub urgent_pointer: NetValue<u16>,
}

#[derive(Clone, Copy)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct TcpV4PseudoHeader {
    pub source_address: NetValue<u32>,
    pub destination_address: NetValue<u32>,
    pub _zero: u8,
    pub protocol: IpProtocol,
    pub tcp_length: NetValue<u16>,
}

#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Pod, bytemuck::Zeroable))]
#[repr(C, packed)]
pub struct IcmpV4Frame {
    pub ty: u8,
    pub code: u8,
    pub checksum: NetValue<u16>,
    pub rest: NetValue<u32>,
}

/////////// Helper utilities /////////////

pub struct InetChecksum {
    state: u32,
}

///////////// Implementations ////////////

/////////// Layer 3 protocols ////////////

impl fmt::Display for IpProtocol {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match *self {
            Self::ICMP => f.write_str("ICMP"),
            Self::UDP => f.write_str("UDP"),
            Self::TCP => f.write_str("TCP"),
            _ => f.write_str("<unknown-l4-proto>"),
        }
    }
}

impl Ipv4Frame {
    pub fn header_length(&self) -> usize {
        core::cmp::max(
            (self.version_length & 0xF) << 2,
            size_of::<Ipv4Frame>() as u8,
        ) as usize
    }

    pub fn total_length(&self) -> usize {
        u16::from_network_order(self.total_length) as usize
    }
}

/////////// Layer 4 protocols ////////////

impl UdpFrame {
    pub fn data_length(&self) -> usize {
        (u16::from_network_order(self.length) as usize).saturating_sub(size_of::<UdpFrame>())
    }
}

impl TcpFrame {
    pub fn data_offset(&self) -> usize {
        ((self.data_offset >> 4) as usize * size_of::<u32>()).max(size_of::<TcpFrame>())
    }
}

/////////// Helper utilities /////////////

impl InetChecksum {
    pub fn new() -> Self {
        Self { state: 0 }
    }

    #[cfg(feature = "bytemuck")]
    pub fn add_value<T: bytemuck::Pod>(&mut self, value: &T, reorder: bool) {
        self.add_bytes(bytemuck::bytes_of(value), reorder)
    }

    pub fn add_bytes(&mut self, bytes: &[u8], reorder: bool) {
        let len = bytes.len();
        for i in 0..len / 2 {
            let word = if reorder {
                ((bytes[i * 2] as u16) << 8) | (bytes[i * 2 + 1] as u16)
            } else {
                (bytes[i * 2] as u16) | ((bytes[i * 2 + 1] as u16) << 8)
            };
            self.state = self.state.wrapping_add(word as u32);
        }
        if len % 2 != 0 {
            let word = if reorder {
                (bytes[len - 1] as u16) << 8
            } else {
                todo!()
            };
            self.state = self.state.wrapping_add(word as u32);
        }
    }

    pub fn finish(self) -> u16 {
        let sum = (self.state >> 16) + (self.state & 0xFFFF);

        (!(sum & 0xFFFF)) as u16
    }
}