yggdrasil/userspace/netutils/src/dhcp_client.rs

#![feature(yggdrasil_os)]

use std::os::{
    fd::AsRawFd,
    yggdrasil::io::{poll::PollChannel, net::raw_socket::RawSocket, timer::TimerFd},
};
use std::{io, mem::size_of, process::ExitCode, time::Duration};

use bytemuck::{Pod, Zeroable};
use netutils::{netconfig::NetConfig, proto::parse_udp_protocol, Error};
use yggdrasil_abi::net::protocols::{
    EtherType, EthernetFrame, InetChecksum, IpProtocol, Ipv4Frame, UdpFrame,
};
use yggdrasil_abi::net::types::NetValueImpl;
use yggdrasil_abi::net::{IpAddr, Ipv4Addr, MacAddress, SubnetAddr, SubnetV4Addr};

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C, packed)]
struct DhcpMessageHeader {
    op: u8,
    htype: u8,
    hlen: u8,
    hops: u8,
    xid: u32,
    secs: u16,
    flags: u16,
    ciaddr: u32,
    yiaddr: u32,
    siaddr: u32,
    giaddr: u32,
    chaddr: [u8; 16],
    sname: [u8; 64],
    file: [u8; 128],
    cookie: u32,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum DhcpMessageType {
    Offer,
    Discover,
    Request,
    Acknowledge,
}

#[derive(Debug)]
enum DhcpOption {
    MessageType(DhcpMessageType),
    SubnetMask(u32),
    Router(Ipv4Addr),
    ServerId(Ipv4Addr),
    // ParameterRequestList(Vec<u8>),
    RequestedAddress(Ipv4Addr),
}

#[derive(Debug)]
enum ParsedOption {
    Option(usize, DhcpOption),
    Unknown(usize),
    Pad,
    End,
}

#[derive(Debug)]
struct DhcpMessage {
    header: DhcpMessageHeader,
    options: Vec<DhcpOption>,
}

#[derive(Debug)]
struct DhcpOffer {
    router_address: Ipv4Addr,
    server_address: Ipv4Addr,
    your_address: Ipv4Addr,
    subnet_mask: u32,
}

fn pad_to_align(buffer: &mut Vec<u8>) {
    let offset = buffer.len() % 4;
    for _ in offset..4 {
        buffer.push(0);
    }
}

impl From<DhcpMessageType> for u8 {
    fn from(value: DhcpMessageType) -> u8 {
        match value {
            DhcpMessageType::Discover => 0x01,
            DhcpMessageType::Offer => 0x02,
            DhcpMessageType::Request => 0x03,
            DhcpMessageType::Acknowledge => 0x05,
        }
    }
}

impl TryFrom<u8> for DhcpMessageType {
    type Error = ();

    fn try_from(value: u8) -> Result<DhcpMessageType, Self::Error> {
        match value {
            0x01 => Ok(DhcpMessageType::Discover),
            0x02 => Ok(DhcpMessageType::Offer),
            0x03 => Ok(DhcpMessageType::Request),
            0x05 => Ok(DhcpMessageType::Acknowledge),
            _ => Err(()),
        }
    }
}

impl DhcpOption {
    fn write(&self, buffer: &mut Vec<u8>) {
        match self {
            Self::MessageType(ty) => {
                buffer.extend_from_slice(&[53, 0x01, (*ty).into()]);
            }
            // Self::ParameterRequestList(list) => todo!(),
            Self::RequestedAddress(addr) => {
                buffer.extend_from_slice(&[50, 0x04]);
                buffer.extend_from_slice(&u32::from(*addr).to_be_bytes());
            }
            Self::ServerId(addr) => {
                buffer.extend_from_slice(&[54, 0x04]);
                buffer.extend_from_slice(&u32::from(*addr).to_be_bytes());
            }
            _ => todo!(),
        }
    }

