2 files changed, 4 insertions, 295 deletions

diff --git a/src/main.rs b/src/main.rs
index 3d86a8e..07e4eec 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -347,223 +347,18 @@ fn write_boot_ack(xid: u32, chaddr: [u8; 16], client_uuid: Option<Vec<u8>>) -> R
     Ok(writer)
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-struct InvalidTftpOp(u16);
-
-impl std::fmt::Display for InvalidTftpOp {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "invalid tftp opcode '{}'", self.0)
-    }
-}
-
-impl std::error::Error for InvalidTftpOp {}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-enum TftpOp {
-    ReadRequest,
-    WriteRequest,
-    Data,
-    Ack,
-    Error,
-    Oack,
-}
-
-impl Into<u16> for TftpOp {
-    fn into(self) -> u16 {
-        match self {
-            TftpOp::ReadRequest => 1,
-            TftpOp::WriteRequest => 2,
-            TftpOp::Data => 3,
-            TftpOp::Ack => 4,
-            TftpOp::Error => 5,
-            TftpOp::Oack => 6,
-        }
-    }
-}
-
-impl TryFrom<u16> for TftpOp {
-    type Error = InvalidTftpOp;
-
-    fn try_from(value: u16) -> std::result::Result<Self, InvalidTftpOp> {
-        match value {
-            1 => Ok(Self::ReadRequest),
-            2 => Ok(Self::WriteRequest),
-            3 => Ok(Self::Data),
-            4 => Ok(Self::Ack),
-            5 => Ok(Self::Error),
-            6 => Ok(Self::Oack),
-            unknown => Err(InvalidTftpOp(unknown)),
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-enum TftpMode {
-    NetAscii,
-    Octet,
-    Mail,
-}
-
-#[derive(Debug)]
-struct TftpRequestPacket {
-    filename: String,
-    mode: String,
-    tsize: Option<u64>,
-    blksize: Option<u64>,
-}
-
-#[derive(Debug)]
-struct TftpDataPacket {
-    block: u16,
-    data: Vec<u8>,
-}
-
-#[derive(Debug)]
-struct TftpAckPacket {
-    block: u16,
-}
-
-#[derive(Debug)]
-struct TftpErrorPacket {
-    code: u16,
-    message: String,
-}
-
-fn tftp_request_packet_parse(cursor: &mut Cursor<&[u8]>) -> Result<TftpRequestPacket> {
-    let filename = read_null_terminated_string(cursor)?;
-    let mode = read_null_terminated_string(cursor)?;
-    let mut tsize = None;
-    let mut blksize = None;
-    while let Ok(opt_name) = read_null_terminated_string(cursor) {
-        if opt_name.is_empty() {
-            break;
-        }
-        let opt_data = read_null_terminated_string(cursor)?;
-        match opt_name.as_str() {
-            "tsize" => tsize = Some(opt_data.parse::<u64>().unwrap()),
-            "blksize" => blksize = Some(opt_data.parse::<u64>().unwrap()),
-            _ => eprintln!("unknown tftp request option '{opt_name}'"),
-        }
-    }
-
-    Ok(TftpRequestPacket {
-        filename,
-        mode,
-        tsize,
-        blksize,
-    })
-}
-
-fn tftp_data_packet_write(writer: &mut Vec<u8>, block: u16, data: Vec<u8>) -> Result<()> {
-    write_u16(writer, TftpOp::Data.into())?;
-    write_u16(writer, block)?;
-    write_buf(writer, &data)?;
-    Ok(())
-}
-
-fn tftp_oack_packet_write(
-    writer: &mut Vec<u8>,
-    tsize: Option<u64>,
-    blksize: Option<u64>,
-) -> Result<()> {
-    write_u16(writer, TftpOp::Oack.into())?;
-
-    // Only include options that were requested by the client
-    if let Some(blksize_val) = blksize {
-        write_buf(writer, b"blksize")?;
-        write_u8(writer, 0)?; // null terminator
-        let blksize_str = blksize_val.to_string();
-        write_buf(writer, blksize_str.as_bytes())?;
-        write_u8(writer, 0)?; // null terminator
-    }
-
-    if let Some(tsize_val) = tsize {
-        write_buf(writer, b"tsize")?;
-        write_u8(writer, 0)?; // null terminator
-        let tsize_str = tsize_val.to_string();
-        write_buf(writer, tsize_str.as_bytes())?;
-        write_u8(writer, 0)?; // null terminator
-    }
-
-    Ok(())
-}
-
-#[derive(Debug)]
-enum TftpPacket {
-    Request(TftpRequestPacket),
-    Data(TftpDataPacket),
-    Ack(TftpAckPacket),
-    Error(TftpErrorPacket),
-}
-
-fn tftp_packet_parse(cursor: &mut Cursor<&[u8]>) -> Result<TftpPacket> {
-    let op = TftpOp::try_from(read_u16(cursor)?).unwrap();
-    match op {
-        TftpOp::ReadRequest | TftpOp::WriteRequest => {
-            let filename = read_null_terminated_string(cursor)?;
-            let mode = read_null_terminated_string(cursor)?;
-            let mut tsize = None;
-            let mut blksize = None;
-
-            while let Ok(opt_name) = read_null_terminated_string(cursor) {
-                if opt_name.is_empty() {
-                    break;
-                }
-                let opt_data = read_null_terminated_string(cursor)?;
-                match opt_name.as_str() {
-                    "tsize" => tsize = Some(opt_data.parse::<u64>().unwrap()),
-                    "blksize" => blksize = Some(opt_data.parse::<u64>().unwrap()),
-                    _ => eprintln!("unknown tftp request option '{opt_name}'"),
-                }
-            }
-
-            Ok(TftpPacket::Request(TftpRequestPacket {
-                filename,
-                mode,
-                tsize,
-                blksize,
-            }))
-        }
-        TftpOp::Data => {
-            let block = read_u16(cursor)?;
-            let mut data = Vec::new();
-            cursor.read_to_end(&mut data)?;
-            Ok(TftpPacket::Data(TftpDataPacket { block, data }))
-        }
-        TftpOp::Ack => {
-            let block = read_u16(cursor)?;
-            Ok(TftpPacket::Ack(TftpAckPacket { block }))
-        }
-        TftpOp::Error => {
-            let code = read_u16(cursor)?;
-            let message = read_null_terminated_string(cursor)?;
-            Ok(TftpPacket::Error(TftpErrorPacket { code, message }))
-        }
-        TftpOp::Oack => {
-            // OACK parsing not implemented for now
-            Err(std::io::Error::new(
-                std::io::ErrorKind::Unsupported,
-                "OACK parsing not implemented",
-            ))
-        }
-    }
-}
-
 fn main() {
     let socket67 = UdpSocket::bind("0.0.0.0:67").unwrap();
     socket67.set_broadcast(true).unwrap();
     socket67.set_nonblocking(true).unwrap();
 
-    let socket69 = UdpSocket::bind("0.0.0.0:69").unwrap();
-    socket69.set_broadcast(false).unwrap();
-    socket69.set_nonblocking(true).unwrap();
-
     let socket4011 = UdpSocket::bind("0.0.0.0:4011").unwrap();
     socket4011.set_broadcast(true).unwrap();
     socket4011.set_nonblocking(true).unwrap();
 
-    let mut last_blksize = 512u64;
-    let mut current_file = String::new();
+    std::thread::spawn(|| {
+        tftp::serve("tftp").unwrap();
+    });
 
     loop {
         let mut buf = [0u8; 1500];
@@ -575,91 +370,6 @@ fn main() {
         } else if let Ok((n, addr)) = socket4011.recv_from(&mut buf) {
             println!("Received {} bytes from {} on port 4011", n, addr);
             handle_packet_4011(&buf[..n], &socket4011, addr);
-        } else if let Ok((n, addr)) = socket69.recv_from(&mut buf) {
-            let mut cursor = Cursor::new(&buf[..n]);
-
-            let packet = tftp_packet_parse(&mut cursor).unwrap();
-            println!("Received TFTP request from {addr}: {packet:#?}");
-
-            let mut response = Vec::default();
-            match packet {
-                TftpPacket::Request(tftp_request_packet) => {
-                    println!(
-                        "Request options: tsize={:?}, blksize={:?}",
-                        tftp_request_packet.tsize, tftp_request_packet.blksize
-                    );
-
-                    let filepath = format!("tftp/{}", tftp_request_packet.filename);
-                    current_file = filepath.clone();
-                    let meta = std::fs::metadata(&filepath).unwrap();
-                    let actual_file_size = meta.len();
-
-                    // Only send OACK if client sent options
-                    if tftp_request_packet.tsize.is_some() || tftp_request_packet.blksize.is_some()
-                    {
-                        if let Some(blksize) = tftp_request_packet.blksize {
-                            last_blksize = blksize;
-                        }
-
-                        let tsize_response = if tftp_request_packet.tsize.is_some() {
-                            Some(actual_file_size)
-                        } else {
-                            None
-                        };
-
-                        tftp_oack_packet_write(
-                            &mut response,
-                            tsize_response,
-                            tftp_request_packet.blksize,
-                        )
-                        .unwrap();
-                    } else {
-                        // No options, send first data block directly
-                        let contents = std::fs::read(&filepath).unwrap();
-                        let block_size = 512;
-                        let first_block = if contents.len() > block_size {
-                            contents[..block_size].to_vec()
-                        } else {
-                            contents
-                        };
-
-                        tftp_data_packet_write(&mut response, 1, first_block).unwrap();
-                    }
-                }
-                TftpPacket::Data(tftp_data_packet) => {
-                    println!("Received DATA packet: block {}", tftp_data_packet.block);
-                }
-                TftpPacket::Ack(tftp_ack_packet) => {
-                    println!("Received ACK packet: block {}", tftp_ack_packet.block);
-
-                    let contents = std::fs::read(&current_file).unwrap();
-                    let next_block = tftp_ack_packet.block + 1;
-                    let start_offset = (next_block - 1) as u64 * last_blksize;
-                    let end_offset = next_block as u64 * last_blksize;
-                    let prev_start_offset = (next_block.saturating_sub(2)) as u64 * last_blksize;
-                    let prev_remain = contents.len() - prev_start_offset as usize;
-                    if prev_remain as u64 >= last_blksize || tftp_ack_packet.block == 0 {
-                        let end = std::cmp::min(end_offset as usize, contents.len());
-                        let block_data = contents[start_offset as usize..end].to_vec();
-                        println!("sending tftp data packet with {} bytes", block_data.len());
-                        tftp_data_packet_write(&mut response, next_block, block_data).unwrap();
-                    }
-                }
-                TftpPacket::Error(tftp_error_packet) => {
-                    println!(
-                        "Received ERROR packet: code {}, message: {}",
-                        tftp_error_packet.code, tftp_error_packet.message
-                    );
-                }
-            }
-            //
-            // let filepath = format!("tftp/{}", request.filename);
-            // let meta = std::fs::metadata(&filepath).unwrap();
-            // let contents = std::fs::read(&filepath).unwrap();
-            // let mut response = Vec::default();
-            if !response.is_empty() {
-                socket69.send_to(&response, addr).unwrap();
-            }
         } else {
             std::thread::sleep(std::time::Duration::from_millis(10));
         }
diff --git a/src/tftp.rs b/src/tftp.rs
index 681d036..d986a44 100644
--- a/src/tftp.rs
+++ b/src/tftp.rs
@@ -305,7 +305,7 @@ pub fn serve(dir: impl AsRef<Path>) -> Result<()> {
                     };
 
                     Some(TftpPacket::OAck(TftpOAckPacket {
-                        tsize: req.tsize,
+                        tsize: tsize_response,
                         blksize: req.blksize,
                     }))
                 } else {
@@ -361,7 +361,6 @@ pub fn serve(dir: impl AsRef<Path>) -> Result<()> {
             println!("Sending to {addr}: {response:#?}");
             response.write(&mut writer).unwrap();
             let (response, _) = writer.split();
-            println!("Sending {} bytes to {addr}: {response:#?}", response.len());
             socket.send_to(&response, addr).unwrap();
         }
     }