Require initial handshake message before accepting connections

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) {