    fn parse(data: &[u8]) -> Option<ParsedOption> {
        if data.is_empty() {
            return None;
        }
        if data[0] == 0 {
            return Some(ParsedOption::Pad);
        }
        if data[0] == 0xFF {
            return Some(ParsedOption::End);
        }

        if data.len() < 2 {
            return None;
        }

        let opcode = data[0];
        let length = data[1] as usize;

        if data.len() < 2 + length {
            return None;
        }

        Some(ParsedOption::Option(
            2 + length,
            match opcode {
                // DHCP message type
                1 if length == 4 => {
                    let value = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
                    DhcpOption::SubnetMask(value)
                }
                3 if length == 4 => {
                    let value = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
                    DhcpOption::Router(value.into())
                }
                53 if length == 1 => {
                    let ty = DhcpMessageType::try_from(data[2]).ok()?;
                    DhcpOption::MessageType(ty)
                }
                54 if length == 4 => {
                    let value = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
                    DhcpOption::ServerId(value.into())
                }
                _ => return Some(ParsedOption::Unknown(2 + length)),
            },
        ))
    }
}

impl DhcpMessage {
    fn to_vec(&self) -> Vec<u8> {
        let mut result = vec![];

        result.extend_from_slice(bytemuck::bytes_of(&self.header));

        for option in self.options.iter() {
            option.write(&mut result);
            pad_to_align(&mut result);
        }

        // End option
        result.push(0xFF);
        pad_to_align(&mut result);

        result
    }

    fn parse(data: &[u8]) -> Option<Self> {
        if data.len() < size_of::<DhcpMessageHeader>() {
            return None;
        }
        let header: &DhcpMessageHeader =
            bytemuck::from_bytes(&data[..size_of::<DhcpMessageHeader>()]);

        let mut options = vec![];

        let mut offset = size_of::<DhcpMessageHeader>();
        loop {
            let data = &data[offset..];
            let option = DhcpOption::parse(data)?;

            match option {
                ParsedOption::Option(len, value) => {
                    options.push(value);
                    offset += len;
                }
                ParsedOption::Unknown(len) => {
                    offset += len;
                }
                ParsedOption::Pad => {
                    offset += 1;
                }
                ParsedOption::End => break,
            }
        }

        Some(Self {
            header: *header,
            options,
        })
    }

    fn message_type(&self) -> Option<DhcpMessageType> {
        self.options.iter().find_map(|option| match option {
            DhcpOption::MessageType(ty) => Some(*ty),
            _ => None,
        })
    }

    fn as_dhcp_acknowledge(&self, requested_address: Ipv4Addr) -> Option<()> {
        let ty = self.message_type()?;

        if ty != DhcpMessageType::Acknowledge {
            return None;
        }
        let yiaddr = u32::from_be(self.header.yiaddr);
        if yiaddr == 0 {
            return None;
        }
        if Ipv4Addr::from(yiaddr) != requested_address {
            return None;
        }

        Some(())
    }

