From 68d05e3a11ee8bf5236cdac87225e3518d918e5a Mon Sep 17 00:00:00 2001
From: altugbakan <altugbakan@gmail.com>
Date: Wed, 8 Mar 2023 21:12:57 +0300
Subject: [PATCH] Add work on worker and server

---
 src/config.rs  |   4 +-
 src/main.rs    |   6 +-
 src/message.rs |  27 ++++++--
 src/packet.rs  |  63 ++++++++++++++++--
 src/server.rs  | 177 ++++++++++++++++++++++++-------------------------
 src/worker.rs  |  94 ++++++++++++++++++++++++--
 6 files changed, 259 insertions(+), 112 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 32ee90d..89332dc 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -41,7 +41,7 @@ impl Config {
                 "-d" | "--directory" => {
                     if let Some(dir_str) = args.next() {
                         if !Path::new(&dir_str).exists() {
-                            return Err(format!("{} does not exist", dir_str).into());
+                            return Err(format!("{dir_str} does not exist").into());
                         }
                         config.directory = PathBuf::from(dir_str);
                     } else {
@@ -60,7 +60,7 @@ impl Config {
                     println!("  -h, --help\t\t\tPrint help information");
                     process::exit(0);
                 }
-                invalid => return Err(format!("invalid flag: {}", invalid).into()),
+                invalid => return Err(format!("invalid flag: {invalid}").into()),
             }
         }
 
diff --git a/src/main.rs b/src/main.rs
index ab31ba9..3b6827d 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,8 +16,10 @@ fn main() {
     });
 
     println!(
-        "Running TFTP Server on {}:{}",
-        config.ip_address, config.port
+        "Running TFTP Server on {}:{} in {}",
+        config.ip_address,
+        config.port,
+        config.directory.display()
     );
 
     server.listen();
diff --git a/src/message.rs b/src/message.rs
index 85ae55d..ed89460 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -3,7 +3,7 @@ use std::{
     net::{SocketAddr, UdpSocket},
 };
 
-use crate::packet::{ErrorCode, Opcode, Option};
+use crate::packet::{ErrorCode, Opcode, Packet, TransferOption};
 
 pub struct Message;
 
@@ -35,24 +35,37 @@ impl Message {
     }
 
     pub fn send_error_to(socket: &UdpSocket, to: &SocketAddr, code: ErrorCode, msg: &str) {
+        eprintln!("{msg}");
         if socket.send_to(&get_error_buf(code, msg), to).is_err() {
             eprintln!("could not send an error message");
         }
     }
 
-    pub fn send_oack_to(
+    pub fn send_oack(
         socket: &UdpSocket,
-        to: &SocketAddr,
-        options: Vec<Option>,
+        options: &Vec<TransferOption>,
     ) -> Result<(), Box<dyn Error>> {
-        todo!();
+        let mut buf = vec![0x00, Opcode::Oack as u8];
 
-        let buf = [];
+        for option in options {
+            buf = [buf, option.as_bytes()].concat();
+        }
 
-        socket.send_to(&buf, to)?;
+        socket.send(&buf)?;
 
         Ok(())
     }
+
+    pub fn receive_ack(socket: &UdpSocket) -> Result<u16, Box<dyn Error>> {
+        let mut buf = [0; 4];
+        socket.recv(&mut buf)?;
+
+        if let Ok(Packet::Ack(block)) = Packet::deserialize(&buf) {
+            Ok(block)
+        } else {
+            Err("invalid ack".into())
+        }
+    }
 }
 
 fn get_error_buf(code: ErrorCode, msg: &str) -> Vec<u8> {
diff --git a/src/packet.rs b/src/packet.rs
index 6b32b8f..1791396 100644
--- a/src/packet.rs
+++ b/src/packet.rs
@@ -5,12 +5,12 @@ pub enum Packet<'a> {
     Rrq {
         filename: String,
         mode: String,
-        options: Vec<Option>,
+        options: Vec<TransferOption>,
     },
     Wrq {
         filename: String,
         mode: String,
-        options: Vec<Option>,
+        options: Vec<TransferOption>,
     },
     Data {
         block_num: u16,
@@ -32,6 +32,7 @@ impl<'a> Packet<'a> {
             Opcode::Data => parse_data(buf),
             Opcode::Ack => parse_ack(buf),
             Opcode::Error => parse_error(buf),
+            _ => Err("invalid packet".into()),
         }
     }
 }
