From 9f0dfe60739152f19e2393c124deec996aa5adad Mon Sep 17 00:00:00 2001
From: altugbakan <altugbakan@gmail.com>
Date: Mon, 20 Mar 2023 18:46:35 +0300
Subject: [PATCH] Fix data retransmission on repeating ACKs

---
 Cargo.lock    |  2 +-
 Cargo.toml    |  2 +-
 src/server.rs |  4 ++--
 src/worker.rs | 21 +++++++++++----------
 4 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 41f82d0..d0ec0d7 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4,4 +4,4 @@ version = 3
 
 [[package]]
 name = "tftpd"
-version = "0.1.4"
+version = "0.1.5"
diff --git a/Cargo.toml b/Cargo.toml
index fa6af94..e9e22c3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "tftpd"
-version = "0.1.4"
+version = "0.1.5"
 authors = ["Altuğ Bakan <mail@alt.ug>"]
 edition = "2021"
 description = "Multithreaded TFTP server daemon"
diff --git a/src/server.rs b/src/server.rs
index e58d859..dcbb4a1 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -97,7 +97,7 @@ impl Server {
                 "file access violation",
             ),
             ErrorCode::FileExists => Ok(Worker::send(
-                self.socket.local_addr().unwrap(),
+                self.socket.local_addr()?,
                 *to,
                 file_path.to_path_buf(),
                 options.to_vec(),
@@ -127,7 +127,7 @@ impl Server {
                 "file access violation",
             ),
             ErrorCode::FileNotFound => Ok(Worker::receive(
-                self.socket.local_addr().unwrap(),
+                self.socket.local_addr()?,
                 *to,
                 file_path.to_path_buf(),
                 options.to_vec(),
diff --git a/src/worker.rs b/src/worker.rs
index d63e8d5..812c0ca 100644
--- a/src/worker.rs
+++ b/src/worker.rs
@@ -5,7 +5,7 @@ use std::{
     net::{SocketAddr, UdpSocket},
     path::PathBuf,
     thread,
-    time::Duration,
+    time::{Duration, SystemTime},
 };
 
 use crate::{ErrorCode, Message, OptionType, Packet, TransferOption};
@@ -111,7 +111,7 @@ fn send_file(
     file_path: &PathBuf,
     options: &mut Vec<TransferOption>,
 ) -> Result<(), Box<dyn Error>> {
-    let mut file = File::open(file_path).unwrap();
+    let mut file = File::open(file_path)?;
     let worker_options = parse_options(options, &WorkType::Send(file.metadata()?.len()))?;
 
     let mut block_number = 1;
@@ -120,8 +120,12 @@ fn send_file(
         let size = file.read(&mut chunk)?;
 
         let mut retry_cnt = 0;
+        let mut time = SystemTime::now() - Duration::from_secs(DEFAULT_TIMEOUT_SECS);
         loop {
-            Message::send_data(socket, block_number, chunk[..size].to_vec())?;
+            if time.elapsed()? >= Duration::from_secs(DEFAULT_TIMEOUT_SECS) {
+                Message::send_data(socket, block_number, chunk[..size].to_vec())?;
+                time = SystemTime::now();
+            }
 
             match Message::recv(socket) {
                 Ok(Packet::Ack(received_block_number)) => {
@@ -131,7 +135,7 @@ fn send_file(
                     }
                 }
                 Ok(Packet::Error { code, msg }) => {
-                    return Err(format!("Received error code {code}, with message {msg}").into());
+                    return Err(format!("Received error code {code}: {msg}").into());
                 }
                 _ => {
                     retry_cnt += 1;
@@ -184,15 +188,12 @@ fn receive_file(
                 Ok(Packet::Error { code, msg }) => {
                     return Err(format!("Received error code {code}: {msg}").into());
                 }
-                Err(err) => {
+                _ => {
                     retry_cnt += 1;
                     if retry_cnt == MAX_RETRIES {
-                        return Err(
-                            format!("Transfer timed out after {MAX_RETRIES} tries: {err}").into(),
-                        );
+                        return Err(format!("Transfer timed out after {MAX_RETRIES} tries").into());
                     }
                 }
-                _ => {}
             }
         }
 
@@ -205,7 +206,7 @@ fn receive_file(
     println!(
         "Received {} from {}",
         file_path.file_name().unwrap().to_str().unwrap(),
-        socket.peer_addr().unwrap()
+        socket.peer_addr()?
     );
     Ok(())
 }