    fn as_dhcp_offer(&self) -> Option<DhcpOffer> {
        let ty = self.message_type()?;
        if ty != DhcpMessageType::Offer {
            return None;
        }
        let yiaddr = u32::from_be(self.header.yiaddr);
        if yiaddr == 0 {
            return None;
        }
        let mut router = None;
        let mut server = None;
        let mut subnet_mask = None;

        for option in self.options.iter() {
            match option {
                DhcpOption::Router(address) if router.is_none() => {
                    router.replace(*address);
                }
                DhcpOption::ServerId(address) if server.is_none() => {
                    server.replace(*address);
                }
                DhcpOption::SubnetMask(mask) if subnet_mask.is_none() => {
                    subnet_mask.replace(*mask);
                }
                _ => (),
            }
        }

        let router_address = router?;
        let server_address = server?;
        let subnet_mask = subnet_mask?;

        Some(DhcpOffer {
            server_address,
            router_address,
            your_address: Ipv4Addr::from(yiaddr),
            subnet_mask,
        })
    }
}

fn send_split_packet(
    socket: &RawSocket,
    l2_frame: &[u8],
    l3_frame: &[u8],
    l4_frame: &[u8],
    l4_data: &[u8],
) -> Result<(), io::Error> {
    let mut buffer =
        Vec::with_capacity(l2_frame.len() + l3_frame.len() + l4_frame.len() + l4_data.len());
    buffer.extend_from_slice(l2_frame);
    buffer.extend_from_slice(l3_frame);
    buffer.extend_from_slice(l4_frame);
    buffer.extend_from_slice(l4_data);

    socket.send(&buffer)?;

    Ok(())
}

fn send_udp_broadcast(
    socket: &RawSocket,
    source_mac: MacAddress,
    source_port: u16,
    destination_port: u16,
    data: &[u8],
) -> Result<(), io::Error> {
    let ip_size = size_of::<Ipv4Frame>() + size_of::<UdpFrame>() + data.len();
    let udp_size = size_of::<UdpFrame>() + data.len();

    let l2_frame = EthernetFrame {
        source_mac,
        destination_mac: MacAddress::BROADCAST,
        ethertype: EtherType::IPV4.to_network_order(),
    };
    let mut l3_frame = Ipv4Frame {
        source_address: u32::from(Ipv4Addr::UNSPECIFIED).to_network_order(),
        destination_address: u32::from(Ipv4Addr::BROADCAST).to_network_order(),
        protocol: IpProtocol::UDP,
        version_length: 0x45,
        total_length: u16::try_from(ip_size).unwrap().to_network_order(),
        dscp_flags: 0,
        flags_frag: 0x4000u16.to_network_order(),
        id: 0u16.to_network_order(),
        header_checksum: 0u16.to_network_order(),
        ttl: 64,
    };
    let l4_frame = UdpFrame {
        source_port: source_port.to_network_order(),
        destination_port: destination_port.to_network_order(),
        length: u16::try_from(udp_size).unwrap().to_network_order(),
        checksum: 0u16.to_network_order(),
    };

    let mut ip_checksum = InetChecksum::new();
    ip_checksum.add_bytes(bytemuck::bytes_of(&l3_frame), true);
    let ip_checksum = ip_checksum.finish();

    l3_frame.header_checksum = ip_checksum.to_network_order();

    send_split_packet(
        socket,
        bytemuck::bytes_of(&l2_frame),
        bytemuck::bytes_of(&l3_frame),
        bytemuck::bytes_of(&l4_frame),
        data,
    )
}

fn wait_for_dhcp_offer(
    poll: &mut PollChannel,
    timer: &mut TimerFd,
    socket: &RawSocket,
) -> Result<DhcpOffer, Error> {
    let mut packet = [0; 4096];

    let socket_fd = socket.as_raw_fd();
    let timer_fd = timer.as_raw_fd();

    timer.start(Duration::from_secs(1))?;

    loop {
        let (fd, result) = poll.wait(None, true)?.unwrap();
        result?;

        match fd {
            fd if fd == socket_fd => {
                let len = socket.recv(&mut packet)?;

                let Some((source, destination, data)) = parse_udp_protocol(&packet[..len]) else {
                    continue;
                };
                if source.port() != 67 || destination.port() != 68 {
                    continue;
                }
                let Some(msg) = DhcpMessage::parse(data) else {
                    continue;
                };
                let Some(offer) = msg.as_dhcp_offer() else {
                    continue;
                };

                return Ok(offer);
            }
            fd if fd == timer_fd => {
                return Err(Error::TimedOut);
            }
            _ => unreachable!(),
        }
    }
}

fn wait_for_dhcp_acknowledge(
    poll: &mut PollChannel,
    timer: &mut TimerFd,
    socket: &RawSocket,
    requested_address: Ipv4Addr,
) -> Result<(), Error> {
    let mut packet = [0; 4096];

    let socket_fd = socket.as_raw_fd();
    let timer_fd = timer.as_raw_fd();

    timer.start(Duration::from_secs(1))?;

    loop {
        let (fd, result) = poll.wait(None, true)?.unwrap();
        result?;

        match fd {
            fd if fd == socket_fd => {
                let len = socket.recv(&mut packet)?;

                let Some((source, destination, data)) = parse_udp_protocol(&packet[..len]) else {
                    continue;
                };
                if source.port() != 67 || destination.port() != 68 {
                    continue;
                }
                let Some(msg) = DhcpMessage::parse(data) else {
                    continue;
                };

                if msg.as_dhcp_acknowledge(requested_address).is_some() {
                    return Ok(());
                }
            }
            fd if fd == timer_fd => {
                return Err(Error::TimedOut);
            }
            _ => unreachable!(),
        }
    }
}

fn attempt_request(
    poll: &mut PollChannel,
    timer: &mut TimerFd,
    socket: &RawSocket,
    iface_mac_bytes: [u8; 6],
) -> Result<DhcpOffer, Error> {
    let iface_mac = MacAddress::from(iface_mac_bytes);

    let transaction_id = rand::random::<u32>();
    let mut message = DhcpMessage {
        header: DhcpMessageHeader {
            op: 1,
            htype: 1,
            hlen: 6,
            xid: transaction_id.to_be(),
            chaddr: [0; 16],
            cookie: 0x63825363u32.to_be(),
            ..DhcpMessageHeader::zeroed()
        },
        options: vec![DhcpOption::MessageType(DhcpMessageType::Discover)],
    };
    message.header.chaddr[..6].copy_from_slice(&iface_mac_bytes);
    let data = message.to_vec();

    send_udp_broadcast(socket, iface_mac, 68, 67, &data)?;

    let offer = wait_for_dhcp_offer(poll, timer, socket)?;

    println!(
        "Got DHCP offer: address={} router={} server={} mask={}",
        offer.your_address,
        offer.router_address,
        offer.server_address,
        Ipv4Addr::from(offer.subnet_mask)
    );

    // Send DHCP request
    let mut message = DhcpMessage {
        header: DhcpMessageHeader {
            op: 1,
            htype: 1,
            hlen: 6,
            xid: transaction_id.to_be(),
            chaddr: [0; 16],
            siaddr: u32::from(offer.server_address).to_be(),
            cookie: 0x63825363u32.to_be(),
            ..DhcpMessageHeader::zeroed()
        },
        options: vec![
            DhcpOption::MessageType(DhcpMessageType::Request),
            DhcpOption::RequestedAddress(offer.your_address),
            DhcpOption::ServerId(offer.server_address),
        ],
    };
    message.header.chaddr[..6].copy_from_slice(&iface_mac_bytes);
    let data = message.to_vec();

    send_udp_broadcast(socket, iface_mac, 68, 67, &data)?;

    wait_for_dhcp_acknowledge(poll, timer, socket, offer.your_address)?;

    println!("Got DHCP acknowledge for {}", offer.your_address);

    Ok(offer)
}

fn request_address(interface: &str) -> Result<DhcpOffer, Error> {
    let mut poll = PollChannel::new()?;
    let mut timer = TimerFd::new(false, false)?;
    let socket = RawSocket::bind(interface)?;
    let iface_mac_bytes = socket.hardware_address()?;

    poll.add(socket.as_raw_fd())?;
    poll.add(timer.as_raw_fd())?;

    for _ in 0..5 {
        match attempt_request(&mut poll, &mut timer, &socket, iface_mac_bytes) {
            Ok(offer) => {
                return Ok(offer);
            }
            Err(error) => {
                eprintln!("Error: {}, trying again", error);
            }
        }
    }

    Err(Error::TimedOut)
}

fn configure_address(interface: &str, offer: DhcpOffer) -> Result<(), Error> {
    let mut netconf = NetConfig::open()?;
    netconf.set_interface_address(interface, IpAddr::V4(offer.your_address))?;
    netconf.add_route(
        interface,
        SubnetAddr::V4(SubnetV4Addr::UNSPECIFIED),
        Some(IpAddr::V4(offer.router_address)),
    )?;
    Ok(())
}

fn main() -> ExitCode {
    let args: Vec<_> = std::env::args().collect();

    if args.len() != 2 {
        eprintln!("Usage: {} INTERFACE", args[0]);
        return ExitCode::FAILURE;
    }

    let offer = match request_address(&args[1]) {
        Ok(offer) => offer,
        Err(error) => {
            eprintln!("Could not obtain an address from DHCP: {}", error);
            return ExitCode::FAILURE;
        }
    };
    match configure_address(&args[1], offer) {
        Ok(_) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("Could not assign offered address to {}: {}", args[1], error);
            ExitCode::FAILURE
        }
    }
}