@@ -42,6 +43,7 @@ pub enum Opcode {
     Data = 0x0003,
     Ack = 0x0004,
     Error = 0x0005,
+    Oack = 0x0006,
 }
 
 impl Opcode {
@@ -52,14 +54,52 @@ impl Opcode {
             0x0003 => Ok(Opcode::Data),
             0x0004 => Ok(Opcode::Ack),
             0x0005 => Ok(Opcode::Error),
+            0x0006 => Ok(Opcode::Oack),
             _ => Err("invalid opcode"),
         }
     }
 }
 
-pub struct Option {
-    option: String,
-    value: String,
+pub struct TransferOption {
+    pub option: OptionType,
+    pub value: usize,
+}
+
+impl TransferOption {
+    pub fn as_bytes(&self) -> Vec<u8> {
+        [
+            self.option.as_str().as_bytes(),
+            &[0x00],
+            self.value.to_string().as_bytes(),
+            &[0x00],
+        ]
+        .concat()
+    }
+}
+
+pub enum OptionType {
+    BlockSize,
+    TransferSize,
+    Timeout,
+}
+
+impl OptionType {
+    fn as_str(&self) -> &'static str {
+        match self {
+            OptionType::BlockSize => "blksize",
+            OptionType::TransferSize => "tsize",
+            OptionType::Timeout => "timeout",
+        }
+    }
+
+    fn from_str(value: &str) -> Result<Self, &'static str> {
+        match value {
+            "blksize" => Ok(OptionType::BlockSize),
+            "tsize" => Ok(OptionType::TransferSize),
+            "timeout" => Ok(OptionType::Timeout),
+            _ => Err("invalid option type".into()),
+        }
+    }
 }
 
 #[repr(u16)]
@@ -105,7 +145,13 @@ fn parse_rq(buf: &[u8], opcode: Opcode) -> Result<Packet, Box<dyn Error>> {
     while zero_index < buf.len() - 1 {
         (option, zero_index) = Convert::to_string(buf, zero_index + 1)?;
         (value, zero_index) = Convert::to_string(buf, zero_index + 1)?;
-        options.push(Option { option, value });
+
+        if let Ok(option) = OptionType::from_str(option.as_str()) {
+            options.push(TransferOption {
+                option,
+                value: value.parse()?,
+            });
+        }
     }
 
     match opcode {
@@ -139,3 +185,8 @@ fn parse_error(buf: &[u8]) -> Result<Packet, Box<dyn Error>> {
     let (msg, _) = Convert::to_string(buf, 4)?;
     Ok(Packet::Error { code, msg })
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+}
diff --git a/src/server.rs b/src/server.rs
index a1cb84d..a396c62 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,4 +1,4 @@
-use crate::packet::{ErrorCode, Packet};
+use crate::packet::{ErrorCode, Packet, TransferOption};
 use crate::{Config, Message, Worker};
 use std::error::Error;
 use std::net::{SocketAddr, UdpSocket};
@@ -28,104 +28,101 @@ impl Server {
             let mut buf = [0; MAX_REQUEST_PACKET_SIZE];
             if let Ok((number_of_bytes, from)) = self.socket.recv_from(&mut buf) {
                 if let Ok(packet) = Packet::deserialize(&buf[..number_of_bytes]) {
-                    match &packet {
-                        Packet::Rrq { filename, .. } => {
-                            match self.check_file_exists(&Path::new(&filename)) {
-                                ErrorCode::FileNotFound => {
-                                    eprintln!("requested file does not exist");
-                                    Message::send_error_to(
-                                        &self.socket,
-                                        &from,
-                                        ErrorCode::FileNotFound,
-                                        "requested file does not exist",
-                                    );
-                                }
-                                ErrorCode::AccessViolation => {
-                                    eprintln!("requested file is not in the directory");
-                                    Message::send_error_to(
-                                        &self.socket,
-                                        &from,
-                                        ErrorCode::AccessViolation,
-                                        "requested file is not in the directory",
-                                    );
-                                }
-                                ErrorCode::FileExists => self
-                                    .handle_rrq(&packet, from)
-                                    .unwrap_or_else(|_| eprintln!("could not handle read request")),
-                                _ => {}
-                            }
-                        }
-                        Packet::Wrq { filename, .. } => {
-                            match self.check_file_exists(&Path::new(&filename)) {
-                                ErrorCode::FileExists => {
-                                    eprintln!("requested file already exists");
-                                    Message::send_error_to(
-                                        &self.socket,
-                                        &from,
-                                        ErrorCode::FileExists,
-                                        "requested file already exists",
-                                    );
-                                }
-                                ErrorCode::AccessViolation => {
-                                    eprintln!("requested file is not in the directory");
-                                    Message::send_error_to(
-                                        &self.socket,
-                                        &from,
-                                        ErrorCode::AccessViolation,
-                                        "requested file is not in the directory",
-                                    );
-                                }
-                                ErrorCode::FileNotFound => {
-                                    self.handle_wrq(&packet, from).unwrap_or_else(|_| {
-                                        eprintln!("could not handle write request")
-                                    })
-                                }
-                                _ => {}
-                            }
-                        }
-                        _ => {
-                            eprintln!("invalid request packet received");
-                            Message::send_error_to(
-                                &self.socket,
-                                &from,
-                                ErrorCode::IllegalOperation,
-                                "invalid request",
-                            );
-                        }
-                    }
-                };
+                    self.handle_packet(&packet, &from)
+                }
             }
         }
     }
 
