dhcp module split done

1 files changed, 34 insertions, 348 deletions

diff --git a/src/main.rs b/src/main.rs
index 321b5ea..51bbd77 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -3,351 +3,17 @@ pub mod dhcp;
 pub mod tftp;
 pub mod wire;
 
-use std::io::{BufRead, Cursor, Read, Result, Write};
-use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};
+use std::{
+    io::Result,
+    net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket},
+};
 
 use ipnet::Ipv4Net;
 
-const FLAG_BROADCAST: u16 = 1 << 15;
+use crate::dhcp::{DhcpOption, DhcpPacket};
 
-const OPTION_CODE_PAD: u8 = 0;
-const OPTION_CODE_END: u8 = 255;
-const OPTION_CODE_VENDOR_CLASS_IDENTIFIER: u8 = 60;
-const OPTION_CODE_USER_CLASS_INFORMATION: u8 = 77;
-
-const MAGIC_COOKIE: [u8; 4] = [0x63, 0x82, 0x53, 0x63];
-
-//const BOOT_FILE_NAME: &[u8] = b"pxelinux.0";
-//const BOOT_FILE_NAME: &[u8] = b"debian-installer/amd64/bootnetx64.efi";
-const BOOT_FILE_NAME: &[u8] = b"ipxe.efi";
-const BOOT_FILE_NAME_IPXE: &[u8] = b"test.ipxe";
-
-const LOCAL_IPV4: Ipv4Addr = Ipv4Addr::new(192, 168, 1, 100);
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-enum BootOp {
-    Request,
-    Reply,
-}
-
-impl BootOp {
-    pub const OP_REQUEST: u8 = 1;
-    pub const OP_REPLY: u8 = 2;
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-enum HardwareType {
-    Ethernet,
-}
-
-impl HardwareType {
-    pub const TYPE_ETHER: u8 = 1;
-    pub const LEN_ETHER: u8 = 6;
-}
-
-#[derive(Debug, Clone)]
-enum DhcpOption {
-    Pad,
-    End,
-    VendorClassIdentifier(String),
-    UserClassInformation(String),
-    Unknown { code: u8, data: Vec<u8> },
-}
-
-#[derive(Debug)]
-struct DhcpPacket {
-    op: BootOp,
-    htype: HardwareType,
-    hlen: u8,
-    hops: u8, // should be zero
-    xid: u32,
-    secs: u16,
-    flags: u16,
-    ciaddr: Ipv4Addr,
-    yiaddr: Ipv4Addr,
-    siaddr: Ipv4Addr,
-    giaddr: Ipv4Addr,
-    chaddr: [u8; 16],
-    sname: [u8; 64],
-    file: [u8; 128],
-    options: Vec<DhcpOption>,
-}
-
-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_null_terminated_vec(cursor: &mut Cursor<&[u8]>) -> Result<Vec<u8>> {
-    let mut buf = Vec::default();
-    cursor.read_until(0, &mut buf)?;
-    buf.pop();
-    Ok(buf)
-}
-
-fn read_null_terminated_string(cursor: &mut Cursor<&[u8]>) -> Result<String> {
-    let buf = read_null_terminated_vec(cursor)?;
-    Ok(String::from_utf8(buf).unwrap())
-}
-
-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> {
-    match read_u8(cursor)? {
-        HardwareType::TYPE_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 {
-        OPTION_CODE_PAD => DhcpOption::Pad,
-        OPTION_CODE_END => DhcpOption::End,
-        OPTION_CODE_VENDOR_CLASS_IDENTIFIER => {
-            DhcpOption::VendorClassIdentifier(read_len8_prefixed_string(cursor)?)
-        }
-        OPTION_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 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)?,
-        hlen: read_u8(&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_arr(&mut cursor)?,
-        file: read_arr(&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)
-}
-
-fn write_buf(writer: &mut Vec<u8>, buf: &[u8]) -> Result<()> {
-    writer.write_all(buf)
-}
-
-fn write_u8(writer: &mut Vec<u8>, v: u8) -> Result<()> {
-    write_buf(writer, &[v])
-}
-
-fn write_u16(writer: &mut Vec<u8>, v: u16) -> Result<()> {
-    let buf = u16::to_be_bytes(v);
-    write_buf(writer, &buf)
-}
-
-fn write_u32(writer: &mut Vec<u8>, v: u32) -> Result<()> {
-    let buf = u32::to_be_bytes(v);
-    write_buf(writer, &buf)
-}
-
-fn write_ipv4(writer: &mut Vec<u8>, v: Ipv4Addr) -> Result<()> {
-    write_buf(writer, &v.octets())
-}
-
-fn write_boot_packet(
-    xid: u32,
-    chaddr: [u8; 16],
-    client_uuid: Option<Vec<u8>>,
-    ipxe: bool,
-) -> Result<Vec<u8>> {
-    let mut writer = Vec::default();
-    write_u8(&mut writer, BootOp::OP_REPLY)?;
-    write_u8(&mut writer, HardwareType::TYPE_ETHER)?;
-    write_u8(&mut writer, 6)?;
-    write_u8(&mut writer, 0)?;
-    write_u32(&mut writer, xid)?;
-    write_u16(&mut writer, 0)?;
-    write_u16(&mut writer, FLAG_BROADCAST)?;
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // ciaddr
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // yiaddr  
-    write_ipv4(&mut writer, LOCAL_IPV4)?; // siaddr (TFTP server)
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // giaddr
-    write_buf(&mut writer, &chaddr)?;
-    write_buf(&mut writer, &[0u8; 64])?;
-    write_buf(&mut writer, &[0u8; 128])?;
-    write_buf(&mut writer, &MAGIC_COOKIE)?;
-
-    // Option 53: DHCP Message Type (DHCPOFFER)
-    write_u8(&mut writer, 53)?;
-    write_u8(&mut writer, 1)?;
-    write_u8(&mut writer, 2)?; // DHCPOFFER
-
-    // Option 54: DHCP Server Identifier
-    write_u8(&mut writer, 54)?;
-    write_u8(&mut writer, 4)?;
-    write_ipv4(&mut writer, LOCAL_IPV4)?; // Your server IP
-
-    // Option 60: Vendor Class Identifier
-    const PXE_CLIENT: &[u8] = b"PXEClient";
-    write_u8(&mut writer, 60)?;
-    write_u8(&mut writer, PXE_CLIENT.len() as u8)?;
-    write_buf(&mut writer, PXE_CLIENT)?;
-
-    // Option 97: Client Machine Identifier (UUID from client)
-    if let Some(uuid) = client_uuid {
-        write_u8(&mut writer, 97)?;
-        write_u8(&mut writer, uuid.len() as u8)?;
-        write_buf(&mut writer, &uuid)?;
-    }
-
-    // TFTP server name
-    const SERVER_NAME: &[u8] = b"diogos-air";
-    write_u8(&mut writer, 66)?;
-    write_u8(&mut writer, SERVER_NAME.len() as u8)?;
-    write_buf(&mut writer, SERVER_NAME)?;
-
-    write_u8(&mut writer, 67)?;
-    if !ipxe {
-        write_u8(&mut writer, BOOT_FILE_NAME.len() as u8)?;
-        write_buf(&mut writer, BOOT_FILE_NAME)?;
-    } else {
-        write_u8(&mut writer, BOOT_FILE_NAME_IPXE.len() as u8)?;
-        write_buf(&mut writer, BOOT_FILE_NAME_IPXE)?;
-    }
-
-    // Option 255: End
-    write_u8(&mut writer, 255)?;
-
-    Ok(writer)
-}
-
-fn write_boot_ack(xid: u32, chaddr: [u8; 16], client_uuid: Option<Vec<u8>>) -> Result<Vec<u8>> {
-    let mut writer = Vec::default();
-    write_u8(&mut writer, BootOp::OP_REPLY)?;
-    write_u8(&mut writer, HardwareType::TYPE_ETHER)?;
-    write_u8(&mut writer, 6)?;
-    write_u8(&mut writer, 0)?;
-    write_u32(&mut writer, xid)?;
-    write_u16(&mut writer, 0)?;
-    write_u16(&mut writer, 0)?;
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // ciaddr
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // yiaddr  
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // siaddr (TFTP server)
-    write_ipv4(&mut writer, Ipv4Addr::UNSPECIFIED)?; // giaddr
-    write_buf(&mut writer, &chaddr)?;
-    write_buf(&mut writer, &[0u8; 64])?;
-    write_buf(&mut writer, &[0u8; 128])?;
-    write_buf(&mut writer, &MAGIC_COOKIE)?;
-
-    // Option 53: DHCP Message Type (DHCPOFFER)
-    write_u8(&mut writer, 53)?;
-    write_u8(&mut writer, 1)?;
-    write_u8(&mut writer, 5)?; // DHCPACK
-
-    // Option 54: DHCP Server Identifier
-    write_u8(&mut writer, 54)?;
-    write_u8(&mut writer, 4)?;
-    write_ipv4(&mut writer, LOCAL_IPV4)?; // Your server IP
-
-    // Option 60: Vendor Class Identifier
-    const PXE_CLIENT: &[u8] = b"PXEClient";
-    write_u8(&mut writer, 60)?;
-    write_u8(&mut writer, PXE_CLIENT.len() as u8)?;
-    write_buf(&mut writer, PXE_CLIENT)?;
-
-    // Option 97: Client Machine Identifier (UUID from client)
-    if let Some(uuid) = client_uuid {
-        write_u8(&mut writer, 97)?;
-        write_u8(&mut writer, uuid.len() as u8)?;
-        write_buf(&mut writer, &uuid)?;
-    }
-
-    // TFTP server name
-    const SERVER_NAME: &[u8] = b"diogos-air";
-    write_u8(&mut writer, 66)?;
-    write_u8(&mut writer, SERVER_NAME.len() as u8)?;
-    write_buf(&mut writer, SERVER_NAME)?;
-
-    // TFTP file name
-    write_u8(&mut writer, 67)?;
-    write_u8(&mut writer, BOOT_FILE_NAME.len() as u8)?;
-    write_buf(&mut writer, BOOT_FILE_NAME)?;
-
-    write_u8(&mut writer, 71)?;
-    write_u8(&mut writer, 4)?;
-    write_buf(&mut writer, &[0, 0, 0, 0])?;
-
-    // Option 255: End
-    write_u8(&mut writer, 255)?;
-
-    Ok(writer)
-}
+const LOCAL_IPV4: Ipv4Addr = Ipv4Addr::new(192, 168, 1, 103);
+const LOCAL_HOSTNAME: &'static str = "Diogos-Air";
 
 #[derive(Debug, Clone)]
 struct InterfaceAddr {
@@ -433,7 +99,7 @@ fn main() {
 }
 
 fn handle_packet(buf: &[u8], socket: &UdpSocket) {
-    match parse_packet(buf) {
+    match dhcp::parse_packet(buf) {
         Ok(packet) => {
             println!("Parsed DHCP packet: XID={:08x}", packet.xid);
 
@@ -464,10 +130,21 @@ fn handle_packet(buf: &[u8], socket: &UdpSocket) {
 
             if is_pxe {
                 println!("Responding to PXE client with DHCPOFFER");
-                let response =
-                    write_boot_packet(packet.xid, packet.chaddr, client_uuid, is_ipxe).unwrap();
+                let mut response_buf = Vec::default();
+                let response = DhcpPacket::new_boot(
+                    packet.xid,
+                    packet.chaddr,
+                    client_uuid.unwrap(),
+                    LOCAL_IPV4,
+                    LOCAL_HOSTNAME.to_string(),
+                    match is_ipxe {
+                        true => "test.ipxe".to_string(),
+                        false => "ipxe.efi".to_string(),
+                    },
+                );
+                response.write(&mut response_buf).unwrap();
                 socket
-                    .send_to(&response, SocketAddrV4::new(Ipv4Addr::BROADCAST, 68))
+                    .send_to(&response_buf, SocketAddrV4::new(Ipv4Addr::BROADCAST, 68))
                     .unwrap();
             } else {
                 println!("Not a PXE client, ignoring");
@@ -480,7 +157,7 @@ fn handle_packet(buf: &[u8], socket: &UdpSocket) {
 }
 
 fn handle_packet_4011(buf: &[u8], socket: &UdpSocket, sender_addr: SocketAddr) {
-    match parse_packet(buf) {
+    match dhcp::parse_packet(buf) {
         Ok(packet) => {
             println!("Parsed DHCP packet on 4011: XID={:08x}", packet.xid);
 
@@ -494,8 +171,17 @@ fn handle_packet_4011(buf: &[u8], socket: &UdpSocket, sender_addr: SocketAddr) {
             }
 
             println!("Responding with DHCPACK");
-            let response = write_boot_ack(packet.xid, packet.chaddr, client_uuid).unwrap();
-            socket.send_to(&response, sender_addr).unwrap();
+            let mut response_buf = Vec::default();
+            let response = DhcpPacket::new_boot_ack(
+                packet.xid,
+                packet.chaddr,
+                client_uuid.unwrap(),
+                LOCAL_IPV4,
+                LOCAL_HOSTNAME.to_string(),
+                "ipxe.efi".to_string(),
+            );
+            response.write(&mut response_buf).unwrap();
+            socket.send_to(&response_buf, sender_addr).unwrap();
         }
         Err(e) => {
             println!("Failed to parse packet on 4011: {}", e);