Clean up unwraps everywhere except main()

src/main.rs

@@ -12,58 +12,74 @@ mod proxy_socket;
 
 use proxy_socket::ProxySocket;
 
-const BUFFER_SIZE: usize = 1024 ^ 2;	 // 1024 ^ 2 is the maximum
+const BUFFER_SIZE: usize = 1024 ^ 2; // 1024 ^ 2 is the maximum
 
 fn valid_length(length: usize) -> bool {
     length > 0 && length <= BUFFER_SIZE
 }
 
-fn read_header() -> usize {
+// Read a header (message size) from stdin (e.g.: from the browser).
+fn read_header() -> Result<usize> {
     let stdin = stdin();
     let mut buf = vec![0; 4];
     let mut handle = stdin.lock();
 
-    handle.read_exact(&mut buf).unwrap();
-    NativeEndian::read_u32(&buf).try_into().unwrap()
+    handle.read_exact(&mut buf)?;
+
+    NativeEndian::read_u32(&buf)
+        .try_into()
+        .map_err(|err| Error::new(ErrorKind::InvalidData, err))
 }
 
-fn read_body<T: Read + Write>(length: usize, socket: &mut ProxySocket<T>) {
+// Handle a whole request/response cycle
+//
+// Read a message body from stdin (e.g.: from the browser), and echo it back to the browser's
+// socket. Then await a response from the socket and relay that back to the browser.
+fn read_body<T: Read + Write>(length: usize, socket: &mut ProxySocket<T>) -> Result<()> {
     let mut buffer = vec![0; length];
     let stdin = stdin();
     let mut handle = stdin.lock();
 
-    if handle.read_exact(&mut buffer).is_ok() && valid_length(length) {
-        socket.write_all(&buffer).unwrap();
-        socket.flush().unwrap();
-        read_response(socket);
+    handle.read_exact(&mut buffer)?;
+
+    if valid_length(length) {
+        socket.write_all(&buffer)?;
+        socket.flush()?;
+        read_response(socket)?;
     }
+
+    Ok(())
 }
 
-fn read_response<T: Read>(socket: &mut ProxySocket<T>) {
+// Read a response (from KP's socket) and echo it back to the browser.
+fn read_response<T: Read>(socket: &mut ProxySocket<T>) -> Result<()>{
     let mut buf = vec![0; BUFFER_SIZE];
     if let Ok(len) = socket.read(&mut buf) {
-        write_response(&buf[0..len]);
+        write_response(&buf[0..len])?;
     }
+
+    Ok(())
 }
 
-fn write_response(buf: &[u8]) {
+// Write a response to stdout (e.g.: to the browser).
+fn write_response(buf: &[u8]) -> Result<()> {
     let stdout = stdout();
     let mut out = stdout.lock();
 
-    out.write_u32::<NativeEndian>(buf.len() as u32).unwrap();
-    out.write_all(buf).unwrap();
-    out.flush().unwrap();
+    out.write_u32::<NativeEndian>(buf.len() as u32)?;
+    out.write_all(buf)?;
+    out.flush()?;
+
+    Ok(())
 }
 
 fn main() {
     let mut socket = proxy_socket::connect(BUFFER_SIZE).unwrap();
 
     // Start thread for user input reading
-    let ui = thread::spawn(move || {
-        loop {
-            let length = read_header();
-            read_body(length, &mut socket);
-        }
+    let ui = thread::spawn(move || loop {
+        let length = read_header().unwrap();
+        read_body(length, &mut socket).unwrap();
     });
 
     let _ui_res = ui.join().unwrap();

src/proxy_socket.rs

@@ -49,7 +49,7 @@ fn get_socket_dirs() -> Vec<PathBuf> {
     let mut dirs = Vec::new();
 
     if !cfg!(target_os = "macos") {
-        if let Ok(dir) = env::var("XDG_RUNTIME_DIR")  {
+        if let Ok(dir) = env::var("XDG_RUNTIME_DIR") {
             let xdg_runtime_dir: PathBuf = dir.into();
 
             // Sandbox-friendly path.
@@ -79,8 +79,9 @@ pub fn connect(buffer_size: usize) -> io::Result<ProxySocket<UnixStream>> {
         .find_map(|dir| UnixStream::connect(dir.join(socket_name)).ok())
         .ok_or_else(|| io::Error::from(io::ErrorKind::NotFound))?;
 
-    socket::setsockopt(s.as_raw_fd(), SndBuf, &buffer_size).expect("setsockopt for SndBuf failed");
-    socket::setsockopt(s.as_raw_fd(), RcvBuf, &buffer_size).expect("setsockopt for RcvBuf failed");
+    socket::setsockopt(s.as_raw_fd(), SndBuf, &buffer_size)?;
+    socket::setsockopt(s.as_raw_fd(), RcvBuf, &buffer_size)?;
+
     let timeout: Option<Duration> = Some(Duration::from_secs(1));
     s.set_read_timeout(timeout)?;
     Ok(ProxySocket { inner: s })