1 files changed, 219 insertions, 14 deletions

diff --git a/src/dhcp.rs b/src/dhcp.rs
index b53ee92..38cc8e4 100644
--- a/src/dhcp.rs
+++ b/src/dhcp.rs
@@ -1,11 +1,12 @@
 use std::{
-    io::{Result, Write},
+    io::{Cursor, Read as _, Result, Write},
     net::Ipv4Addr,
 };
 
 use crate::wire;
 
 const MAGIC_COOKIE: [u8; 4] = [0x63, 0x82, 0x53, 0x63];
+const FLAG_BROADCAST: u16 = 1 << 15;
 
 #[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum BootOp {
@@ -55,14 +56,17 @@ impl From<HardwareType> for u8 {
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum DhcpMessageType {
+    Offer,
     Ack,
 }
 
 impl DhcpMessageType {
+    pub const CODE_OFFER: u8 = 2;
     pub const CODE_ACK: u8 = 5;
 
     pub fn code(&self) -> u8 {
         match self {
+            DhcpMessageType::Offer => Self::CODE_OFFER,
             DhcpMessageType::Ack => Self::CODE_ACK,
         }
     }
@@ -78,7 +82,7 @@ pub enum DhcpOption {
     TftpServerName(String),
     TftpFileName(String),
     UserClassInformation(String),
-    ClientMachineIdentifier(String),
+    ClientMachineIdentifier(Vec<u8>),
     Unknown { code: u8, data: Vec<u8> },
 }
 
@@ -113,6 +117,7 @@ impl DhcpOption {
 pub struct DhcpPacket {
     pub op: BootOp,
     pub htype: HardwareType,
+    pub hops: u8,
     pub xid: u32,
     pub secs: u16,
     pub flags: u16,
@@ -122,9 +127,9 @@ pub struct DhcpPacket {
     pub giaddr: Ipv4Addr,
     pub chaddr: [u8; 16],
     // server host name
-    pub sname: Option<String>,
+    pub sname: String,
     // boot file name
-    pub file: Option<String>,
+    pub file: String,
     pub options: Vec<DhcpOption>,
 }
 
@@ -133,6 +138,7 @@ impl Default for DhcpPacket {
         Self {
             op: Default::default(),
             htype: Default::default(),
+            hops: Default::default(),
             xid: Default::default(),
             secs: Default::default(),
             flags: Default::default(),
@@ -148,11 +154,209 @@ impl Default for DhcpPacket {
     }
 }
 
+impl DhcpPacket {
+    pub fn new_boot(
+        xid: u32,
+        chaddr: [u8; 16],
+        client_uuid: Vec<u8>,
+        local_ip: Ipv4Addr,
+        local_hostname: String,
+        filename: String,
+    ) -> Self {
+        Self {
+            op: BootOp::Reply,
+            htype: HardwareType::Ethernet,
+            hops: Default::default(),
+            xid,
+            secs: Default::default(),
+            flags: FLAG_BROADCAST,
+            ciaddr: Ipv4Addr::UNSPECIFIED,
+            yiaddr: Ipv4Addr::UNSPECIFIED,
+            siaddr: local_ip,
+            giaddr: Ipv4Addr::UNSPECIFIED,
+            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),
+            ],
+        }
+    }
+
+    pub fn new_boot_ack(
+        xid: u32,
+        chaddr: [u8; 16],
+        client_uuid: Vec<u8>,
+        local_ip: Ipv4Addr,
+        hostname: String,
+        filename: String,
+    ) -> Self {
+        Self {
+            op: BootOp::Reply,
+            htype: HardwareType::Ethernet,
+            hops: 0,
+            xid,
+            secs: 0,
+            flags: 0,
+            ciaddr: Ipv4Addr::UNSPECIFIED,
+            yiaddr: Ipv4Addr::UNSPECIFIED,
+            siaddr: Ipv4Addr::UNSPECIFIED,
+            giaddr: Ipv4Addr::UNSPECIFIED,
+            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),
+            ],
+        }
+    }
+
+    pub fn write<W: Write>(&self, writer: W) -> Result<()> {
+        write_packet(writer, self)
+    }
+}
+
+fn read_u8(cursor: &mut Cursor<&[u8]>) -> Result<u8> {
+    let mut buf = [0u8; 1];
+    cursor.read_exact(&mut buf)?;
+    Ok(buf[0])
+}
+
+fn read_u16(cursor: &mut Cursor<&[u8]>) -> Result<u16> {
+    let mut buf = [0u8; 2];
+    cursor.read_exact(&mut buf)?;
+    Ok(u16::from_be_bytes(buf))
+}
+
+fn read_u32(cursor: &mut Cursor<&[u8]>) -> Result<u32> {
+    let mut buf = [0u8; 4];
+    cursor.read_exact(&mut buf)?;
+    Ok(u32::from_be_bytes(buf))
+}
+
+fn read_arr<const N: usize>(cursor: &mut Cursor<&[u8]>) -> Result<[u8; N]> {
+    let mut buf = [0u8; N];
+    cursor.read_exact(&mut buf)?;
+    Ok(buf)
+}
+
+fn read_len8_prefixed_vec(cursor: &mut Cursor<&[u8]>) -> Result<Vec<u8>> {
+    let len = read_u8(cursor)?;
+    let mut buf = vec![0u8; len as usize];
+    cursor.read_exact(&mut buf)?;
+    Ok(buf)
+}
+
+fn read_len8_prefixed_string(cursor: &mut Cursor<&[u8]>) -> Result<String> {
+    let buf = read_len8_prefixed_vec(cursor)?;
+    Ok(String::from_utf8(buf).unwrap())
+}
+
+fn read_ipv4(cursor: &mut Cursor<&[u8]>) -> Result<Ipv4Addr> {
+    Ok(Ipv4Addr::from_octets(read_arr(cursor)?))
+}
+
+fn read_op(cursor: &mut Cursor<&[u8]>) -> Result<BootOp> {
+    let v = read_u8(cursor)?;
+    match v {
+        BootOp::OP_REQUEST => Ok(BootOp::Request),
+        BootOp::OP_REPLY => Ok(BootOp::Reply),
+        _ => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "invalid boot op",
+        )),
+    }
+}
+
+fn read_htype(cursor: &mut Cursor<&[u8]>) -> Result<HardwareType> {
+    let ty = read_u8(cursor)?;
+    let len = read_u8(cursor)?;
+    match (ty, len) {
+        (HardwareType::TYPE_ETHER, HardwareType::LEN_ETHER) => Ok(HardwareType::Ethernet),
+        _ => Err(std::io::Error::new(
+            std::io::ErrorKind::InvalidData,
+            "invalid hardware type",
+        )),
+    }
+}
+
+fn read_option(cursor: &mut Cursor<&[u8]>) -> Result<DhcpOption> {
+    let code = read_u8(cursor)?;
+    Ok(match code {
+        DhcpOption::CODE_PAD => DhcpOption::Pad,
+        DhcpOption::CODE_END => DhcpOption::End,
+        DhcpOption::CODE_VENDOR_CLASS_IDENTIFIER => {
+            DhcpOption::VendorClassIdentifier(read_len8_prefixed_string(cursor)?)
+        }
+        DhcpOption::CODE_USER_CLASS_INFORMATION => {
+            DhcpOption::UserClassInformation(read_len8_prefixed_string(cursor)?)
+        }
+        _ => {
+            let len = read_u8(cursor)?;
+            let mut data = vec![0u8; usize::from(len)];
+            cursor.read_exact(&mut data)?;
+            DhcpOption::Unknown { code, data }
+        }
+    })
+}
+
+fn read_sname(cursor: &mut Cursor<&[u8]>) -> Result<String> {
+    let arr = read_arr::<64>(cursor)?;
+    let sname = std::str::from_utf8(&arr).unwrap();
+    Ok(sname.to_string())
+}
+
+fn read_filename(cursor: &mut Cursor<&[u8]>) -> Result<String> {
+    let arr = read_arr::<128>(cursor)?;
+    let filename = std::str::from_utf8(&arr).unwrap();
+    Ok(filename.to_string())
+}
+
+pub fn parse_packet(buf: &[u8]) -> Result<DhcpPacket> {
+    let mut cursor = Cursor::new(buf);
+    let mut packet = DhcpPacket {
+        op: read_op(&mut cursor)?,
+        htype: read_htype(&mut cursor)?,
+        hops: read_u8(&mut cursor)?,
+        xid: read_u32(&mut cursor)?,
+        secs: read_u16(&mut cursor)?,
+        flags: read_u16(&mut cursor)?,
+        ciaddr: read_ipv4(&mut cursor)?,
+        yiaddr: read_ipv4(&mut cursor)?,
+        siaddr: read_ipv4(&mut cursor)?,
+        giaddr: read_ipv4(&mut cursor)?,
+        chaddr: read_arr(&mut cursor)?,
+        sname: read_sname(&mut cursor)?,
+        file: read_filename(&mut cursor)?,
+        options: Default::default(),
+    };
+
+    let magic = read_arr::<4>(&mut cursor)?;
+    assert_eq!(magic, MAGIC_COOKIE);
+
+    while cursor.position() < buf.len() as u64 {
+        let option = read_option(&mut cursor)?;
+        packet.options.push(option);
+    }
+
+    Ok(packet)
+}
+
 pub fn write_packet<W: Write>(mut writer: W, packet: &DhcpPacket) -> Result<()> {
     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())?;
