wiprfc7440

1 files changed, 78 insertions, 53 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 64291cb..f3647d9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -41,6 +41,7 @@ pub struct TftpOptions {
     blksize: usize,
     timeout: Duration,
     tsize: u64,
+    windowsize: u16,
 }
 
 #[derive(Clone, Copy)]
@@ -55,6 +56,7 @@ fn default_options() -> TftpOptions {
         blksize: 512,
         timeout: Duration::from_secs(3),
         tsize: 0,
+        windowsize: 1,
     }
 }
 
@@ -240,12 +242,12 @@ impl Tftp {
         std::io::Error::new(kind, error)
     }
 
-    fn wait_for_ack(&self, sock: &UdpSocket, expected_block: u16) -> Result<bool, io::Error> {
+    fn wait_for_ack(&self, sock: &UdpSocket) -> Result<Option<u16>, io::Error> {
         let mut buf = [0; 512];
         let len = match sock.recv(&mut buf) {
             Ok(l) => l,
             Err(ref error) if [io::ErrorKind::WouldBlock, io::ErrorKind::TimedOut].contains(&error.kind()) => {
-                return Ok(false);
+                return Ok(None);
             }
             Err(err) => return Err(err),
         };
@@ -257,13 +259,13 @@ impl Tftp {
         let opcode = u16::from_be_bytes([buf[0], buf[1]]);
         let block_nr = u16::from_be_bytes([buf[2], buf[3]]);
 
-        if opcode == Opcode::ACK as u16 && block_nr == expected_block {
-            return Ok(true);
+        if opcode == Opcode::ACK as u16 {
+            return Ok(Some(block_nr));
         } else if opcode == Opcode::ERROR as u16 {
             return Err(self.parse_error(&buf[4..]));
         }
 
-        Ok(false)
+        Ok(None)
     }
 
     pub fn ack_options(&self, sock: &UdpSocket, options: &HashMap<String, String>, ackwait: bool) -> Result<(), io::Error> {
@@ -287,9 +289,9 @@ impl Tftp {
             if !ackwait {
                 return Ok(());
             }
-            match self.wait_for_ack(&sock, 0) {
-                Ok(true) => return Ok(()),
-                Ok(false) => continue,
+            match self.wait_for_ack(&sock) {
+                Ok(Some(0)) => return Ok(()),
+                Ok(_) => continue,
                 Err(e) => return Err(e),
             };
         }
@@ -339,6 +341,13 @@ impl Tftp {
                     }
                     _ => false,
                 },
+                "windowsize" => match val.parse() {
+                    Ok(t) if t >= 1 => {
+                        self.options.windowsize = t;
+                        true
+                    }
+                    _ => false,
+                },
                 _ => false,
             }
         });
@@ -408,7 +417,7 @@ impl Tftp {
         Ok(())
     }
 
-    pub fn send_ack(&self, sock: &UdpSocket, block_nr: u16) -> Result<(), io::Error> {
+    fn send_ack(&self, sock: &UdpSocket, block_nr: u16) -> Result<(), io::Error> {
         self._send_ack(sock, None, block_nr)
     }
 
@@ -416,11 +425,39 @@ impl Tftp {
         self._send_ack(sock, Some(cl), block_nr)
     }
 
+    fn send_block(&self, socket: &UdpSocket, block: Vec<u8>, block_nr: u16, overflow: &mut Vec<u8>) -> Result<usize, io::Error> {
+        let mut len = block.len();
+        let mut block = block;
+        match self.mode {
+            Mode::OCTET => {},
+            Mode::NETASCII => {
+                overflow.extend(octet_to_netascii(&block));
+                block = overflow.clone();
+                if overflow.len() > self.options.blksize {
+                    *overflow = block.split_off(self.options.blksize);
+                } else {
+                    overflow.clear();
+                }
+                len = block.len();
+            }
+        }
+
+        let mut sendbuf = Vec::with_capacity(4 + len);
+        sendbuf.extend((Opcode::DATA as u16).to_be_bytes().iter());
+        sendbuf.extend(block_nr.to_be_bytes().iter());
+        sendbuf.extend(block.iter());
+        socket.send(&sendbuf)?;
+
+        Ok(len)
+    }
+
     pub fn send_file(&self, socket: &UdpSocket, file: &mut File) -> Result<(), io::Error> {
         let mut block_nr: u16 = 1;
-        let mut transferred = 0;
-        let mut prog_update = 0;
-        let tsize = self.transfer_size(file);
+        let mut acked_block: u16 = 0;
+        //let mut transferred = 0;
+        //let mut prog_update = 0;
+        let mut seek_back = None;
+        //let tsize = self.transfer_size(file);
 
         /* holds bytes from netascii conversion that did not fit in tx buffer */
         let mut overflow = Vec::with_capacity(2 * self.options.blksize);
@@ -436,49 +473,32 @@ impl Tftp {
                 }
             };
 
-            /* take care of netascii conversion */
-            let mut databuf = filebuf[0..len].to_vec();
-            match self.mode {
-                Mode::OCTET => {},
-                Mode::NETASCII => {
-                    overflow.extend(octet_to_netascii(&databuf));
-                    databuf = overflow.clone();
-                    if overflow.len() > self.options.blksize {
-                        overflow = databuf.split_off(self.options.blksize);
-                    } else {
-                        overflow.clear();
+            len = self.send_block(&socket, filebuf[0..len].to_vec(), block_nr, &mut overflow)?;
+            if block_nr - acked_block == self.options.windowsize {
+                match self.wait_for_ack(&socket) {
+                    Ok(Some(b)) if b > acked_block => {
+                        acked_block = b;
+                    },
+                    Ok(Some(b)) => {
+                        // seek/rewind to last acked block and 'continue'
+                        // retry counter (for timeouts)
+                        block_nr = b;
+                        seek_back = Some(42);
+                    },
+                    Ok(_) => {
+                        // TODO: retry counter
+                        continue;
                     }
-                    len = databuf.len();
+                    Err(e) => return Err(e),
                 }
             }
 
-            let mut sendbuf = Vec::with_capacity(4 + len);
-            sendbuf.extend((Opcode::DATA as u16).to_be_bytes().iter());
-            sendbuf.extend(block_nr.to_be_bytes().iter());
-            sendbuf.extend(databuf.iter());
-
-            let mut acked = false;
-            for _ in 1..5 {
-                /* try a couple of times to send data, in case of timeouts
-                   or re-ack of previous data */
-                socket.send(&sendbuf)?;
-                match self.wait_for_ack(&socket, block_nr) {
-                    Ok(true) => {
-                        acked = true;
-                        break;
-                    }
-                    Ok(false) => continue,
-                    Err(e) => return Err(e),
-                };
-            }
-            if !acked {
-                return Err(io::Error::new(io::ErrorKind::TimedOut, "ack timeout"));
-            }
 
-            transferred += len as u64;
+            // TODO: update for windowsize
+            /*transferred += len as u64;
             if let Some(cb) = self.progress_cb {
                 prog_update = cb(transferred, tsize, prog_update);
-            }
+            }*/
 
             if len < self.options.blksize {
                 /* this was the last block */
@@ -493,13 +513,14 @@ impl Tftp {
 
     pub fn recv_file(&self, sock: &UdpSocket, file: &mut File) -> Result<(), io::Error> {
         let mut block_nr: u16 = 1;
+        let mut acked_block: u16 = 0;
         let mut prog_update = 0;
         let mut transferred = 0;
         let mut netascii_state = false;
         let tsize = self.transfer_size(file);
 
         loop {
-            let mut buf = vec![0; 4 + self.options.blksize + 1]; // +1 for later size check
+            let mut buf = vec![0; 4 + self.options.blksize];
             let mut len = 0;
 
             for _ in 1..5 {
@@ -514,8 +535,7 @@ impl Tftp {
                 };
                 break;
             }
-            if len < 4 || len > 4 + self.options.blksize {
-                /* max size: 2 + 2 + blksize */
+            if len < 4 {
                 return Err(io::Error::new(io::ErrorKind::InvalidInput, "unexpected size"));
             }
 
@@ -546,10 +566,15 @@ impl Tftp {
                 prog_update = cb(transferred, tsize, prog_update);
             }
 
-            self.send_ack(&sock, block_nr)?;
+            let is_last_block = len < 4 + self.options.blksize;
+
+            if block_nr - acked_block == self.options.windowsize || is_last_block {
+                self.send_ack(&sock, block_nr)?;
+                acked_block = block_nr;
+            }
             block_nr = block_nr.wrapping_add(1);
 
-            if len < 4 + self.options.blksize {
+            if is_last_block {
                 break;
             }
         }