From 960ea9072b03cf265a113c2ff978c0d2fb4735a3 Mon Sep 17 00:00:00 2001 From: Reiner Herrmann Date: Sun, 3 Mar 2019 12:57:18 +0100 Subject: Implement Option Extension for client --- src/lib.rs | 37 ++++++++++++++++-------- src/tftpc.rs | 92 +++++++++++++++++++++++++++++++++++++++++++++--------------- src/tftpd.rs | 6 ++-- 3 files changed, 99 insertions(+), 36 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 9d47cff..6a45eeb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,7 +3,7 @@ * License: GPL-3+ */ -use std::net::UdpSocket; +use std::net::{SocketAddr,UdpSocket}; use std::fs::File; use std::collections::HashMap; use std::path::{Path,PathBuf}; @@ -60,6 +60,13 @@ impl Tftp { return Some((val, len)); } + pub fn append_option(&self, buf: &mut Vec, key: &str, val: &str) { + buf.extend(key.bytes()); + buf.push(0x00); + buf.extend(val.bytes()); + buf.push(0x00); + } + fn wait_for_ack(&self, sock: &UdpSocket, expected_block: u16) -> Result { let mut buf = [0; 4]; match sock.recv(&mut buf) { @@ -80,7 +87,7 @@ impl Tftp { Ok(false) } - fn ack_options(&self, sock: &UdpSocket, options: &HashMap, ackwait: bool) -> Result<(), io::Error> { + pub 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 */ @@ -93,10 +100,7 @@ impl Tftp { buf.extend((Opcodes::OACK as u16).to_be_bytes().iter()); for (key, val) in options { - buf.extend(key.bytes()); - buf.push(0x00); - buf.extend(val.bytes()); - buf.push(0x00); + self.append_option(&mut buf, key, val); } for _ in 1..5 { @@ -114,7 +118,7 @@ impl Tftp { 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> { + pub fn init_tftp_options(&mut self, sock: &UdpSocket, options: &mut HashMap) -> Result<(), io::Error> { self.options = default_options(); options.retain(|key, val| { @@ -144,12 +148,10 @@ impl Tftp { 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 { + pub fn parse_options(&self, buf: &[u8]) -> HashMap { let mut options = HashMap::new(); let mut pos = 0; @@ -205,15 +207,26 @@ impl Tftp { Ok(()) } - pub fn send_ack(&self, sock: &UdpSocket, block_nr: u16) -> Result<(), io::Error> { + fn _send_ack(&self, sock: &UdpSocket, cl: Option, block_nr: u16) -> Result<(), io::Error> { let mut buf = Vec::with_capacity(4); buf.extend((Opcodes::ACK as u16).to_be_bytes().iter()); buf.extend(block_nr.to_be_bytes().iter()); - sock.send(&buf)?; + match cl { + Some(remote) => { sock.send_to(&buf, remote)?; } + None => { sock.send(&buf)?; } + } Ok(()) } + pub fn send_ack(&self, sock: &UdpSocket, block_nr: u16) -> Result<(), io::Error> { + self._send_ack(sock, None, block_nr) + } + + pub fn send_ack_to(&self, sock: &UdpSocket, cl: SocketAddr, block_nr: u16) -> Result<(), io::Error> { + self._send_ack(sock, Some(cl), block_nr) + } + pub fn send_file(&self, socket: &UdpSocket, file: &mut File) -> Result<(), io::Error> { let mut block_nr: u16 = 1; diff --git a/src/tftpc.rs b/src/tftpc.rs index 8815bb7..fdc3853 100644 --- a/src/tftpc.rs +++ b/src/tftpc.rs @@ -9,7 +9,6 @@ use std::path::{Path,PathBuf}; use std::env; use std::io; use std::time::Duration; -use std::os::unix::ffi::OsStrExt; /* for converting filename into bytes */ extern crate getopts; use getopts::Options; @@ -40,46 +39,89 @@ impl Tftpc { } } - fn wait_for_response(&self, sock: &UdpSocket, expected_opcode: u16, expected_block: u16) -> Option { + fn wait_for_option_ack(&mut self, sock: &UdpSocket) -> Option { + let mut buf = [0; 512]; + match sock.peek_from(&mut buf) { + Ok(_) => (), + Err(_) => return None, + }; + let opcode = u16::from_be_bytes([buf[0], buf[1]]); + if opcode != rtftp::Opcodes::OACK as u16 { + return None; + } + + let (len, remote) = match sock.recv_from(&mut buf) { + Ok(args) => args, + Err(_) => return None, + }; + + let mut options = self.tftp.parse_options(&buf[2 .. len]); + match self.tftp.init_tftp_options(&sock, &mut options) { + Ok(_) => {}, + Err(_) => return None, + } + + Some(remote) + } + + fn wait_for_response(&self, sock: &UdpSocket, expected_opcode: rtftp::Opcodes, expected_block: u16, expected_remote: Option) -> Option { let mut buf = [0; 4]; let (len, remote) = match sock.peek_from(&mut buf) { Ok(args) => args, Err(_) => return None, }; + if let Some(rem) = expected_remote { + /* verify we got a response from the same client that sent + an optional previous option ack */ + if rem != remote { + return None; + } + } + let opcode = u16::from_be_bytes([buf[0], buf[1]]); let block_nr = u16::from_be_bytes([buf[2], buf[3]]); /* first data packet is expected to be block 1 */ - if len != 4 || opcode != expected_opcode || block_nr != expected_block { + if len != 4 || opcode != expected_opcode as u16 || block_nr != expected_block { return None; } Some(remote) } - fn handle_wrq(&self, sock: &UdpSocket) -> Result<(), io::Error> { + fn append_option_req(&self, buf: &mut Vec) { + self.tftp.append_option(buf, "blksize", &format!("{}", 1428)); + self.tftp.append_option(buf, "timeout", &format!("{}", 3)); + } + + fn handle_wrq(&mut self, sock: &UdpSocket) -> Result<(), io::Error> { let mut file = match File::open(self.conf.filename.as_path()) { Ok(f) => f, Err(err) => return Err(err), }; let filename = match self.conf.filename.file_name() { - Some(f) => f, + Some(f) => match f.to_str() { + Some(s) => s, + None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")), + } None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")), }; let mut buf = Vec::with_capacity(512); - buf.extend([0x00, 0x02].iter()); - buf.extend(filename.as_bytes()); - buf.push(0x00); - buf.extend("octet".bytes()); - buf.push(0x00); + 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); let mut remote = None; for _ in 1 .. 3 { sock.send_to(&buf, self.conf.remote)?; - remote = self.wait_for_response(&sock, rtftp::Opcodes::ACK as u16, 0); - if let Some(_) = remote { + remote = self.wait_for_option_ack(&sock); + if remote.is_none() { + /* for WRQ either OACK or ACK is replied */ + remote = self.wait_for_response(&sock, rtftp::Opcodes::ACK, 0, None); + } + if remote.is_some() { break; } } @@ -100,7 +142,7 @@ impl Tftpc { Ok(()) } - fn handle_rrq(&self, sock: &UdpSocket) -> Result<(), io::Error> { + fn handle_rrq(&mut self, sock: &UdpSocket) -> Result<(), io::Error> { let filename = match self.conf.filename.file_name() { Some(f) => f, None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")), @@ -110,19 +152,26 @@ impl Tftpc { Ok(f) => f, Err(err) => return Err(err), }; + let filename = match self.conf.filename.to_str() { + Some(f) => f, + None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")), + }; let mut buf = Vec::with_capacity(512); - buf.extend([0x00, 0x01].iter()); - buf.extend(self.conf.filename.as_os_str().as_bytes()); - buf.push(0x00); - buf.extend("octet".bytes()); - buf.push(0x00); + 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); let mut remote = None; for _ in 1 .. 3 { sock.send_to(&buf, self.conf.remote)?; - remote = self.wait_for_response(&sock, rtftp::Opcodes::DATA as u16, 1); - if let Some(_) = remote { + let oack_remote = self.wait_for_option_ack(&sock); + if let Some(r) = oack_remote { + /* for RRQ the received OACKs need to be acked */ + self.tftp.send_ack_to(&sock, r, 0)?; + } + remote = self.wait_for_response(&sock, rtftp::Opcodes::DATA, 1, oack_remote); + if remote.is_some() { break; } } @@ -143,10 +192,9 @@ impl Tftpc { Ok(()) } - pub fn start(&self) { + pub fn start(&mut self) { let socket = UdpSocket::bind("[::]:0").expect("binding failed"); socket.set_read_timeout(Some(Duration::from_secs(5))).expect("setting socket timeout failed"); - //socket.connect(self.conf.remote).expect("conneting to remote failed"); let err = match self.conf.mode { Mode::RRQ => self.handle_rrq(&socket), diff --git a/src/tftpd.rs b/src/tftpd.rs index 48e374e..c424f7b 100644 --- a/src/tftpd.rs +++ b/src/tftpd.rs @@ -75,7 +75,8 @@ 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, false)?; + self.tftp.init_tftp_options(&socket, &mut options)?; + self.tftp.ack_options(&socket, &options, false)?; match mode.as_ref() { "octet" => (), @@ -121,7 +122,8 @@ 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, true)?; + self.tftp.init_tftp_options(&socket, &mut options)?; + self.tftp.ack_options(&socket, &options, true)?; match mode.as_ref() { "octet" => (), -- cgit v1.2.3