diff --git a/src/main.rs b/src/main.rs index eb7fdeb..25d9eb7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,7 @@ use rand::{thread_rng, RngCore}; use std::env; use std::fs::remove_file; -use std::io::{BufRead, BufReader, Read, Write}; +use std::io::{BufRead, BufReader, Write}; use std::os::unix::net::{UnixListener, UnixStream}; use std::path::PathBuf; use std::process::Command; @@ -61,46 +61,40 @@ fn create_socket() -> PathBuf { temp_dir } -fn handle_client(mut stream: UnixStream) { - let start_time = SystemTime::now(); - +fn handle_client(stream: UnixStream) { const TIMEOUT: Duration = Duration::from_secs(3); - stream - .set_read_timeout(Some(Duration::from_millis(1))) - .unwrap(); + let reader = BufReader::new(stream.try_clone().unwrap()); - loop { - let mut response = String::new(); + stream.set_read_timeout(Some(TIMEOUT)).unwrap(); - // Change to use buffered, read line until timeout, and catch timeout as below (hopefully can change timeout after initial read) + let mut lines = reader.lines(); - match stream.read_to_string(&mut response) { - Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {} - Err(e) => { - println!("Error reading {}", e); - } - _ => {} - } - - if response.trim_end() == "handshake" { - println!("Got correct handshake"); - break; - } - - if SystemTime::now().duration_since(start_time).unwrap() >= TIMEOUT { - println!("Timeout expired, killing connection"); + match lines.next() { + Some(Err(e)) if e.kind() == std::io::ErrorKind::WouldBlock => { + println!("Timeout"); + return; + } + Some(Err(e)) => { + println!("Error reading {}", e); + return; + } + Some(Ok(msg)) => { + if msg != "handshake" { + println!("Got incorrect handshake: {:?}", msg); + return; + } + println!("Got correct handshake"); + } + None => { return; } - - // thread::sleep(Duration::from_secs(1)); } + // Restore blocking socket for increased efficiency stream.set_read_timeout(None).unwrap(); - let reader = BufReader::new(stream); - - for line in reader.lines() { + for line in lines { dbg!(line.unwrap()); }