From ecd3f18f71cc1bf96c7500a3fa40c5fafdff1947 Mon Sep 17 00:00:00 2001 From: bnewbold Date: Thu, 2 Jun 2016 22:19:26 -0400 Subject: massively improve io error handling Slay many unwrap()'s --- src/client.rs | 25 ++++++++++++++--------- src/common.rs | 64 ++++++++++++++++++++++++++++++++++------------------------- src/server.rs | 13 ++++++++---- 3 files changed, 62 insertions(+), 40 deletions(-) diff --git a/src/client.rs b/src/client.rs index f26c35d..79a95f7 100644 --- a/src/client.rs +++ b/src/client.rs @@ -6,7 +6,9 @@ use std::string::String; use std::str::{self, FromStr}; use std::env; use std::net; +use std::io; use std::process; +use std::error::Error; use std::process::exit; use std::process::Command; use getopts::Options; @@ -70,23 +72,28 @@ pub fn run_client(host: &str, local_file: &str, remote_file: &str, remote_is_dir }; let mut stream: UdtStream = UdtStream::new(socket); + let io_result: io::Result<()>; if !no_crypto { let mut stream = SecretStream::new(stream); stream.key = string2key(remote_secret).unwrap(); stream.read_nonce = string2nonce(remote_write_nonce).unwrap(); stream.write_nonce = string2nonce(remote_read_nonce).unwrap(); if is_recv { - common::sink_files(&mut stream, local_file, remote_is_dir) + io_result = common::sink_files(&mut stream, local_file, remote_is_dir); } else { - common::source_files(&mut stream, local_file, remote_is_dir) + io_result = common::source_files(&mut stream, local_file, remote_is_dir); } } else { if is_recv { - common::sink_files(&mut stream, local_file, remote_is_dir) + io_result = common::sink_files(&mut stream, local_file, remote_is_dir); } else { - common::source_files(&mut stream, local_file, remote_is_dir) + io_result = common::source_files(&mut stream, local_file, remote_is_dir); } } + match io_result { + Ok(_) => Ok(()), + Err(e) => Err(e.description().to_string()), + } } fn usage_client(opts: Options) { @@ -154,7 +161,7 @@ pub fn main_client() { let mut stream: UdtStream = UdtStream::new(socket); println!("opened socket"); - let mut ret: Result<(), String>; + let mut ret: io::Result<()>; if !no_crypto { let mut stream = SecretStream::new(stream); stream.key = string2key(&matches.opt_str("key").unwrap()).unwrap(); @@ -165,7 +172,7 @@ pub fn main_client() { } else if matches.opt_present("t") { ret = common::sink_files(&mut stream, &matches.opt_str("t").unwrap(), dir_mode); } else { - ret = Err("Didn't Run".to_string()); + unreachable!(); } } else { if matches.opt_present("f") { @@ -173,14 +180,14 @@ pub fn main_client() { } else if matches.opt_present("t") { ret = common::sink_files(&mut stream, &matches.opt_str("t").unwrap(), dir_mode); } else { - ret = Err("Didn't Run".to_string()); + unreachable!(); } } match ret { Ok(_) => { exit(0); }, - Err(msg) => { - writeln!(&mut ::std::io::stderr(), "{}", msg).unwrap(); + Err(err) => { + writeln!(&mut ::std::io::stderr(), "{}", err.description()).unwrap(); exit(-1); } } diff --git a/src/common.rs b/src/common.rs index 9aa477f..facd527 100644 --- a/src/common.rs +++ b/src/common.rs @@ -9,50 +9,54 @@ use std::io::{Read, Write, BufRead, BufReader}; use std::process::exit; use std::net; -pub fn source_files(stream: &mut S, file_path: &str, recursive: bool) -> Result<(), String> { +fn fake_io_error(msg: &str) -> io::Result<()> { + Err(io::Error::new(io::ErrorKind::Other, msg)) +} + +pub fn source_files(stream: &mut S, file_path: &str, recursive: bool) -> io::Result<()> { println!("SOURCE FILE: {}", file_path); if recursive { unimplemented!(); } - let mut f = File::open(file_path).unwrap(); - let metadata = f.metadata().unwrap(); + let mut f = try!(File::open(file_path)); + let metadata = try!(f.metadata()); assert!(metadata.is_file()); let fmode: u32 = metadata.permissions().mode(); let flen: usize = metadata.len() as usize; // Format as 4 digits octal, left-padding with zero let line = format!("C{:0>4o} {} {}\n", fmode, flen, file_path); - stream.write_all(line.as_bytes()).unwrap(); + try!(stream.write_all(line.as_bytes())); println!("{}", line); let mut byte_buf = [0; 1]; - stream.read_exact(&mut byte_buf).unwrap(); + try!(stream.read_exact(&mut byte_buf)); let reply = byte_buf[0]; match reply { 0 => {}, // Success, pass 1 | 2 => { // Warning unimplemented!(); }, - _ => { return Err("Unexpected status char!".to_string()); }, + _ => { return fake_io_error("Unexpected status char!"); }, }; let mut buf = [0; 4096]; let mut sent: usize = 0; while sent < flen { - let rlen = f.read(&mut buf).unwrap(); + let rlen = try!(f.read(&mut buf)); assert!(rlen > 0); let mut wbuf = &mut buf[..rlen]; - stream.write_all(&wbuf).unwrap(); + try!(stream.write_all(&wbuf)); sent += rlen; //println!("sent: {}", sent); } // f.close(); XXX: - stream.read_exact(&mut byte_buf).unwrap(); + try!(stream.read_exact(&mut byte_buf)); let reply = byte_buf[0]; match reply { 0 => {}, // Success, pass 1 | 2 => { // Warning unimplemented!(); }, - _ => { return Err("Unexpected status char!".to_string()) }, + _ => { return fake_io_error("Unexpected status char!"); }, }; Ok(()) } @@ -62,7 +66,7 @@ fn raw_read_line(stream: &mut S) -> io::Result { let mut s = String::new(); let mut byte_buf = [0]; loop { - stream.read_exact(&mut byte_buf).unwrap(); + try!(stream.read_exact(&mut byte_buf)); if byte_buf[0] == '\n' as u8 { return Ok(s); } @@ -73,48 +77,54 @@ fn raw_read_line(stream: &mut S) -> io::Result { // TODO: it would be nice to be able to do BufReader/BufWriter on stream. This would require // implementations of Read and Write on immutable references to stream (a la TcpStream, File, et // al) -pub fn sink_files(stream: &mut S, file_path: &str, recursive: bool) -> Result<(), String> { - println!("SINK FILE: {}", file_path); +pub fn sink_files(stream: &mut S, file_path: &str, recursive: bool) -> io::Result<()> { + info!("SINK FILE: {}", file_path); if recursive { unimplemented!(); } - let mut f = File::create(file_path).unwrap(); + let mut f = try!(File::create(file_path)); let mut byte_buf = [0; 1]; let mut buf = [0; 4096]; - stream.read_exact(&mut byte_buf).unwrap(); + try!(stream.read_exact(&mut byte_buf)); let msg_type = byte_buf[0]; match msg_type as char { 'C' => { - println!("Going to create!"); + info!("Going to create!"); }, 'D' => { unimplemented!(); }, 'E' => { unimplemented!(); }, 'T' => { unimplemented!(); }, - _ => { return Err(format!("Unexpected message type: {}", msg_type)); }, + _ => { return fake_io_error(&format!("Unexpected message type: {}", msg_type)); }, }; - let line = raw_read_line(stream).unwrap(); + let line = try!(raw_read_line(stream)); println!("got msg: {}", line); - stream.write(&[0]).unwrap(); + try!(stream.write(&[0])); let line: Vec<&str> = line.split_whitespace().collect(); assert!(line.len() == 3); - let fmode: u32 = u32::from_str_radix(line[0], 8).unwrap(); - let flen: usize = line[1].parse::().unwrap(); + let fmode: u32 = match u32::from_str_radix(line[0], 8) { + Ok(x) => x, + Err(_) => { return fake_io_error("unparsable file mode in ucp 'C' message"); }, + }; + let flen: usize = match line[1].parse::() { + Ok(x) => x, + Err(_) => { return fake_io_error("unparsable file length in ucp 'C' message"); }, + }; let fpath = Path::new(line[2]); // TODO: I've disabled set_len; is this best practice? scp doesn't do it. - //f.set_len(flen as u64).unwrap(); - fs::set_permissions(file_path, PermissionsExt::from_mode(fmode)).unwrap(); + //try!(f.set_len(flen as u64)); + try!(fs::set_permissions(file_path, PermissionsExt::from_mode(fmode))); let mut received: usize = 0; while received < flen { - let rlen = stream.read(&mut buf).unwrap(); + let rlen = try!(stream.read(&mut buf)); //println!("recieved: {}", rlen); assert!(rlen > 0); let mut wbuf = &mut buf[..rlen]; - f.write_all(&wbuf).unwrap(); + try!(f.write_all(&wbuf)); received += rlen; } - f.sync_all().unwrap(); + try!(f.sync_all()); // f.close(); XXX: closes automatically? - stream.write(&[0]).unwrap(); + try!(stream.write(&[0])); Ok(()) } diff --git a/src/server.rs b/src/server.rs index cb57499..cedcf8c 100644 --- a/src/server.rs +++ b/src/server.rs @@ -123,23 +123,28 @@ fn run_server(path: &str, is_recv: bool, recursive: bool, daemonize: bool, no_cr info!("Got connection from {}", socket.getpeername().unwrap()); let mut stream: UdtStream = UdtStream::new(socket); + let io_result: io::Result<()>; if !no_crypto { let mut stream = SecretStream::new(stream); stream.key = secret_key; stream.read_nonce = read_nonce; stream.write_nonce = write_nonce; if is_recv { - common::sink_files(&mut stream, path, recursive) + io_result = common::sink_files(&mut stream, path, recursive) } else { - common::source_files(&mut stream, path, recursive) + io_result = common::source_files(&mut stream, path, recursive) } } else { if is_recv { - common::sink_files(&mut stream, path, recursive) + io_result = common::sink_files(&mut stream, path, recursive) } else { - common::source_files(&mut stream, path, recursive) + io_result = common::source_files(&mut stream, path, recursive) } } + match io_result { + Ok(_) => Ok(()), + Err(e) => Err(e.description().to_string()), + } } fn usage_server(opts: Options) { -- cgit v1.2.3