Use buffer to ensure only 1 line is read

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());
     }