diff --git a/src/main.rs b/src/main.rs index 04515bc..04f3467 100644 --- a/src/main.rs +++ b/src/main.rs @@ -64,7 +64,38 @@ fn create_socket() -> PathBuf { fn handle_client(stream: UnixStream) { let reader = BufReader::new(stream); - for line in reader.lines() { + let start_time = SystemTime::now(); + + const TIMEOUT: Duration = Duration::from_secs(3); + let handshake_msg = String::from("handshake"); + + let mut lines = reader.lines(); + + match lines.next() { + Some(Ok(msg)) if msg == handshake_msg => { + if SystemTime::now().duration_since(start_time).unwrap() >= TIMEOUT { + println!("Client took too long to send first message"); + return; + } + } + Some(Ok(msg)) => { + println!("First line isn't handshake: {}", msg); + return; + } + Some(Err(e)) => { + println!("Failed to get first line: {}", e); + return; + } + None => { + println!("Client terminated before first message"); + // If the stream ends here, abort + return; + } + } + + println!("Got correct handshake"); + + for line in lines { dbg!(line.unwrap()); } @@ -74,7 +105,11 @@ fn handle_client(stream: UnixStream) { fn client(path: String) { let mut stream = UnixStream::connect(path).unwrap(); - stream.write_all(b"Hello World").unwrap(); + writeln!(stream, "handshake").unwrap(); + + thread::sleep(Duration::from_secs(1)); + + writeln!(stream, "Hello world").unwrap(); } fn monitor_child(socket_path: PathBuf) {