initial working version

1 files changed, 369 insertions, 29 deletions

diff --git a/src/dhcp.rs b/src/dhcp.rs
index 38cc8e4..51680d1 100644
--- a/src/dhcp.rs
+++ b/src/dhcp.rs
@@ -1,6 +1,7 @@
 use std::{
     io::{Cursor, Read as _, Result, Write},
     net::Ipv4Addr,
+    str::FromStr,
 };
 
 use crate::wire;
@@ -8,6 +9,11 @@ use crate::wire;
 const MAGIC_COOKIE: [u8; 4] = [0x63, 0x82, 0x53, 0x63];
 const FLAG_BROADCAST: u16 = 1 << 15;
 
+pub const VENDOR_CLASS_PXE_CLIENT: &'static [u8] = b"PXEClient";
+pub const VENDOR_CLASS_PXE_SERVER: &'static [u8] = b"PXEServer";
+
+pub const USER_CLASS_IPXE: &'static [u8] = b"iPXE";
+
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum BootOp {
     #[default]
@@ -29,6 +35,8 @@ impl From<BootOp> for u8 {
     }
 }
 
+pub type HardwareAddress = [u8; 16];
+
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum HardwareType {
     #[default]
@@ -78,10 +86,11 @@ pub enum DhcpOption {
     End,
     MessageType(DhcpMessageType),
     ServerIdentifier(Ipv4Addr),
-    VendorClassIdentifier(String),
+    VendorClassIdentifier(Vec<u8>),
     TftpServerName(String),
     TftpFileName(String),
-    UserClassInformation(String),
+    UserClassInformation(Vec<u8>),
+    ClientSystemArchitecture(SystemArchitecture),
     ClientMachineIdentifier(Vec<u8>),
     Unknown { code: u8, data: Vec<u8> },
 }
@@ -95,6 +104,7 @@ impl DhcpOption {
     pub const CODE_TFTP_SERVER_NAME: u8 = 66;
     pub const CODE_TFTP_FILE_NAME: u8 = 67;
     pub const CODE_USER_CLASS_INFORMATION: u8 = 77;
+    pub const CODE_CLIENT_SYSTEM_ARCHITECTURE: u8 = 93;
     pub const CODE_CLIENT_MACHINE_IDENTIFIER: u8 = 97;
 
     pub fn code(&self) -> u8 {
@@ -107,12 +117,294 @@ impl DhcpOption {
             DhcpOption::TftpServerName(_) => Self::CODE_TFTP_SERVER_NAME,
             DhcpOption::TftpFileName(_) => Self::CODE_TFTP_FILE_NAME,
             DhcpOption::UserClassInformation(_) => Self::CODE_USER_CLASS_INFORMATION,
+            DhcpOption::ClientSystemArchitecture(_) => Self::CODE_CLIENT_SYSTEM_ARCHITECTURE,
             DhcpOption::ClientMachineIdentifier(_) => Self::CODE_CLIENT_MACHINE_IDENTIFIER,
             DhcpOption::Unknown { code, .. } => *code,
         }
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum SystemArchitecture {
+    IntelX86pc,
+    NECPC98,
+    EfiItanium,
+    DecAlpha,
+    ArcX86,
+    IntelLeanClient,
+    EfiIA32,
+    EfiBC,
+    EfiXscale,
+    EfiX86_64,
+    EfiARM32,
+    EfiARM64,
+    EfiARM32Http,
+    EfiARM64Http,
+    ARM32Uboot,
+    ARM64Uboot,
+    Unknown(u16),
+}
+
+impl SystemArchitecture {
+    pub const CODE_INTEL_X86_PC: u16 = 0;
+    pub const CODE_NEC_PC98: u16 = 1;
+    pub const CODE_EFI_ITANIUM: u16 = 2;
+    pub const CODE_DEC_ALPHA: u16 = 3;
+    pub const CODE_ARC_X86: u16 = 4;
+    pub const CODE_INTEL_LEAN_CLIENT: u16 = 5;
+    pub const CODE_EFI_IA32: u16 = 6;
+    pub const CODE_EFI_BC: u16 = 7;
+    pub const CODE_EFI_XSCALE: u16 = 8;
+    pub const CODE_EFI_X86_64: u16 = 9;
+    pub const CODE_EFI_ARM32: u16 = 10;
+    pub const CODE_EFI_ARM64: u16 = 11;
+    pub const CODE_EFI_ARM32_HTTP: u16 = 18;
+    pub const CODE_EFI_ARM64_HTTP: u16 = 19;
+    pub const CODE_ARM32_UBOOT: u16 = 21;
+    pub const CODE_ARM64_UBOOT: u16 = 22;
+}
+
+impl From<u16> for SystemArchitecture {
+    fn from(value: u16) -> Self {
+        match value {
+            Self::CODE_INTEL_X86_PC => SystemArchitecture::IntelX86pc,
+            Self::CODE_NEC_PC98 => SystemArchitecture::NECPC98,
+            Self::CODE_EFI_ITANIUM => SystemArchitecture::EfiItanium,
+            Self::CODE_DEC_ALPHA => SystemArchitecture::DecAlpha,
+            Self::CODE_ARC_X86 => SystemArchitecture::ArcX86,
+            Self::CODE_INTEL_LEAN_CLIENT => SystemArchitecture::IntelLeanClient,
+            Self::CODE_EFI_IA32 => SystemArchitecture::EfiIA32,
+            Self::CODE_EFI_BC => SystemArchitecture::EfiBC,
+            Self::CODE_EFI_XSCALE => SystemArchitecture::EfiXscale,
+            Self::CODE_EFI_X86_64 => SystemArchitecture::EfiX86_64,
+            Self::CODE_EFI_ARM32 => SystemArchitecture::EfiARM32,
+            Self::CODE_EFI_ARM64 => SystemArchitecture::EfiARM64,
+            Self::CODE_EFI_ARM32_HTTP => SystemArchitecture::EfiARM32Http,
+            Self::CODE_EFI_ARM64_HTTP => SystemArchitecture::EfiARM64Http,
+            Self::CODE_ARM32_UBOOT => SystemArchitecture::ARM32Uboot,
+            Self::CODE_ARM64_UBOOT => SystemArchitecture::ARM64Uboot,
+            _ => SystemArchitecture::Unknown(value),
+        }
+    }
+}
+
+impl From<SystemArchitecture> for u16 {
+    fn from(value: SystemArchitecture) -> Self {
+        match value {
+            SystemArchitecture::IntelX86pc => SystemArchitecture::CODE_INTEL_X86_PC,
+            SystemArchitecture::NECPC98 => SystemArchitecture::CODE_NEC_PC98,
+            SystemArchitecture::EfiItanium => SystemArchitecture::CODE_EFI_ITANIUM,
+            SystemArchitecture::DecAlpha => SystemArchitecture::CODE_DEC_ALPHA,
+            SystemArchitecture::ArcX86 => SystemArchitecture::CODE_ARC_X86,
+            SystemArchitecture::IntelLeanClient => SystemArchitecture::CODE_INTEL_LEAN_CLIENT,
+            SystemArchitecture::EfiIA32 => SystemArchitecture::CODE_EFI_IA32,
+            SystemArchitecture::EfiBC => SystemArchitecture::CODE_EFI_BC,
+            SystemArchitecture::EfiXscale => SystemArchitecture::CODE_EFI_XSCALE,
+            SystemArchitecture::EfiX86_64 => SystemArchitecture::CODE_EFI_X86_64,
+            SystemArchitecture::EfiARM32 => SystemArchitecture::CODE_EFI_ARM32,
+            SystemArchitecture::EfiARM64 => SystemArchitecture::CODE_EFI_ARM64,
+            SystemArchitecture::EfiARM32Http => SystemArchitecture::CODE_EFI_ARM32_HTTP,
+            SystemArchitecture::EfiARM64Http => SystemArchitecture::CODE_EFI_ARM64_HTTP,
+            SystemArchitecture::ARM32Uboot => SystemArchitecture::CODE_ARM32_UBOOT,
+            SystemArchitecture::ARM64Uboot => SystemArchitecture::CODE_ARM64_UBOOT,
+            SystemArchitecture::Unknown(code) => code,
+        }
+    }
+}
+
+impl FromStr for SystemArchitecture {
+    type Err = <u16 as FromStr>::Err;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        s.parse::<u16>().map(From::from)
+    }
+}
+
+#[derive(Debug)]
+pub struct InvalidPxeClassIdentifierKind(String);
+
+impl InvalidPxeClassIdentifierKind {
+    fn new(kind: impl Into<String>) -> Self {
+        Self(kind.into())
+    }
+}
+
+impl std::fmt::Display for InvalidPxeClassIdentifierKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "invalid pxe class identifier kind '{}', expected 'PXEClient' or 'PXEServer'",
+            self.0
+        )
+    }
+}
+
+impl std::error::Error for InvalidPxeClassIdentifierKind {}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum PxeClassIdentifierKind {
+    Client,
+    Server,
+}
+
+impl PxeClassIdentifierKind {
+    pub const KIND_CLIENT: &'static str = "PXEClient";
+    pub const KIND_SERVER: &'static str = "PXEServer";
+}
+
+impl FromStr for PxeClassIdentifierKind {
+    type Err = InvalidPxeClassIdentifierKind;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s {
+            Self::KIND_CLIENT => Ok(Self::Client),
+            Self::KIND_SERVER => Ok(Self::Server),
+            _ => Err(InvalidPxeClassIdentifierKind::new(s)),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub struct InvalidPxeClassIdentifier(String, String);
+
+impl InvalidPxeClassIdentifier {
+    fn new(class: impl Into<String>, reason: impl Into<String>) -> Self {
+        Self(class.into(), reason.into())
+    }
+}
+
+impl std::fmt::Display for InvalidPxeClassIdentifier {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "invalid pxe class identifier '{}': {}", self.0, self.1)
+    }
+}
+
+impl std::error::Error for InvalidPxeClassIdentifier {}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum PxeClassIdentifier {
+    Client(PxeClassIdentifierClient),
+    Server(PxeClassIdentifierServer),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct PxeClassIdentifierServer;
+
+impl std::fmt::Display for PxeClassIdentifierServer {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.write_str("PXEServer")
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct PxeClassIdentifierClient {
+    pub architecture: SystemArchitecture,
+    pub undi_major: u16,
+    pub undi_minor: u16,
+}
+
+impl std::fmt::Display for PxeClassIdentifierClient {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "PXEClient:Arch:{:05}:UNDI:{:03}{:03}",
+            u16::from(self.architecture),
+            self.undi_major,
+            self.undi_minor
+        )
+    }
+}
+
+impl std::fmt::Display for PxeClassIdentifier {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PxeClassIdentifier::Client(client) => client.fmt(f),
+            PxeClassIdentifier::Server(server) => server.fmt(f),
+        }
+    }
+}
+
+impl TryFrom<&[u8]> for PxeClassIdentifier {
+    type Error = InvalidPxeClassIdentifier;
+
+    fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> {
+        let str = std::str::from_utf8(value).map_err(|err| {
+            InvalidPxeClassIdentifier::new(
+                format!("{value:?}"),
+                format!("invalid utf-8 string: {err}"),
+            )
+        })?;
+        str.parse()
+    }
+}
+
+impl FromStr for PxeClassIdentifier {
+    type Err = InvalidPxeClassIdentifier;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        let mut parts = s.split(":");
+        let make_err = |reason: String| InvalidPxeClassIdentifier::new(s, reason);
+
+        let kind = match parts.next() {
+            Some(kind) => kind
+                .parse::<PxeClassIdentifierKind>()
+                .map_err(|err| make_err(err.to_string()))?,
+            None => return Err(make_err("missing class kind".to_string())),
+        };
+
+        if kind == PxeClassIdentifierKind::Server {
+            if parts.next().is_some() {
+                return Err(make_err("invalid class".to_string()));
+            }
+            return Ok(Self::Server(PxeClassIdentifierServer));
+        }
+
+        if !parts.next().map(|s| s == "Arch").unwrap_or(false) {
+            return Err(make_err("invalid class".to_string()));
+        }
+
+        let architecture = match parts.next() {
+            Some(arch) => arch
+                .parse::<SystemArchitecture>()
+                .map_err(|err| make_err(err.to_string()))?,
+            None => return Err(make_err("missing architecture".to_string())),
+        };
+
+        if !parts.next().map(|s| s == "UNDI").unwrap_or(false) {
+            return Err(make_err("invalid class".to_string()));
+        }
+
+        let undi_str = match parts.next() {
+            Some(undi_str) => undi_str,
+            None => return Err(make_err("missing undi version".to_string())),
+        };
+
+        if undi_str.len() != 6 {
+            return Err(make_err("invalid undi version length".to_string()));
+        }
+
+        let (undi_major_str, undi_minor_str) = undi_str.split_at_checked(3).unwrap();
+
+        let undi_major = undi_major_str
+            .parse::<u16>()
+            .map_err(|err| make_err(err.to_string()))?;
+
+        let undi_minor = undi_minor_str
+            .parse::<u16>()
+            .map_err(|err| make_err(err.to_string()))?;
+
+        if parts.next().is_some() {
+            return Err(make_err("invalid class".to_string()));
+        }
+
+        Ok(Self::Client(PxeClassIdentifierClient {
+            architecture,
+            undi_major,
+            undi_minor,
+        }))
+    }
+}
+
 #[derive(Debug)]
 pub struct DhcpPacket {
     pub op: BootOp,
@@ -125,7 +417,7 @@ pub struct DhcpPacket {
     pub yiaddr: Ipv4Addr,
     pub siaddr: Ipv4Addr,
     pub giaddr: Ipv4Addr,
-    pub chaddr: [u8; 16],
+    pub chaddr: HardwareAddress,
     // server host name
     pub sname: String,
     // boot file name
@@ -158,11 +450,21 @@ impl DhcpPacket {
     pub fn new_boot(
         xid: u32,
         chaddr: [u8; 16],
-        client_uuid: Vec<u8>,
+        client_uuid: Option<Vec<u8>>,
         local_ip: Ipv4Addr,
         local_hostname: String,
         filename: String,
     ) -> Self {
+        let mut options = vec![
+            DhcpOption::MessageType(DhcpMessageType::Offer),
+            DhcpOption::ServerIdentifier(local_ip),
+            DhcpOption::VendorClassIdentifier(b"PXEClient".to_vec()),
+            DhcpOption::TftpServerName(local_hostname),
+            DhcpOption::TftpFileName(filename),
+        ];
+        if let Some(uuid) = client_uuid {
+            options.push(DhcpOption::ClientMachineIdentifier(uuid));
+        }
         Self {
             op: BootOp::Reply,
             htype: HardwareType::Ethernet,
@@ -177,25 +479,28 @@ impl DhcpPacket {
             chaddr,
             sname: Default::default(),
             file: Default::default(),
-            options: vec![
-                DhcpOption::MessageType(DhcpMessageType::Offer),
-                DhcpOption::ServerIdentifier(local_ip),
-                DhcpOption::VendorClassIdentifier("PXEClient".to_string()),
-                DhcpOption::ClientMachineIdentifier(client_uuid),
-                DhcpOption::TftpServerName(local_hostname),
-                DhcpOption::TftpFileName(filename),
-            ],
+            options,
         }
     }
 
     pub fn new_boot_ack(
         xid: u32,
         chaddr: [u8; 16],
-        client_uuid: Vec<u8>,
+        client_uuid: Option<Vec<u8>>,
         local_ip: Ipv4Addr,
         hostname: String,
         filename: String,
     ) -> Self {
+        let mut options = vec![
+            DhcpOption::MessageType(DhcpMessageType::Ack),
+            DhcpOption::ServerIdentifier(local_ip),
+            DhcpOption::VendorClassIdentifier(b"PXEClient".to_vec()),
+            DhcpOption::TftpServerName(hostname),
+            DhcpOption::TftpFileName(filename),
+        ];
+        if let Some(uuid) = client_uuid {
+            options.push(DhcpOption::ClientMachineIdentifier(uuid));
+        }
         Self {
             op: BootOp::Reply,
             htype: HardwareType::Ethernet,
@@ -210,14 +515,7 @@ impl DhcpPacket {
             chaddr,
             sname: Default::default(),
             file: Default::default(),
-            options: vec![
-                DhcpOption::MessageType(DhcpMessageType::Ack),
-                DhcpOption::ServerIdentifier(local_ip),
-                DhcpOption::VendorClassIdentifier("PXEClient".to_string()),
-                DhcpOption::ClientMachineIdentifier(client_uuid),
-                DhcpOption::TftpServerName(hostname),
-                DhcpOption::TftpFileName(filename),
-            ],
+            options,
         }
     }
 
@@ -296,10 +594,23 @@ fn read_option(cursor: &mut Cursor<&[u8]>) -> Result<DhcpOption> {
         DhcpOption::CODE_PAD => DhcpOption::Pad,
         DhcpOption::CODE_END => DhcpOption::End,
         DhcpOption::CODE_VENDOR_CLASS_IDENTIFIER => {
-            DhcpOption::VendorClassIdentifier(read_len8_prefixed_string(cursor)?)
+            DhcpOption::VendorClassIdentifier(read_len8_prefixed_vec(cursor)?)
         }
         DhcpOption::CODE_USER_CLASS_INFORMATION => {
-            DhcpOption::UserClassInformation(read_len8_prefixed_string(cursor)?)
+            DhcpOption::UserClassInformation(read_len8_prefixed_vec(cursor)?)
+        }
+        DhcpOption::CODE_CLIENT_SYSTEM_ARCHITECTURE => {
+            let len = read_u8(cursor)?;
+            assert_eq!(len, 2);
+
+            let mut buf = [0u8; 2];
+            cursor.read_exact(&mut buf)?;
+
+            let arch = SystemArchitecture::from(u16::from_be_bytes(buf));
+            DhcpOption::ClientSystemArchitecture(arch)
+        }
+        DhcpOption::CODE_CLIENT_MACHINE_IDENTIFIER => {
+            DhcpOption::ClientMachineIdentifier(read_len8_prefixed_vec(cursor)?)
         }
         _ => {
             let len = read_u8(cursor)?;
@@ -353,6 +664,20 @@ pub fn parse_packet(buf: &[u8]) -> Result<DhcpPacket> {
 }
 
 pub fn write_packet<W: Write>(mut writer: W, packet: &DhcpPacket) -> Result<()> {
+    if packet.sname.len() >= 64 {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "sname cannot be longer than 64 bytes",
+        ));
+    }
+
+    if packet.file.len() >= 128 {
+        return Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidInput,
+            "filename cannot be longer than 128 bytes",
+        ));
+    }
+
     wire::write_u8(&mut writer, u8::from(packet.op))?;
     wire::write_u8(&mut writer, u8::from(packet.htype))?;
     wire::write_u8(&mut writer, packet.htype.hardware_len())?;
@@ -365,10 +690,21 @@ pub fn write_packet<W: Write>(mut writer: W, packet: &DhcpPacket) -> Result<()>
     wire::write_ipv4(&mut writer, packet.siaddr)?;
     wire::write_ipv4(&mut writer, packet.giaddr)?;
     wire::write(&mut writer, &packet.chaddr)?;
-    //wire::write_null_terminated_string(&mut writer, &packet.sname)?;
-    //wire::write_null_terminated_string(&mut writer, &packet.file)?;
-    wire::write(&mut writer, &vec![0u8; 64])?;
-    wire::write(&mut writer, &vec![0u8; 128])?;
+
+    let sname_bytes = packet.sname.as_bytes();
+    wire::write(&mut writer, sname_bytes)?;
+    for _ in 0..(64 - sname_bytes.len()) {
+        wire::write_u8(&mut writer, 0)?;
+    }
+
+    let file_bytes = packet.file.as_bytes();
+    wire::write(&mut writer, file_bytes)?;
+    for _ in 0..(128 - file_bytes.len()) {
+        wire::write_u8(&mut writer, 0)?;
+    }
+
+    // wire::write(&mut writer, &vec![0u8; 64])?;
+    // wire::write(&mut writer, &vec![0u8; 128])?;
     wire::write(&mut writer, &MAGIC_COOKIE)?;
     for option in &packet.options {
         write_option(&mut writer, option)?;
@@ -390,12 +726,16 @@ pub fn write_option<W: Write>(mut writer: W, option: &DhcpOption) -> Result<()>
             wire::write_ipv4(&mut writer, *ip)?;
         }
         DhcpOption::VendorClassIdentifier(vendor_class) => {
-            write_option_len_prefixed_string(&mut writer, &vendor_class)?
+            write_option_len_prefixed_buf(&mut writer, &vendor_class)?
         }
         DhcpOption::TftpServerName(name) => write_option_len_prefixed_string(&mut writer, &name)?,
         DhcpOption::TftpFileName(name) => write_option_len_prefixed_string(&mut writer, &name)?,
         DhcpOption::UserClassInformation(user_class) => {
-            write_option_len_prefixed_string(&mut writer, &user_class)?
+            write_option_len_prefixed_buf(&mut writer, &user_class)?
+        }
+        DhcpOption::ClientSystemArchitecture(arch) => {
+            wire::write_u8(&mut writer, 2)?;
+            wire::write_u16(&mut writer, u16::from(*arch))?;
         }
         DhcpOption::ClientMachineIdentifier(identifier) => {
             write_option_len_prefixed_buf(&mut writer, &identifier)?