-    wire::write_u8(&mut writer, 0)?; // hops
+    wire::write_u8(&mut writer, packet.hops)?;
     wire::write_u32(&mut writer, packet.xid)?;
     wire::write_u16(&mut writer, packet.secs)?;
     wire::write_u16(&mut writer, packet.flags)?;
@@ -161,14 +365,10 @@ 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)?;
-    match &packet.sname {
-        Some(name) => wire::write_null_terminated_string(&mut writer, &name)?,
-        None => wire::write_null_terminated_string(&mut writer, "")?,
-    };
-    match &packet.file {
-        Some(name) => wire::write_null_terminated_string(&mut writer, &name)?,
-        None => wire::write_null_terminated_string(&mut writer, "")?,
-    };
+    //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])?;
     wire::write(&mut writer, &MAGIC_COOKIE)?;
     for option in &packet.options {
         write_option(&mut writer, option)?;
@@ -198,7 +398,7 @@ pub fn write_option<W: Write>(mut writer: W, option: &DhcpOption) -> Result<()>
             write_option_len_prefixed_string(&mut writer, &user_class)?
         }
         DhcpOption::ClientMachineIdentifier(identifier) => {
-            write_option_len_prefixed_string(&mut writer, &identifier)?
+            write_option_len_prefixed_buf(&mut writer, &identifier)?
         }
         DhcpOption::Unknown { data, .. } => {
             wire::write_u8(&mut writer, u8::try_from(data.len()).unwrap())?;
@@ -212,3 +412,8 @@ fn write_option_len_prefixed_string<W: Write>(mut writer: W, s: &str) -> Result<
     wire::write_u8(&mut writer, u8::try_from(s.len()).unwrap())?;
     wire::write(&mut writer, s.as_bytes())
 }
+
+fn write_option_len_prefixed_buf<W: Write>(mut writer: W, s: &[u8]) -> Result<()> {
+    wire::write_u8(&mut writer, u8::try_from(s.len()).unwrap())?;
+    wire::write(&mut writer, s)
+}