-    fn handle_rrq(&self, packet: &Packet, to: SocketAddr) -> Result<(), Box<dyn Error>> {
-        if let Packet::Rrq {
-            filename,
-            mode: _,
-            options,
-        } = packet
-        {
-            let worker = Worker::new(self.socket.local_addr().unwrap(), to)?;
-            worker.send_file(Path::new(&filename), options)?;
-        } else {
-            return Err("invalid read request packet".into());
+    fn handle_packet(&self, packet: &Packet, from: &SocketAddr) {
+        match &packet {
+            Packet::Rrq {
+                filename, options, ..
+            } => self.validate_rrq(filename, options, from),
+            Packet::Wrq {
+                filename, options, ..
+            } => self.validate_wrq(filename, options, from),
+            _ => {
+                Message::send_error_to(
+                    &self.socket,
+                    from,
+                    ErrorCode::IllegalOperation,
+                    "invalid request",
+                );
+            }
         }
+    }
+
+    fn validate_rrq(&self, filename: &String, options: &Vec<TransferOption>, to: &SocketAddr) {
+        match self.check_file_exists(&Path::new(&filename)) {
+            ErrorCode::FileNotFound => {
+                Message::send_error_to(
+                    &self.socket,
+                    to,
+                    ErrorCode::FileNotFound,
+                    "requested file does not exist",
+                );
+            }
+            ErrorCode::AccessViolation => {
+                Message::send_error_to(
+                    &self.socket,
+                    to,
+                    ErrorCode::AccessViolation,
+                    "requested file is not in the directory",
+                );
+            }
+            ErrorCode::FileExists => self
+                .handle_rrq(filename, options, to)
+                .unwrap_or_else(|err| eprintln!("could not handle read request: {err}")),
+            _ => {}
+        }
+    }
+
+    fn validate_wrq(&self, filename: &String, options: &Vec<TransferOption>, to: &SocketAddr) {
+        match self.check_file_exists(&Path::new(&filename)) {
+            ErrorCode::FileExists => {
+                Message::send_error_to(
+                    &self.socket,
+                    to,
+                    ErrorCode::FileExists,
+                    "requested file already exists",
+                );
+            }
+            ErrorCode::AccessViolation => {
+                Message::send_error_to(
+                    &self.socket,
+                    to,
+                    ErrorCode::AccessViolation,
+                    "requested file is not in the directory",
+                );
+            }
+            ErrorCode::FileNotFound => self
+                .handle_wrq(filename, options, to)
+                .unwrap_or_else(|err| eprintln!("could not handle write request: {err}")),
+            _ => {}
+        };
+    }
+
+    fn handle_rrq(
+        &self,
+        filename: &String,
+        options: &Vec<TransferOption>,
+        to: &SocketAddr,
+    ) -> Result<(), Box<dyn Error>> {
+        let mut worker = Worker::new(&self.socket.local_addr().unwrap(), to)?;
+        worker.send_file(Path::new(&filename), options)?;
 
         Ok(())
     }
 
-    fn handle_wrq(&self, packet: &Packet, to: SocketAddr) -> Result<(), Box<dyn Error>> {
-        if let Packet::Wrq {
-            filename,
-            mode: _,
-            options,
-        } = packet
-        {
-            let worker = Worker::new(self.socket.local_addr().unwrap(), to)?;
-            worker.receive_file(Path::new(&filename), options)?;
-        } else {
-            return Err("invalid write request packet".into());
-        }
+    fn handle_wrq(
+        &self,
+        filename: &String,
+        options: &Vec<TransferOption>,
+        to: &SocketAddr,
+    ) -> Result<(), Box<dyn Error>> {
+        let mut worker = Worker::new(&self.socket.local_addr().unwrap(), to)?;
+        worker.receive_file(Path::new(&filename), options)?;
 
         Ok(())
     }
diff --git a/src/worker.rs b/src/worker.rs
index 7353364..7e2e00e 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -1,27 +1,111 @@
 use std::{
     error::Error,
+    fs::File,
+    io::Read,
     net::{SocketAddr, UdpSocket},
     path::Path,
+    time::Duration,
 };
 
-use crate::packet::Option;
+use crate::{
+    packet::{OptionType, TransferOption},
+    Message,
+};
 
 pub struct Worker {
     socket: UdpSocket,
+    blk_size: usize,
+    t_size: usize,
+    timeout: usize,
 }
 
+const MAX_RETRIES: u32 = 6;
+
 impl Worker {
-    pub fn new(addr: SocketAddr, remote: SocketAddr) -> Result<Worker, Box<dyn Error>> {
+    pub fn new(addr: &SocketAddr, remote: &SocketAddr) -> Result<Worker, Box<dyn Error>> {
         let socket = UdpSocket::bind(SocketAddr::from((addr.ip(), 0)))?;
         socket.connect(remote)?;
-        Ok(Worker { socket })
+        Ok(Worker {
+            socket,
+            blk_size: 512,
+            t_size: 0,
+            timeout: 5,
+        })
     }
 
-    pub fn send_file(&self, file: &Path, options: &Vec<Option>) -> Result<(), Box<dyn Error>> {
+    pub fn send_file(
+        &mut self,
+        file: &Path,
+        options: &Vec<TransferOption>,
+    ) -> Result<(), Box<dyn Error>> {
+        let mut file = File::open(file).unwrap();
+
+        self.parse_options(options, Some(&file));
+        Message::send_oack(&self.socket, options)?;
+
+        self.socket
+            .set_write_timeout(Some(Duration::from_secs(self.timeout as u64)))?;
+
+        let mut retry_cnt = 0;
+        loop {
+            let mut chunk = Vec::with_capacity(self.blk_size);
+            let size = file
+                .by_ref()
+                .take(self.blk_size as u64)
+                .read_to_end(&mut chunk)?;
+
+            loop {
+                if Message::send_data(&self.socket, &chunk).is_err() {
+                    return Err(format!("failed to send data").into());
+                }
+
+                if let Ok(block) = Message::receive_ack(&self.socket) {
+                    todo!("handle block number");
+                } else {
+                    retry_cnt += 1;
+                    if retry_cnt == MAX_RETRIES {
+                        return Err(format!("transfer timed out after {MAX_RETRIES} tries").into());
+                    }
+                }
+            }
+
+            if size < self.blk_size {
+                break;
+            };
+        }
+
         Ok(())
     }
 
-    pub fn receive_file(&self, file: &Path, options: &Vec<Option>) -> Result<(), Box<dyn Error>> {
+    pub fn receive_file(
+        &mut self,
+        file: &Path,
+        options: &Vec<TransferOption>,
+    ) -> Result<(), Box<dyn Error>> {
+        let mut file = File::open(file).unwrap();
+
+        self.parse_options(options, Some(&file));
+        Message::send_oack(&self.socket, options)?;
+
+        todo!("file receiving");
+
         Ok(())
     }
+
+    fn parse_options(&mut self, options: &Vec<TransferOption>, file: Option<&File>) {
+        for option in options {
+            let TransferOption { option, value } = option;
+
+            match option {
+                OptionType::BlockSize => self.blk_size = *value,
+                OptionType::TransferSize => {
+                    self.t_size = match file {
+                        Some(file) => file.metadata().unwrap().len() as usize,
+                        None => *value,
+                    }
+                }
+                OptionType::Timeout => self.timeout = *value,
+            }
+        }
+    }
 }