From 65bf28fbfcd91f430a8f3d1dbf669d7c398aac50 Mon Sep 17 00:00:00 2001 From: Reiner Herrmann Date: Sat, 2 Mar 2019 18:29:45 +0100 Subject: Move main tftp module into lib to be used as crate Otherwise compiler throws warnings about unused symbols, if not every public function is used in every module importing it. --- src/lib.rs | 304 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ src/tftp.rs | 304 ----------------------------------------------------------- src/tftpc.rs | 6 +- src/tftpd.rs | 6 +- 4 files changed, 310 insertions(+), 310 deletions(-) create mode 100644 src/lib.rs delete mode 100644 src/tftp.rs diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..b1ff232 --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,304 @@ +/* + * Copyright 2019 Reiner Herrmann + * License: GPL-3+ + */ + +use std::net::UdpSocket; +use std::fs::File; +use std::collections::HashMap; +use std::path::{Path,PathBuf}; +use std::io; +use std::io::prelude::*; +use std::time::Duration; + +pub struct TftpOptions { + blksize: usize, + timeout: u8, +} + +pub struct Tftp { + options: TftpOptions, +} + +fn default_options() -> TftpOptions { + TftpOptions { + blksize: 512, + timeout: 3, + } +} + +impl Tftp { + + pub fn new() -> Tftp { + Tftp{ + options: default_options(), + } + } + + 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 { + 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, ackwait: bool) -> Result<(), io::Error> { + if options.is_empty() { + if !ackwait { + /* it's a WRQ, send normal ack to start transfer */ + self.send_ack(&sock, 0)?; + } + 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(&mut self, sock: &UdpSocket, options: &mut HashMap, ackwait: bool) -> Result<(), io::Error> { + self.options = default_options(); + + options.retain(|key, val| { + let val = val.to_lowercase(); + match key.to_lowercase().as_str() { + "blksize" => { + match val.parse() { + Ok(b) if b >= 8 && b <= 65464 => { + self.options.blksize = b; + true + } + _ => false + } + } + "timeout" => { + match val.parse() { + Ok(t) if t >= 1 => { + self.options.timeout = t; + true + } + _ => false + } + } + _ => false + } + }); + + sock.set_read_timeout(Some(Duration::from_secs(self.options.timeout as u64)))?; + + self.ack_options(&sock, &options, ackwait)?; + + return Ok(()); + } + + fn parse_options(&self, buf: &[u8]) -> HashMap { + 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), 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((m,l)) => (m.to_lowercase(), l), + 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, file: &mut File) -> Result<(), io::Error> { + let mut block_nr: u16 = 1; + + loop { + let mut filebuf = vec![0; self.options.blksize]; + 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 < self.options.blksize { + /* 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, file: &mut File) -> Result<(), io::Error> { + let mut block_nr: u16 = 1; + + loop { + let mut buf = vec![0; 4 + self.options.blksize + 1]; // +1 for later size check + let mut len = 0; + + for _ in 1..5 { + len = match sock.recv(&mut buf) { + Ok(n) => n, + Err(ref error) if [io::ErrorKind::WouldBlock, io::ErrorKind::TimedOut].contains(&error.kind()) => { + /* re-ack previous and try to recv again */ + self.send_ack(&sock, block_nr - 1)?; + continue; + } + Err(err) => return Err(err), + }; + break; + } + if len < 4 || len > 4 + self.options.blksize { + /* max size: 2 + 2 + blksize */ + 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")), + }; + if u16::from_be_bytes([buf[2], buf[3]]) != block_nr { + /* already received or packets were missed, re-acknowledge */ + self.send_ack(&sock, block_nr - 1)?; + continue; + } + + let databuf = &buf[4..len]; + file.write_all(databuf)?; + + self.send_ack(&sock, block_nr)?; + block_nr = block_nr.wrapping_add(1); + + if len < 4 + self.options.blksize { + break; + } + } + + file.flush()?; + + Ok(()) + } + +} diff --git a/src/tftp.rs b/src/tftp.rs deleted file mode 100644 index b1ff232..0000000 --- a/src/tftp.rs +++ /dev/null @@ -1,304 +0,0 @@ -/* - * Copyright 2019 Reiner Herrmann - * License: GPL-3+ - */ - -use std::net::UdpSocket; -use std::fs::File; -use std::collections::HashMap; -use std::path::{Path,PathBuf}; -use std::io; -use std::io::prelude::*; -use std::time::Duration; - -pub struct TftpOptions { - blksize: usize, - timeout: u8, -} - -pub struct Tftp { - options: TftpOptions, -} - -fn default_options() -> TftpOptions { - TftpOptions { - blksize: 512, - timeout: 3, - } -} - -impl Tftp { - - pub fn new() -> Tftp { - Tftp{ - options: default_options(), - } - } - - 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 { - 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, ackwait: bool) -> Result<(), io::Error> { - if options.is_empty() { - if !ackwait { - /* it's a WRQ, send normal ack to start transfer */ - self.send_ack(&sock, 0)?; - } - 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(&mut self, sock: &UdpSocket, options: &mut HashMap, ackwait: bool) -> Result<(), io::Error> { - self.options = default_options(); - - options.retain(|key, val| { - let val = val.to_lowercase(); - match key.to_lowercase().as_str() { - "blksize" => { - match val.parse() { - Ok(b) if b >= 8 && b <= 65464 => { - self.options.blksize = b; - true - } - _ => false - } - } - "timeout" => { - match val.parse() { - Ok(t) if t >= 1 => { - self.options.timeout = t; - true - } - _ => false - } - } - _ => false - } - }); - - sock.set_read_timeout(Some(Duration::from_secs(self.options.timeout as u64)))?; - - self.ack_options(&sock, &options, ackwait)?; - - return Ok(()); - } - - fn parse_options(&self, buf: &[u8]) -> HashMap { - 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), 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((m,l)) => (m.to_lowercase(), l), - 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, file: &mut File) -> Result<(), io::Error> { - let mut block_nr: u16 = 1; - - loop { - let mut filebuf = vec![0; self.options.blksize]; - 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 < self.options.blksize { - /* 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, file: &mut File) -> Result<(), io::Error> { - let mut block_nr: u16 = 1; - - loop { - let mut buf = vec![0; 4 + self.options.blksize + 1]; // +1 for later size check - let mut len = 0; - - for _ in 1..5 { - len = match sock.recv(&mut buf) { - Ok(n) => n, - Err(ref error) if [io::ErrorKind::WouldBlock, io::ErrorKind::TimedOut].contains(&error.kind()) => { - /* re-ack previous and try to recv again */ - self.send_ack(&sock, block_nr - 1)?; - continue; - } - Err(err) => return Err(err), - }; - break; - } - if len < 4 || len > 4 + self.options.blksize { - /* max size: 2 + 2 + blksize */ - 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")), - }; - if u16::from_be_bytes([buf[2], buf[3]]) != block_nr { - /* already received or packets were missed, re-acknowledge */ - self.send_ack(&sock, block_nr - 1)?; - continue; - } - - let databuf = &buf[4..len]; - file.write_all(databuf)?; - - self.send_ack(&sock, block_nr)?; - block_nr = block_nr.wrapping_add(1); - - if len < 4 + self.options.blksize { - break; - } - } - - file.flush()?; - - Ok(()) - } - -} diff --git a/src/tftpc.rs b/src/tftpc.rs index 3283f41..15cadc3 100644 --- a/src/tftpc.rs +++ b/src/tftpc.rs @@ -14,7 +14,7 @@ use std::os::unix::ffi::OsStrExt; /* for converting filename into bytes */ extern crate getopts; use getopts::Options; -mod tftp; +extern crate rtftp; enum Mode { RRQ, @@ -28,14 +28,14 @@ struct Configuration { } struct Tftpc { - tftp: tftp::Tftp, + tftp: rtftp::Tftp, conf: Configuration, } impl Tftpc { pub fn new(conf: Configuration) -> Tftpc { Tftpc { - tftp: tftp::Tftp::new(), + tftp: rtftp::Tftp::new(), conf: conf, } } diff --git a/src/tftpd.rs b/src/tftpd.rs index ce108de..4773664 100644 --- a/src/tftpd.rs +++ b/src/tftpd.rs @@ -18,7 +18,7 @@ use nix::unistd::{Gid,Uid,setresgid,setresuid}; extern crate getopts; use getopts::Options; -mod tftp; +extern crate rtftp; struct Configuration { port: u16, @@ -30,14 +30,14 @@ struct Configuration { } struct Tftpd { - tftp: tftp::Tftp, + tftp: rtftp::Tftp, conf: Configuration, } impl Tftpd { pub fn new(conf: Configuration) -> Tftpd { Tftpd{ - tftp: tftp::Tftp::new(), + tftp: rtftp::Tftp::new(), conf: conf, } } -- cgit v1.2.3