From 59e0c934e22a9564e84378a2a468fee054d5be18 Mon Sep 17 00:00:00 2001
From: Jake Howard <git@theorangeone.net>
Date: Sat, 19 Aug 2023 21:37:39 +0100
Subject: [PATCH] Require initial handshake message before accepting
 connections

---
 src/main.rs | 39 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 37 insertions(+), 2 deletions(-)

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