2 files changed, 52 insertions, 2 deletions

diff --git a/src/lib.rs b/src/lib.rs
index ade6954..5fb8c1d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -24,11 +24,12 @@ pub enum Opcodes {
 pub struct TftpOptions {
     blksize: usize,
     timeout: u8,
-    tsize: usize,
+    tsize: u64,
 }
 
 pub struct Tftp {
     options: TftpOptions,
+    progress_cb: Option<fn(u64, u64, u64) -> u64>,
 }
 
 fn default_options() -> TftpOptions {
@@ -43,6 +44,7 @@ impl Default for Tftp {
     fn default() -> Tftp {
         Tftp {
             options: default_options(),
+            progress_cb: None,
         }
     }
 }
@@ -67,6 +69,17 @@ impl Tftp {
         Some((val, len))
     }
 
+    pub fn set_progress_callback(&mut self, cb: fn(u64, u64, u64) -> u64) {
+        self.progress_cb = Some(cb);
+    }
+
+    fn transfer_size(&self, file: &File) -> u64 {
+        match file.metadata() {
+            Ok(ref m) if m.len() > 0 => m.len(),
+            _ => self.options.tsize,
+        }
+    }
+
     pub fn append_option(&self, buf: &mut Vec<u8>, key: &str, val: &str) {
         buf.extend(key.bytes());
         buf.push(0x00);
@@ -278,6 +291,9 @@ impl Tftp {
 
     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);
 
         loop {
             let mut filebuf = vec![0; self.options.blksize];
@@ -313,6 +329,11 @@ impl Tftp {
                 return Err(io::Error::new(io::ErrorKind::TimedOut, "ack timeout"));
             }
 
+            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 */
                 break;
@@ -326,6 +347,9 @@ impl Tftp {
 
     pub fn recv_file(&self, sock: &UdpSocket, file: &mut File) -> Result<(), io::Error> {
         let mut block_nr: u16 = 1;
+        let mut prog_update = 0;
+        let mut transferred = 0;
+        let tsize = self.transfer_size(file);
 
         loop {
             let mut buf = vec![0; 4 + self.options.blksize + 1]; // +1 for later size check
@@ -362,6 +386,11 @@ impl Tftp {
             let databuf = &buf[4..len];
             file.write_all(databuf)?;
 
+            transferred += (len - 4) as u64;
+            if let Some(cb) = self.progress_cb {
+                prog_update = cb(transferred, tsize, prog_update);
+            }
+
             self.send_ack(&sock, block_nr)?;
             block_nr = block_nr.wrapping_add(1);
 
diff --git a/src/tftpc.rs b/src/tftpc.rs
index c6b4c0f..689d055 100644
--- a/src/tftpc.rs
+++ b/src/tftpc.rs
@@ -5,7 +5,7 @@
 
 use std::env;
 use std::fs::File;
-use std::io;
+use std::io::{self, Write};
 use std::net::{SocketAddr, ToSocketAddrs, UdpSocket};
 use std::path::{Path, PathBuf};
 use std::time::Duration;
@@ -32,6 +32,26 @@ struct Tftpc {
     conf: Configuration,
 }
 
+fn update_progress(current: u64, total: u64, last: u64) -> u64 {
+    if total == 0 {
+        /* unknown; remote does not support tsize */
+        return 0;
+    }
+    let onepercent = total / 100;
+    if current < total && current < last + onepercent {
+        /* not enough progress to warrant an update */
+        return last;
+    }
+
+    let percent = 100 * current / total;
+    print!("\r {}% ", percent);
+    io::stdout().flush().expect("flushing stdout failed");
+    if current == total {
+        print!("\r");
+    }
+    current
+}
+
 impl Tftpc {
     pub fn new(conf: Configuration) -> Tftpc {
         Tftpc {
@@ -207,6 +227,7 @@ impl Tftpc {
     }
 
     pub fn start(&mut self) {
+        self.tftp.set_progress_callback(update_progress);
         let socket = UdpSocket::bind("[::]:0").expect("binding failed");
         socket.set_read_timeout(Some(Duration::from_secs(5))).expect("setting socket timeout failed");