Implement Transfer Size Option (RFC 2349, part 2)

3 files changed, 28 insertions, 5 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 6a45eeb..f9fc57e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -24,6 +24,7 @@ pub enum Opcodes {
 pub struct TftpOptions {
     blksize: usize,
     timeout: u8,
+    tsize: usize,
 }
 
 pub struct Tftp {
@@ -34,6 +35,7 @@ fn default_options() -> TftpOptions {
     TftpOptions {
         blksize: 512,
         timeout: 3,
+        tsize: 0,
     }
 }
 
@@ -142,6 +144,15 @@ impl Tftp {
                         _ => false
                     }
                 }
+                "tsize" => {
+                    match val.parse() {
+                        Ok(t) => {
+                            self.options.tsize = t;
+                            true
+                        }
+                        _ => false
+                    }
+                }
                 _ => false
             }
         });
diff --git a/src/tftpc.rs b/src/tftpc.rs
index fdc3853..294475c 100644
--- a/src/tftpc.rs
+++ b/src/tftpc.rs
@@ -90,9 +90,10 @@ impl Tftpc {
         Some(remote)
     }
 
-    fn append_option_req(&self, buf: &mut Vec<u8>) {
+    fn append_option_req(&self, buf: &mut Vec<u8>, fsize: u64) {
         self.tftp.append_option(buf, "blksize", &format!("{}", 1428));
         self.tftp.append_option(buf, "timeout", &format!("{}", 3));
+        self.tftp.append_option(buf, "tsize", &format!("{}", fsize));
     }
 
     fn handle_wrq(&mut self, sock: &UdpSocket) -> Result<(), io::Error> {
@@ -107,11 +108,18 @@ impl Tftpc {
             }
             None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")),
         };
+        let metadata = match file.metadata() {
+            Ok(m) => m,
+            Err(_) => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")),
+        };
+        if !metadata.is_file() {
+            return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename"));
+        }
 
         let mut buf = Vec::with_capacity(512);
         buf.extend((rtftp::Opcodes::WRQ as u16).to_be_bytes().iter());
         self.tftp.append_option(&mut buf, filename, "octet");
-        self.append_option_req(&mut buf);
+        self.append_option_req(&mut buf, metadata.len());
 
         let mut remote = None;
         for _ in 1 .. 3 {
@@ -160,7 +168,7 @@ impl Tftpc {
         let mut buf = Vec::with_capacity(512);
         buf.extend((rtftp::Opcodes::RRQ as u16).to_be_bytes().iter());
         self.tftp.append_option(&mut buf, filename, "octet");
-        self.append_option_req(&mut buf);
+        self.append_option_req(&mut buf, 0);
 
         let mut remote = None;
         for _ in 1 .. 3 {
diff --git a/src/tftpd.rs b/src/tftpd.rs
index c424f7b..c85603d 100644
--- a/src/tftpd.rs
+++ b/src/tftpd.rs
@@ -76,7 +76,6 @@ impl Tftpd {
     fn handle_wrq(&mut self, socket: &UdpSocket, cl: &SocketAddr, buf: &[u8]) -> Result<(), io::Error> {
         let (filename, mode, mut options) = self.tftp.parse_file_mode_options(buf)?;
         self.tftp.init_tftp_options(&socket, &mut options)?;
-        self.tftp.ack_options(&socket, &options, false)?;
 
         match mode.as_ref() {
             "octet" => (),
@@ -109,6 +108,7 @@ impl Tftpd {
             }
         };
 
+        self.tftp.ack_options(&socket, &options, false)?;
         match self.tftp.recv_file(&socket, &mut file) {
             Ok(_) => println!("Received {} from {}.", path.display(), cl),
             Err(ref err) => {
@@ -123,7 +123,6 @@ impl Tftpd {
     fn handle_rrq(&mut self, socket: &UdpSocket, cl: &SocketAddr, buf: &[u8]) -> Result<(), io::Error> {
         let (filename, mode, mut options) = self.tftp.parse_file_mode_options(buf)?;
         self.tftp.init_tftp_options(&socket, &mut options)?;
-        self.tftp.ack_options(&socket, &options, true)?;
 
         match mode.as_ref() {
             "octet" => (),
@@ -157,6 +156,11 @@ impl Tftpd {
             self.tftp.send_error(&socket, 1, "File not found")?;
             return Err(io::Error::new(io::ErrorKind::NotFound, "file not found"));
         }
+
+        if let Some(opt) = options.get_mut("tsize") {
+            *opt = file.metadata()?.len().to_string();
+        }
+        self.tftp.ack_options(&socket, &options, true)?;
         match self.tftp.send_file(&socket, &mut file) {
             Ok(_) => println!("Sent {} to {}.", path.display(), cl),
             Err(err) => println!("Sending {} to {} failed ({}).", path.display(), cl, err.to_string()),