diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/lib.rs | 166 | ||||
| -rw-r--r-- | src/tftpc.rs | 68 | ||||
| -rw-r--r-- | src/tftpd.rs | 14 |
3 files changed, 202 insertions, 46 deletions
@@ -20,7 +20,7 @@ pub static VERSION: Option<&str> = option_env!("CARGO_PKG_VERSION"); type ProgressCallback = fn(cur: u64, total: u64, state: u64) -> u64; #[repr(u16)] -pub enum Opcodes { +pub enum Opcode { RRQ = 0x01, WRQ = 0x02, DATA = 0x03, @@ -30,6 +30,13 @@ pub enum Opcodes { } #[derive(Clone, Copy)] +#[repr(u8)] +pub enum Mode { + OCTET, + NETASCII, +} + +#[derive(Clone, Copy)] pub struct TftpOptions { blksize: usize, timeout: u8, @@ -39,6 +46,7 @@ pub struct TftpOptions { #[derive(Clone, Copy)] pub struct Tftp { options: TftpOptions, + mode: Mode, progress_cb: Option<ProgressCallback>, } @@ -50,10 +58,46 @@ fn default_options() -> TftpOptions { } } +fn netascii_to_octet(buf: &[u8], previous_cr: bool) -> (Vec<u8>, bool) { + let mut out = Vec::with_capacity(buf.len()); + + let mut prev_cr = previous_cr; + for b in buf { + match *b { + b'\r' => { + if prev_cr { + out.push(b'\r'); + } + prev_cr = true; + continue; + } + b'\0' if prev_cr => out.push(b'\r'), + b'\n' if prev_cr => out.push(b'\n'), + _ => out.push(*b), + } + prev_cr = false; + } + (out, prev_cr) +} + +fn octet_to_netascii(buf: &[u8]) -> Vec<u8> { + let mut out = Vec::with_capacity(2 * buf.len()); + + for b in buf { + match *b { + b'\r' => out.extend(b"\r\0"), + b'\n' => out.extend(b"\r\n"), + _ => out.push(*b), + } + } + out +} + impl Default for Tftp { fn default() -> Tftp { Tftp { options: default_options(), + mode: Mode::OCTET, progress_cb: None, } } @@ -64,6 +108,36 @@ impl Tftp { Default::default() } + pub fn transfersize(&self, file: &mut File) -> Result<u64, io::Error> { + match self.mode { + Mode::OCTET => return Ok(file.metadata().expect("failed to get metadata").len()), + Mode::NETASCII => {}, + } + + let mut total_size = 0; + loop { + let mut buf = [0; 4096]; + let size = match file.read(&mut buf) { + Ok(0) => break, + Ok(s) => s, + Err(ref err) if err.kind() == io::ErrorKind::Interrupted => continue, + Err(err) => return Err(err), + }; + total_size += size as u64; + /* each \r and \n will take two bytes in netascii output */ + total_size += buf[0..size].iter() + .filter(|&x| *x == b'\r' || *x == b'\n') + .count() as u64; + } + + file.seek(io::SeekFrom::Start(0))?; + Ok(total_size) + } + + pub fn set_mode(&mut self, mode: Mode) { + self.mode = mode; + } + fn get_tftp_str(&self, buf: &[u8]) -> Option<String> { let mut iter = buf.iter(); @@ -106,7 +180,7 @@ impl Tftp { } let opcode = u16::from_be_bytes([buf[0], buf[1]]); - if opcode != Opcodes::ERROR as u16 { + if opcode != Opcode::ERROR as u16 { return std::io::Error::new(kind, error); } @@ -147,9 +221,9 @@ impl Tftp { let opcode = u16::from_be_bytes([buf[0], buf[1]]); let block_nr = u16::from_be_bytes([buf[2], buf[3]]); - if opcode == Opcodes::ACK as u16 && block_nr == expected_block { + if opcode == Opcode::ACK as u16 && block_nr == expected_block { return Ok(true); - } else if opcode == Opcodes::ERROR as u16 { + } else if opcode == Opcode::ERROR as u16 { return Err(self.parse_error(&buf[4..])); } @@ -166,7 +240,7 @@ impl Tftp { } let mut buf = Vec::with_capacity(512); - buf.extend((Opcodes::OACK as u16).to_be_bytes().iter()); + buf.extend((Opcode::OACK as u16).to_be_bytes().iter()); for (key, val) in options { self.append_option(&mut buf, key, val); @@ -271,7 +345,7 @@ impl Tftp { pub fn send_error(&self, socket: &UdpSocket, code: u16, msg: &str) -> Result<(), io::Error> { let mut buf = Vec::with_capacity(512); - buf.extend((Opcodes::ERROR as u16).to_be_bytes().iter()); + buf.extend((Opcode::ERROR as u16).to_be_bytes().iter()); buf.extend(code.to_be_bytes().iter()); buf.extend(msg.as_bytes()); @@ -281,7 +355,7 @@ impl Tftp { fn _send_ack(&self, sock: &UdpSocket, cl: Option<SocketAddr>, 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((Opcode::ACK as u16).to_be_bytes().iter()); buf.extend(block_nr.to_be_bytes().iter()); match cl { @@ -305,9 +379,12 @@ impl Tftp { let mut prog_update = 0; let tsize = self.transfer_size(file); + /* holds bytes from netascii conversion that did not fit in tx buffer */ + let mut overflow = Vec::with_capacity(2 * self.options.blksize); + loop { - let mut filebuf = vec![0; self.options.blksize]; - let len = match file.read(&mut filebuf) { + let mut filebuf = vec![0; self.options.blksize - overflow.len()]; + let mut len = match file.read(&mut filebuf) { Ok(n) => n, Err(ref error) if error.kind() == io::ErrorKind::Interrupted => continue, /* retry */ Err(err) => { @@ -316,10 +393,26 @@ impl Tftp { } }; + /* take care of netascii conversion */ + let mut databuf = filebuf[0..len].to_vec(); + match self.mode { + Mode::OCTET => {}, + Mode::NETASCII => { + overflow.extend(octet_to_netascii(&databuf)); + databuf = overflow.clone(); + if overflow.len() > self.options.blksize { + overflow = databuf.split_off(self.options.blksize); + } else { + overflow.clear(); + } + len = databuf.len(); + } + } + let mut sendbuf = Vec::with_capacity(4 + len); - sendbuf.extend((Opcodes::DATA as u16).to_be_bytes().iter()); + sendbuf.extend((Opcode::DATA as u16).to_be_bytes().iter()); sendbuf.extend(block_nr.to_be_bytes().iter()); - sendbuf.extend(filebuf[0..len].iter()); + sendbuf.extend(databuf.iter()); let mut acked = false; for _ in 1..5 { @@ -359,6 +452,7 @@ impl Tftp { let mut block_nr: u16 = 1; let mut prog_update = 0; let mut transferred = 0; + let mut netascii_state = false; let tsize = self.transfer_size(file); loop { @@ -383,8 +477,8 @@ impl Tftp { } match u16::from_be_bytes([buf[0], buf[1]]) { // opcode - opc if opc == Opcodes::DATA as u16 => (), - opc if opc == Opcodes::ERROR as u16 => return Err(self.parse_error(&buf[..len])), + opc if opc == Opcode::DATA as u16 => (), + opc if opc == Opcode::ERROR as u16 => return Err(self.parse_error(&buf[..len])), _ => return Err(io::Error::new(io::ErrorKind::Other, "unexpected opcode")), }; if u16::from_be_bytes([buf[2], buf[3]]) != block_nr { @@ -393,8 +487,16 @@ impl Tftp { continue; } - let databuf = &buf[4..len]; - file.write_all(databuf)?; + let mut databuf = buf[4..len].to_vec(); + match self.mode { + Mode::OCTET => {}, + Mode::NETASCII => { + let (converted, state) = netascii_to_octet(&databuf, netascii_state); + databuf = converted; + netascii_state = state; + } + } + file.write_all(&databuf)?; transferred += (len - 4) as u64; if let Some(cb) = self.progress_cb { @@ -409,6 +511,11 @@ impl Tftp { } } + if netascii_state { + /* the file ended with an incomplete \r encoding */ + file.write(&[b'\r'])?; + } + file.flush()?; Ok(()) @@ -492,4 +599,33 @@ mod tests { tftp.append_option(&mut buf, "key", "value"); assert_eq!(buf, "key\x00value\x00".as_bytes()); } + + #[test] + fn test_netascii_to_octet() { + assert_eq!(netascii_to_octet(b"\r\nfoo\r\0bar", false), (b"\nfoo\rbar".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\r\0", false), (b"\r".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\r\n", false), (b"\n".to_vec(), false)); + assert_eq!(netascii_to_octet(b"", false), (b"".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\n\0\n\0", false), (b"\n\0\n\0".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\r\r\n", false), (b"\r\n".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\r\n\r\n", false), (b"\n\n".to_vec(), false)); + assert_eq!(netascii_to_octet(b"test\r\0", false), (b"test\r".to_vec(), false)); + assert_eq!(netascii_to_octet(b"test\r", false), (b"test".to_vec(), true)); + assert_eq!(netascii_to_octet(b"\r", false), (b"".to_vec(), true)); + assert_eq!(netascii_to_octet(b"\0test", true), (b"\rtest".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\ntest", true), (b"\ntest".to_vec(), false)); + assert_eq!(netascii_to_octet(b"\n\r", true), (b"\n".to_vec(), true)); + assert_eq!(netascii_to_octet(b"", true), (b"".to_vec(), true)); + assert_eq!(netascii_to_octet(b"\r", true), (b"\r".to_vec(), true)); + } + + #[test] + fn test_octet_to_netascii() { + assert_eq!(octet_to_netascii(b"foobar"), b"foobar"); + assert_eq!(octet_to_netascii(b"foo\rbar\n"), b"foo\r\0bar\r\n"); + assert_eq!(octet_to_netascii(b"\r\n"), b"\r\0\r\n"); + assert_eq!(octet_to_netascii(b"\r\r\n\n"), b"\r\0\r\0\r\n\r\n"); + assert_eq!(octet_to_netascii(b"\r\0\r\n"), b"\r\0\0\r\0\r\n"); + assert_eq!(octet_to_netascii(b""), b""); + } } diff --git a/src/tftpc.rs b/src/tftpc.rs index adac705..24073d2 100644 --- a/src/tftpc.rs +++ b/src/tftpc.rs @@ -15,13 +15,14 @@ use getopts::Options; extern crate rtftp; -enum Mode { +enum Operation { RRQ, WRQ, } struct Configuration { - mode: Mode, + operation: Operation, + mode: rtftp::Mode, filename: PathBuf, remote: SocketAddr, blksize: usize, @@ -46,7 +47,7 @@ fn update_progress(current: u64, total: u64, last: u64) -> u64 { let percent = 100 * current / total; print!("\r {}% ", percent); io::stdout().flush().expect("flushing stdout failed"); - if current == total { + if current >= total { print!("\r"); } current @@ -54,8 +55,10 @@ fn update_progress(current: u64, total: u64, last: u64) -> u64 { impl Tftpc { pub fn new(conf: Configuration) -> Tftpc { + let mut tftp = rtftp::Tftp::new(); + tftp.set_mode(conf.mode); Tftpc { - tftp: rtftp::Tftp::new(), + tftp, conf, } } @@ -67,7 +70,7 @@ impl Tftpc { Err(_) => return None, }; let opcode = u16::from_be_bytes([buf[0], buf[1]]); - if opcode != rtftp::Opcodes::OACK as u16 { + if opcode != rtftp::Opcode::OACK as u16 { return None; } @@ -85,7 +88,7 @@ impl Tftpc { Some(remote) } - fn wait_for_response(&self, sock: &UdpSocket, expected_opcode: rtftp::Opcodes, expected_block: u16, expected_remote: Option<SocketAddr>) -> Result<Option<SocketAddr>, std::io::Error> { + fn wait_for_response(&self, sock: &UdpSocket, expected_opcode: rtftp::Opcode, expected_block: u16, expected_remote: Option<SocketAddr>) -> Result<Option<SocketAddr>, std::io::Error> { let mut buf = [0; 4]; let (len, remote) = match sock.peek_from(&mut buf) { Ok(args) => args, @@ -103,7 +106,7 @@ impl Tftpc { let opcode = u16::from_be_bytes([buf[0], buf[1]]); let block_nr = u16::from_be_bytes([buf[2], buf[3]]); - if opcode == rtftp::Opcodes::ERROR as u16 { + if opcode == rtftp::Opcode::ERROR as u16 { let mut buf = [0; 512]; let len = sock.recv(&mut buf)?; return Err(self.tftp.parse_error(&buf[..len])); @@ -123,6 +126,19 @@ impl Tftpc { self.tftp.append_option(buf, "tsize", &format!("{}", fsize)); } + fn init_req(&self, opcode: rtftp::Opcode, filename: &str, size: u64) -> Vec<u8> { + let mut buf = Vec::with_capacity(512); + buf.extend((opcode as u16).to_be_bytes().iter()); + let mode_str = match self.conf.mode { + rtftp::Mode::OCTET => "octet", + rtftp::Mode::NETASCII => "netascii", + }; + self.tftp.append_option(&mut buf, filename, mode_str); + self.append_option_req(&mut buf, size); + + buf + } + fn handle_wrq(&mut self, sock: &UdpSocket) -> Result<String, io::Error> { let mut file = match File::open(self.conf.filename.as_path()) { Ok(f) => f, @@ -144,10 +160,8 @@ impl Tftpc { return Err(err_invalidpath); } - 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, metadata.len()); + let tsize = self.tftp.transfersize(&mut file)?; + let buf = self.init_req(rtftp::Opcode::WRQ, filename, tsize); let mut remote = None; for _ in 1..3 { @@ -155,7 +169,7 @@ impl Tftpc { 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)?; + remote = self.wait_for_response(&sock, rtftp::Opcode::ACK, 0, None)?; } if remote.is_some() { break; @@ -192,10 +206,7 @@ impl Tftpc { None => return Err(io::Error::new(io::ErrorKind::InvalidInput, "Invalid path/filename")), }; - 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, 0); + let buf = self.init_req(rtftp::Opcode::RRQ, filename, 0); let mut remote = None; for _ in 1..3 { @@ -205,7 +216,7 @@ impl Tftpc { /* 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)?; + remote = self.wait_for_response(&sock, rtftp::Opcode::DATA, 1, oack_remote)?; if remote.is_some() { break; } @@ -231,9 +242,9 @@ impl Tftpc { let socket = UdpSocket::bind("[::]:0").expect("binding failed"); socket.set_read_timeout(Some(Duration::from_secs(5))).expect("setting socket timeout failed"); - let err = match self.conf.mode { - Mode::RRQ => self.handle_rrq(&socket), - Mode::WRQ => self.handle_wrq(&socket), + let err = match self.conf.operation { + Operation::RRQ => self.handle_rrq(&socket), + Operation::WRQ => self.handle_wrq(&socket), }; match err { Ok(msg) => println!("{}", msg), @@ -255,7 +266,8 @@ fn usage(opts: Options, program: String, error: Option<String>) { fn parse_commandline(args: &[String]) -> Result<Configuration, &str> { let program = args[0].clone(); - let mut mode = None; + let mut operation = None; + let mut mode = rtftp::Mode::OCTET; let mut filename = None; let mut blksize = 1428; @@ -264,6 +276,7 @@ fn parse_commandline(args: &[String]) -> Result<Configuration, &str> { opts.optopt("g", "get", "download file from remote server", "FILE"); opts.optopt("p", "put", "upload file to remote server", "FILE"); opts.optopt("b", "blksize", format!("negotiate a different block size (default: {})", blksize).as_ref(), "SIZE"); + opts.optflag("n", "netascii","use netascii mode (instead of octet)"); let matches = match opts.parse(&args[1..]) { Ok(m) => m, Err(err) => { @@ -277,19 +290,23 @@ fn parse_commandline(args: &[String]) -> Result<Configuration, &str> { } if let Some(f) = matches.opt_str("g") { - mode = Some(Mode::RRQ); + operation = Some(Operation::RRQ); filename = Some(Path::new(&f).to_path_buf()); } if let Some(f) = matches.opt_str("p") { - mode = Some(Mode::WRQ); + operation = Some(Operation::WRQ); filename = Some(Path::new(&f).to_path_buf()); } - if mode.is_none() || (matches.opt_present("g") && matches.opt_present("p")) { + if operation.is_none() || (matches.opt_present("g") && matches.opt_present("p")) { usage(opts, program, Some("Exactly one of g (get) and p (put) required".to_string())); return Err("get put"); } + if matches.opt_present("n") { + mode = rtftp::Mode::NETASCII; + } + let remote_in = matches.free[0].as_str(); let remote = match remote_in.to_socket_addrs() { Ok(mut i) => i.next(), @@ -311,7 +328,8 @@ fn parse_commandline(args: &[String]) -> Result<Configuration, &str> { }; Ok(Configuration { - mode: mode.unwrap(), + operation: operation.unwrap(), + mode, filename: filename.unwrap(), remote: remote.unwrap(), blksize, diff --git a/src/tftpd.rs b/src/tftpd.rs index 024b105..ecb9b37 100644 --- a/src/tftpd.rs +++ b/src/tftpd.rs @@ -92,7 +92,8 @@ impl Tftpd { self.tftp.init_tftp_options(&socket, &mut options)?; match mode.as_ref() { - "octet" => (), + "octet" => self.tftp.set_mode(rtftp::Mode::OCTET), + "netascii" => self.tftp.set_mode(rtftp::Mode::NETASCII), _ => { self.tftp.send_error(&socket, 0, "Unsupported mode")?; return Err(io::Error::new(io::ErrorKind::Other, "unsupported mode")); @@ -138,7 +139,8 @@ impl Tftpd { self.tftp.init_tftp_options(&socket, &mut options)?; match mode.as_ref() { - "octet" => (), + "octet" => self.tftp.set_mode(rtftp::Mode::OCTET), + "netascii" => self.tftp.set_mode(rtftp::Mode::NETASCII), _ => { self.tftp.send_error(&socket, 0, "Unsupported mode")?; return Err(io::Error::new(io::ErrorKind::Other, "unsupported mode")); @@ -173,7 +175,7 @@ impl Tftpd { } if let Some(opt) = options.get_mut("tsize") { - *opt = file.metadata()?.len().to_string(); + *opt = self.tftp.transfersize(&mut file)?.to_string(); } self.tftp.ack_options(&socket, &options, true)?; match self.tftp.send_file(&socket, &mut file) { @@ -191,7 +193,7 @@ impl Tftpd { socket.connect(cl)?; match u16::from_be_bytes([buf[0], buf[1]]) { // opcode - o if o == rtftp::Opcodes::RRQ as u16 => { + o if o == rtftp::Opcode::RRQ as u16 => { if self.conf.wo { self.tftp.send_error(&socket, 4, "reading not allowed")?; Err(io::Error::new(io::ErrorKind::Other, "unallowed mode")) @@ -199,7 +201,7 @@ impl Tftpd { self.handle_rrq(&socket, &cl, &buf[2..]) } } - o if o == rtftp::Opcodes::WRQ as u16 => { + o if o == rtftp::Opcode::WRQ as u16 => { if self.conf.ro { self.tftp.send_error(&socket, 4, "writing not allowed")?; Err(io::Error::new(io::ErrorKind::Other, "unallowed mode")) @@ -207,7 +209,7 @@ impl Tftpd { self.handle_wrq(&socket, &cl, &buf[2..]) } } - o if o == rtftp::Opcodes::ERROR as u16 => Ok(format!("Received ERROR from {}", cl)), + o if o == rtftp::Opcode::ERROR as u16 => Ok(format!("Received ERROR from {}", cl)), _ => { self.tftp.send_error(&socket, 4, "Unexpected opcode")?; Err(io::Error::new(io::ErrorKind::Other, "unexpected opcode")) |