1 files changed, 294 insertions, 0 deletions
diff --git a/src/tftp.rs b/src/tftp.rs
new file mode 100644
index 0000000..67dcd1a
--- /dev/null
+++ b/src/tftp.rs
@@ -0,0 +1,294 @@
+/*
+ * Copyright 2019 Reiner Herrmann <reiner@reiner-h.de>
+ * License: GPL-3+
+ */
+
+use std::net::UdpSocket;
+use std::fs::OpenOptions;
+use std::fs::File;
+use std::collections::HashMap;
+use std::path::{Path,PathBuf};
+use std::io;
+use std::io::prelude::*;
+
+pub struct TftpOptions {
+
+}
+
+pub struct Tftp {
+}
+
+impl Tftp {
+    pub fn new() -> Tftp {
+        Tftp{}
+    }
+
+    fn get_tftp_str(&self, buf: &[u8]) -> Option<(String, usize)> {
+        let mut iter = buf.iter();
+
+        let len = match iter.position(|&x| x == 0) {
+            Some(l) => l,
+            None => return None,
+        };
+        let val = match String::from_utf8(buf[0 .. len].to_vec()) {
+            Ok(v) => v,
+            Err(_) => return None,
+        };
+
+        return Some((val, len));
+    }
+
+    fn wait_for_ack(&self, sock: &UdpSocket, expected_block: u16) -> Result<bool, io::Error> {
+        let mut buf = [0; 4];
+        match sock.recv(&mut buf) {
+            Ok(_) => (),
+            Err(ref error) if [io::ErrorKind::WouldBlock, io::ErrorKind::TimedOut].contains(&error.kind()) => {
+                return Ok(false);
+            }
+            Err(err) => return Err(err),
+        };
+
+        let opcode = u16::from_be_bytes([buf[0], buf[1]]);
+        let block_nr = u16::from_be_bytes([buf[2], buf[3]]);
+
+        if opcode == 4 && block_nr == expected_block {
+            return Ok(true)
+        }
+
+        Ok(false)
+    }
+
+    fn ack_options(&self, sock: &UdpSocket, options: &HashMap<String, String>, ackwait: bool) -> Result<(), io::Error> {
+        if options.is_empty() {
+            return Ok(())
+        }
+
+        let mut buf = Vec::with_capacity(512);
+        buf.extend([0x00, 0x06].iter());  // opcode
+
+        for (key, val) in options {
+            buf.extend(key.bytes());
+            buf.push(0x00);
+            buf.extend(val.bytes());
+            buf.push(0x00);
+        }
+
+        for _ in 1..5 {
+            sock.send(&buf)?;
+            if !ackwait {
+                return Ok(());
+            }
+            match self.wait_for_ack(&sock, 0) {
+                Ok(true) => return Ok(()),
+                Ok(false) => continue,
+                Err(e) => return Err(e),
+            };
+        }
+
+        Err(io::Error::new(io::ErrorKind::TimedOut, "ack timeout"))
+    }
+
+    pub fn init_tftp_options(&self, sock: &UdpSocket, options: &mut HashMap<String, String>, ackwait: bool) -> Result<TftpOptions, io::Error> {
+        let tftpopts = TftpOptions {};
+
+        options.retain(|key, _val| {
+            match key.as_str() {
+                "placeholder_option" => {
+                    true
+                }
+                _ => false
+            }
+        });
+
+        self.ack_options(&sock, &options, ackwait)?;
+
+        return Ok(tftpopts);
+    }
+
+    fn parse_options(&self, buf: &[u8]) -> HashMap<String, String> {
+        let mut options = HashMap::new();
+
+        let mut pos = 0;
+        loop {
+            let (key, len) = match self.get_tftp_str(&buf[pos ..]) {
+                Some(args) => args,
+                None => break,
+            };
+            pos += len + 1;
+
+            let (val, len) = match self.get_tftp_str(&buf[pos ..]) {
+                Some(args) => args,
+                None => break,
+            };
+            pos += len + 1;
+
+            options.insert(key, val);
+        }
+
+        return options;
+    }
+
+    pub fn parse_file_mode_options(&self, buf: &[u8]) -> Result<(PathBuf, String, HashMap<String, String>), io::Error> {
+        let dataerr = io::Error::new(io::ErrorKind::InvalidData, "invalid data received");
+
+        let mut pos = 0;
+        let (filename, len) = match self.get_tftp_str(&buf[pos ..]) {
+            Some(args) => args,
+            None => return Err(dataerr),
+        };
+        pos += len + 1;
+
+        let filename = Path::new(&filename);
+
+        let (mode, len) = match self.get_tftp_str(&buf[pos ..]) {
+            Some(args) => args,
+            None => return Err(dataerr),
+        };
+        pos += len + 1;
+
+        let options = self.parse_options(&buf[pos ..]);
+
+        Ok((filename.to_path_buf(), mode, options))
+    }
+
+    pub fn send_error(&self, socket: &UdpSocket, code: u16, msg: &str) -> Result<(), io::Error> {
+        let mut buf = vec![0x00, 0x05];  // opcode
+        buf.extend(code.to_be_bytes().iter());
+        buf.extend(msg.as_bytes());
+
+        socket.send(&buf)?;
+        Ok(())
+    }
+
+    pub fn send_ack(&self, sock: &UdpSocket, block_nr: u16) -> Result<(), io::Error> {
+        let mut buf = vec![0x00, 0x04];  // opcode
+        buf.extend(block_nr.to_be_bytes().iter());
+
+        sock.send(&buf)?;
+
+        Ok(())
+    }
+
+    pub fn send_file(&self, socket: &UdpSocket, path: &Path) -> Result<(), io::Error> {
+        let mut file = match File::open(path) {
+            Ok(f) => f,
+            Err(ref error) if error.kind() == io::ErrorKind::NotFound => {
+                self.send_error(&socket, 1, "File not found")?;
+                return Err(io::Error::new(io::ErrorKind::NotFound, "file not found"));
+            },
+            Err(_) => {
+                self.send_error(&socket, 2, "Permission denied")?;
+                return Err(io::Error::new(io::ErrorKind::PermissionDenied, "permission denied"));
+            }
+        };
+        if !file.metadata()?.is_file() {
+            self.send_error(&socket, 1, "File not found")?;
+            return Err(io::Error::new(io::ErrorKind::NotFound, "file not found"));
+        }
+
+        let mut block_nr: u16 = 1;
+
+        loop {
+            let mut filebuf = [0; 512];
+            let len = match file.read(&mut filebuf) {
+                Ok(n) => n,
+                Err(ref error) if error.kind() == io::ErrorKind::Interrupted => continue, /* retry */
+                Err(err) => {
+                    self.send_error(&socket, 0, "File reading error")?;
+                    return Err(err);
+                }
+            };
+
+            let mut sendbuf = Vec::with_capacity(4 + len);
+            sendbuf.extend([0x00, 0x03].iter());  // opcode
+            sendbuf.extend(block_nr.to_be_bytes().iter());
+            sendbuf.extend(filebuf[0..len].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"))
+            }
+
+            if len < 512 {
+                /* this was the last block */
+                break;
+            }
+
+            /* increment with rollover on overflow */
+            block_nr = block_nr.wrapping_add(1);
+        }
+        Ok(())
+    }
+
+    pub fn recv_file(&self, sock: &UdpSocket, path: &PathBuf) -> Result<(), io::Error> {
+        let mut file = match OpenOptions::new().write(true).create_new(true).open(path) {
+            Ok(f) => f,
+            Err(ref err) if err.kind() == io::ErrorKind::AlreadyExists => {
+                return Err(io::Error::new(err.kind(), "already exists"));
+            },
+            Err(_) => return Err(io::Error::new(io::ErrorKind::PermissionDenied, "permission denied")),
+        };
+
+        let mut block_nr = 0;
+
+        loop {
+            let mut buf = [0; 1024];
+            let mut len = 0;
+
+            for _ in 1..5 {
+                self.send_ack(&sock, block_nr)?;
+                len = match sock.recv(&mut buf) {
+                    Ok(n) => n,
+                    Err(ref error) if [io::ErrorKind::WouldBlock, io::ErrorKind::TimedOut].contains(&error.kind()) => {
+                        /* re-ack and try to recv again */
+                        continue;
+                    }
+                    Err(err) => return Err(err),
+                };
+                break;
+            }
+            if len > 516 || len < 4 {
+                /* max size: 2 + 2 + 512 */
+                return Err(io::Error::new(io::ErrorKind::InvalidInput, "unexpected size"));
+            }
+
+            let _opcode = match u16::from_be_bytes([buf[0], buf[1]]) {
+                3 /* DATA */ => (),
+                _ => return Err(io::Error::new(io::ErrorKind::Other, "unexpected opcode")),
+            };
+            let nr = u16::from_be_bytes([buf[2], buf[3]]);
+            if nr != block_nr.wrapping_add(1) {
+                /* already received or packets were missed, re-acknowledge */
+                continue;
+            }
+            block_nr = nr;
+
+            let databuf = &buf[4..len];
+            file.write_all(databuf)?;
+
+            if len < 516 {
+                break;
+            }
+        }
+
+        file.flush()?;
+
+        self.send_ack(&sock, block_nr)?;
+
+        Ok(())
+    }
+
+}