From 85ff31c0b630f628d8aa8c3f5cd528337046e154 Mon Sep 17 00:00:00 2001 From: Hugo Osvaldo Barrera Date: Tue, 7 Jun 2022 18:40:09 +0200 Subject: [PATCH 1/2] Tidy up imports Mostly for `io::`. --- src/main.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/main.rs b/src/main.rs index 37b5fda..d77445b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,12 @@ extern crate byteorder; -extern crate nix; #[cfg(windows)] extern crate named_pipe; +extern crate nix; -use std::convert::TryInto; -use std::io::{self, Read, Write}; -use std::thread; use byteorder::{ByteOrder, NativeEndian, WriteBytesExt}; +use std::convert::TryInto; +use std::io::{stdin, stdout, Error, ErrorKind, Read, Result, Write}; +use std::thread; mod proxy_socket; @@ -19,7 +19,7 @@ fn valid_length(length: usize) -> bool { } fn read_header() -> usize { - let stdin = io::stdin(); + let stdin = stdin(); let mut buf = vec![0; 4]; let mut handle = stdin.lock(); @@ -29,7 +29,7 @@ fn read_header() -> usize { fn read_body(length: usize, socket: &mut ProxySocket) { let mut buffer = vec![0; length]; - let stdin = io::stdin(); + let stdin = stdin(); let mut handle = stdin.lock(); if handle.read_exact(&mut buffer).is_ok() && valid_length(length) { @@ -47,7 +47,7 @@ fn read_response(socket: &mut ProxySocket) { } fn write_response(buf: &[u8]) { - let stdout = io::stdout(); + let stdout = stdout(); let mut out = stdout.lock(); out.write_u32::(buf.len() as u32).unwrap(); From 22c65c01786779ee9cfdd6b9aee44d9c730fa947 Mon Sep 17 00:00:00 2001 From: Hugo Osvaldo Barrera Date: Tue, 7 Jun 2022 18:49:27 +0200 Subject: [PATCH 2/2] Clean up `unwrap`s everywhere except `main()` --- src/main.rs | 56 +++++++++++++++++++++++++++++---------------- src/proxy_socket.rs | 7 +++--- 2 files changed, 40 insertions(+), 23 deletions(-) diff --git a/src/main.rs b/src/main.rs index d77445b..b175d6b 100644 --- a/src/main.rs +++ b/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 { 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(length: usize, socket: &mut ProxySocket) { +// 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(length: usize, socket: &mut ProxySocket) -> 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(socket: &mut ProxySocket) { +// Read a response (from KP's socket) and echo it back to the browser. +fn read_response(socket: &mut ProxySocket) -> 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::(buf.len() as u32).unwrap(); - out.write_all(buf).unwrap(); - out.flush().unwrap(); + out.write_u32::(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(); diff --git a/src/proxy_socket.rs b/src/proxy_socket.rs index fa6fd33..4e6668c 100644 --- a/src/proxy_socket.rs +++ b/src/proxy_socket.rs @@ -49,7 +49,7 @@ fn get_socket_dirs() -> Vec { 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> { .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 = Some(Duration::from_secs(1)); s.set_read_timeout(timeout)?; Ok(ProxySocket { inner: s })