1 files changed, 287 insertions, 90 deletions

diff --git a/src/tftp.rs b/src/tftp.rs
index d986a44..72bac22 100644
--- a/src/tftp.rs
+++ b/src/tftp.rs
@@ -1,7 +1,7 @@
 use std::{
+    collections::HashMap,
     io::{Cursor, Read as _, Result, Write},
-    net::UdpSocket,
-    path::{Path, PathBuf},
+    net::SocketAddr,
     str::FromStr,
 };
 
@@ -9,6 +9,9 @@ use crate::wire;
 
 pub const PORT: u16 = 69;
 
+const DEFAULT_BLOCK_SIZE: u64 = 512;
+const MAX_BLOCK_SIZE: usize = 2048;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct InvalidTftpOp(u16);
 
@@ -91,7 +94,7 @@ impl FromStr for TftpMode {
 
 #[derive(Debug)]
 pub enum TftpPacket {
-    Request(TftpRequestPacket),
+    ReadRequest(TftpReadRequestPacket),
     Data(TftpDataPacket),
     Ack(TftpAckPacket),
     OAck(TftpOAckPacket),
@@ -99,9 +102,13 @@ pub enum TftpPacket {
 }
 
 impl TftpPacket {
+    pub fn parse(buf: &[u8]) -> Result<Self> {
+        parse_packet(buf)
+    }
+
     pub fn write<W: Write>(&self, writer: W) -> Result<()> {
         match self {
-            TftpPacket::Request(p) => p.write(writer),
+            TftpPacket::ReadRequest(p) => p.write(writer),
             TftpPacket::Data(p) => p.write(writer),
             TftpPacket::Ack(p) => p.write(writer),
             TftpPacket::OAck(p) => p.write(writer),
@@ -111,14 +118,14 @@ impl TftpPacket {
 }
 
 #[derive(Debug)]
-pub struct TftpRequestPacket {
+pub struct TftpReadRequestPacket {
     pub filename: String,
     pub mode: TftpMode,
     pub tsize: Option<u64>,
     pub blksize: Option<u64>,
 }
 
-impl TftpRequestPacket {
+impl TftpReadRequestPacket {
     pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
         todo!()
     }
@@ -183,16 +190,85 @@ impl TftpOAckPacket {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum TftpErrorCode {
+    Undefined,
+    FileNotFound,
+    AccessViolation,
+    DiskFull,
+    IllegalOperation,
+    UnknownTransferId,
+    FileAreadyExists,
+    NoSuchUser,
+    Unknown(u16),
+}
+
+impl std::fmt::Display for TftpErrorCode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{:?} ({})", *self, u16::from(*self))
+    }
+}
+
+impl TftpErrorCode {
+    pub const CODE_UNDEFINED: u16 = 0;
+    pub const CODE_FILE_NOT_FOUND: u16 = 1;
+    pub const CODE_ACCESS_VIOLATION: u16 = 2;
+    pub const CODE_DISK_FULL: u16 = 3;
+    pub const CODE_ILLEGAL_OPERATION: u16 = 4;
+    pub const CODE_UNKNOWN_TRANSFER_ID: u16 = 5;
+    pub const CODE_FILE_ALREADY_EXISTS: u16 = 6;
+    pub const CODE_NO_SUCH_USER: u16 = 7;
+}
+
+impl From<u16> for TftpErrorCode {
+    fn from(value: u16) -> Self {
+        match value {
+            Self::CODE_UNDEFINED => Self::Undefined,
+            Self::CODE_FILE_NOT_FOUND => Self::FileNotFound,
+            Self::CODE_ACCESS_VIOLATION => Self::AccessViolation,
+            Self::CODE_DISK_FULL => Self::DiskFull,
+            Self::CODE_ILLEGAL_OPERATION => Self::IllegalOperation,
+            Self::CODE_UNKNOWN_TRANSFER_ID => Self::UnknownTransferId,
+            Self::CODE_FILE_ALREADY_EXISTS => Self::FileAreadyExists,
+            Self::CODE_NO_SUCH_USER => Self::NoSuchUser,
+            unknown => Self::Unknown(unknown),
+        }
+    }
+}
+
+impl From<TftpErrorCode> for u16 {
+    fn from(value: TftpErrorCode) -> Self {
+        match value {
+            TftpErrorCode::Undefined => TftpErrorCode::CODE_UNDEFINED,
+            TftpErrorCode::FileNotFound => TftpErrorCode::CODE_FILE_NOT_FOUND,
+            TftpErrorCode::AccessViolation => TftpErrorCode::CODE_ACCESS_VIOLATION,
+            TftpErrorCode::DiskFull => TftpErrorCode::CODE_DISK_FULL,
+            TftpErrorCode::IllegalOperation => TftpErrorCode::CODE_ILLEGAL_OPERATION,
+            TftpErrorCode::UnknownTransferId => TftpErrorCode::CODE_UNKNOWN_TRANSFER_ID,
+            TftpErrorCode::FileAreadyExists => TftpErrorCode::CODE_FILE_ALREADY_EXISTS,
+            TftpErrorCode::NoSuchUser => TftpErrorCode::CODE_NO_SUCH_USER,
+            TftpErrorCode::Unknown(code) => code,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct TftpErrorPacket {
-    pub code: u16,
+    pub code: TftpErrorCode,
     pub message: String,
 }
 
 impl TftpErrorPacket {
+    pub fn new(code: TftpErrorCode, message: impl Into<String>) -> Self {
+        Self {
+            code,
+            message: message.into(),
+        }
+    }
+
     pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
         wire::write_u16(&mut writer, TftpOp::Error.into())?;
-        wire::write_u16(&mut writer, self.code)?;
+        wire::write_u16(&mut writer, u16::from(self.code))?;
         wire::write_null_terminated_string(&mut writer, &self.message)?;
         Ok(())
     }
@@ -224,7 +300,7 @@ pub fn parse_packet(buf: &[u8]) -> Result<TftpPacket> {
                 }
             }
 
-            Ok(TftpPacket::Request(TftpRequestPacket {
+            Ok(TftpPacket::ReadRequest(TftpReadRequestPacket {
                 filename,
                 mode,
                 tsize,
@@ -243,7 +319,7 @@ pub fn parse_packet(buf: &[u8]) -> Result<TftpPacket> {
             Ok(TftpPacket::Ack(TftpAckPacket { block }))
         }
         TftpOp::Error => {
-            let code = wire::read_u16(&mut cursor)?;
+            let code = TftpErrorCode::from(wire::read_u16(&mut cursor)?);
             let message = wire::read_null_terminated_string(&mut cursor)?;
             Ok(TftpPacket::Error(TftpErrorPacket { code, message }))
         }
@@ -267,101 +343,222 @@ pub fn parse_packet(buf: &[u8]) -> Result<TftpPacket> {
     }
 }
 
-pub fn serve(dir: impl AsRef<Path>) -> Result<()> {
-    let dir = dir.as_ref();
-    let socket = UdpSocket::bind(format!("0.0.0.0:{PORT}"))?;
+pub trait FileSystem {
+    fn stat(&self, filename: &str) -> Result<u64>;
+    fn read(&self, filename: &str, offset: u64, buf: &mut [u8]) -> Result<u64>;
+}
 
-    // TODO: this needs to be done per addr
-    let mut last_blksize = 512u64;
-    let mut current_file = PathBuf::default();
+#[derive(Debug)]
+pub struct StaticFileSystem {
+    files: &'static [(&'static str, &'static [u8])],
+}
+
+impl StaticFileSystem {
+    pub fn new(files: &'static [(&'static str, &'static [u8])]) -> Self {
+        Self { files }
+    }
 
-    loop {
-        let mut buf = [0u8; 1500];
-        let (n, addr) = socket.recv_from(&mut buf)?;
-        let packet = parse_packet(&buf[..n]).unwrap();
+    fn find_file(&self, filename: &str) -> Result<&'static [u8]> {
+        self.files
+            .iter()
+            .find(|(name, _)| *name == filename)
+            .map(|(_, contents)| *contents)
+            .ok_or_else(|| std::io::Error::new(std::io::ErrorKind::NotFound, "file not found"))
+    }
+}
+
+impl FileSystem for StaticFileSystem {
+    fn stat(&self, filename: &str) -> Result<u64> {
+        let file = self.find_file(filename)?;
+        Ok(u64::try_from(file.len()).unwrap())
+    }
+
+    fn read(&self, filename: &str, offset: u64, buf: &mut [u8]) -> Result<u64> {
+        let file = self.find_file(filename)?;
+        let offset = usize::try_from(offset).unwrap();
+        if offset >= file.len() {
+            return Ok(0);
+        }
+
+        let rem = &file[offset..];
+        let copy_n = rem.len().min(buf.len());
+        buf[..copy_n].copy_from_slice(&rem[..copy_n]);
+        Ok(u64::try_from(copy_n).unwrap())
+    }
+}
 
-        let response = match packet {
-            TftpPacket::Request(req) => {
+#[derive(Debug)]
+struct Client {
+    blksize: u64,
+    filename: Option<String>,
+}
+
+impl Default for Client {
+    fn default() -> Self {
+        Self {
+            blksize: DEFAULT_BLOCK_SIZE,
+            filename: None,
+        }
+    }
+}
+
+pub enum ServerCommand {
+    Send(TftpPacket),
+    Ignore,
+}
+
+impl ServerCommand {
+    fn error(code: TftpErrorCode, message: impl Into<String>) -> Self {
+        Self::Send(TftpPacket::Error(TftpErrorPacket::new(code, message)))
+    }
+}
+
+#[derive(Debug, Default)]
+pub struct Server {
+    clients: HashMap<SocketAddr, Client>,
+}
+
+impl Server {
+    pub fn process(
+        &mut self,
+        fs: &dyn FileSystem,
+        source: SocketAddr,
+        buf: &[u8],
+    ) -> ServerCommand {
+        let packet = match TftpPacket::parse(buf) {
+            Ok(packet) => packet,
+            Err(err) => {
+                return ServerCommand::Send(TftpPacket::Error(TftpErrorPacket::new(
+                    TftpErrorCode::Undefined,
+                    format!("invalid packet: {err}"),
+                )));
+            }
+        };
+
+        match packet {
+            TftpPacket::ReadRequest(req) => self.process_read_req(fs, source, &req),
+            TftpPacket::Ack(ack) => self.process_ack(fs, source, &ack),
+            TftpPacket::Error(err) => {
                 println!(
-                    "Request options: tsize={:?}, blksize={:?}",
-                    req.tsize, req.blksize
+                    "received error from client {}: ({}) {}",
+                    source, err.code, err.message
                 );
+                self.clients.remove(&source);
+                ServerCommand::Ignore
+            }
+            TftpPacket::Data(_) | TftpPacket::OAck(_) => ServerCommand::Ignore,
+        }
+    }
 
-                let filepath = dir.join(req.filename);
-                current_file = filepath.clone();
-                let meta = std::fs::metadata(&filepath).unwrap();
-                let actual_file_size = meta.len();
+    fn process_read_req(
+        &mut self,
+        fs: &dyn FileSystem,
+        source: SocketAddr,
+        req: &TftpReadRequestPacket,
+    ) -> ServerCommand {
+        println!(
+            "Request options: tsize={:?}, blksize={:?}",
+            req.tsize, req.blksize
+        );
+
+        let client = self.clients.entry(source).or_default();
+        client.filename = Some(req.filename.clone());
+
+        // Only send OACK if client sent options
+        if req.tsize.is_some() || req.blksize.is_some() {
+            if let Some(blksize) = req.blksize {
+                client.blksize = blksize;
+            }
 
-                // Only send OACK if client sent options
-                if req.tsize.is_some() || req.blksize.is_some() {
-                    if let Some(blksize) = req.blksize {
-                        last_blksize = blksize;
+            let tsize_response = if req.tsize.is_some() {
+                let filesize = match fs.stat(&req.filename) {
+                    Ok(filesize) => filesize,
+                    Err(err) => {
+                        return ServerCommand::error(
+                            TftpErrorCode::Undefined,
+                            format!("failed to obtain file size: {}", err),
+                        );
                     }
+                };
 
-                    let tsize_response = if req.tsize.is_some() {
-                        Some(actual_file_size)
-                    } else {
-                        None
-                    };
-
-                    Some(TftpPacket::OAck(TftpOAckPacket {
-                        tsize: tsize_response,
-                        blksize: req.blksize,
-                    }))
-                } 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
-                    };
-
-                    Some(TftpPacket::Data(TftpDataPacket {
-                        block: 1,
-                        data: first_block,
-                    }))
-                }
-            }
-            TftpPacket::Data(dat) => unimplemented!(),
-            TftpPacket::Ack(ack) => {
-                println!("Received ACK packet: block {}", ack.block);
-
-                let contents = std::fs::read(&current_file).unwrap();
-                let next_block = ack.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 || ack.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());
-                    Some(TftpPacket::Data(TftpDataPacket {
-                        block: next_block,
-                        data: block_data,
-                    }))
-                } else {
-                    None
+                Some(filesize)
+            } else {
+                None
+            };
+
+            ServerCommand::Send(TftpPacket::OAck(TftpOAckPacket {
+                tsize: tsize_response,
+                blksize: req.blksize,
+            }))
+        } else {
+            // No options, send first data block directly
+            let options = self.clients.entry(source).or_default();
+            let block_size = usize::try_from(options.blksize).unwrap();
+
+            assert!(block_size <= MAX_BLOCK_SIZE);
+            let mut contents = [0u8; MAX_BLOCK_SIZE];
+            let contents = &mut contents[..block_size];
+
+            let n = match fs.read(&req.filename, 0, contents) {
+                Ok(n) => usize::try_from(n).unwrap(),
+                Err(err) => {
+                    return ServerCommand::error(
+                        TftpErrorCode::Undefined,
+                        format!("failed to read file contents: {}", err),
+                    );
                 }
+            };
+
+            ServerCommand::Send(TftpPacket::Data(TftpDataPacket {
+                block: 1,
+                data: contents[..n].to_vec(),
+            }))
+        }
+    }
+
+    fn process_ack(
+        &mut self,
+        fs: &dyn FileSystem,
+        source: SocketAddr,
+        ack: &TftpAckPacket,
+    ) -> ServerCommand {
+        println!("Received ACK packet: block {}", ack.block);
+
+        let client = self.clients.entry(source).or_default();
+        let filename = match &client.filename {
+            Some(filename) => filename,
+            None => {
+                return ServerCommand::error(
+                    TftpErrorCode::Undefined,
+                    "unknown filename for client",
+                );
             }
-            TftpPacket::OAck(ack) => todo!(),
-            TftpPacket::Error(err) => {
-                println!(
-                    "Received ERROR packet: code {}, message: {}",
-                    err.code, err.message
+        };
+        let block_size = usize::try_from(client.blksize).unwrap();
+
+        let next_block = ack.block + 1;
+        let start_offset = (usize::from(next_block) - 1) * block_size;
+
+        let mut contents = [0u8; MAX_BLOCK_SIZE];
+        let contents = &mut contents[..block_size];
+        let n = match fs.read(filename, u64::try_from(start_offset).unwrap(), contents) {
+            Ok(n) => usize::try_from(n).unwrap(),
+            Err(err) => {
+                return ServerCommand::error(
+                    TftpErrorCode::Undefined,
+                    format!("failed to read file contents: {}", err),
                 );
-                None
             }
         };
+        let contents = &contents[..n];
 
-        if let Some(response) = response {
-            let mut writer = Cursor::new(&mut buf[..]);
-            println!("Sending to {addr}: {response:#?}");
-            response.write(&mut writer).unwrap();
-            let (response, _) = writer.split();
-            socket.send_to(&response, addr).unwrap();
+        if contents.is_empty() {
+            return ServerCommand::Ignore;
         }
+
+        ServerCommand::Send(TftpPacket::Data(TftpDataPacket {
+            block: next_block,
+            data: contents.to_vec(),
+        }))
     }
 }