From bf183a05980b70d08f0852875fe4c21c96e33f55 Mon Sep 17 00:00:00 2001
From: Kodi Craft <kodi@kdcf.me>
Date: Mon, 24 Jun 2024 15:58:54 +0200
Subject: [PATCH] Fix crash in server due to overeager parsing

---
 src/lib.rs    |  2 +-
 tests/full.rs | 43 ++++++++++++++++++++++++++++---------------
 2 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 6d52d20..22dbe55 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -320,7 +320,7 @@ fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream
                                 Ok(n) => {
                                     #debug("Received {} bytes (server)", n);
                                     buf.extend_from_slice(&read_buf[..n]);
-                                    loop {
+                                    while buf.len() >= 4 {
                                         let len = u32::from_le_bytes(buf[..4].try_into().expect("Failed to convert bytes to u32"));
                                         if buf.len() >= (4 + len as usize) {
                                             let serialized = std::str::from_utf8(&buf[4..(4 + len as usize)]).expect("Failed to convert bytes to string");
diff --git a/tests/full.rs b/tests/full.rs
index 30fedcc..7c33a4f 100644
--- a/tests/full.rs
+++ b/tests/full.rs
@@ -53,29 +53,42 @@ impl TestProtocolServerTrait for TrivialServer {
     async fn void(&mut self) {}
 }
 
+struct Cleanup {
+    address: String,
+}
+impl Drop for Cleanup {
+    fn drop(&mut self) {
+        std::fs::remove_file(&self.address).unwrap();
+    }
+}
+
 #[tokio::test]
 async fn e2e() {
     init_logger();
     #[cfg(feature = "unix")]
-    let address = "/tmp/eagle-test.sock";
+    let address = format!("/tmp/eagle-test-{}.sock", rand::random::<u64>());
+    #[cfg(feature = "unix")]
+    let _cleanup = Cleanup {
+        address: address.clone(),
+    };
     #[cfg(feature = "tcp")]
     let address = format!("127.0.0.1:{}", 10000 + rand::random::<u64>() % 1000);
     let server_task = tokio::spawn(TestProtocolServer::bind(TrivialServer, address.clone()));
     // Wait for the server to start
     tokio::time::sleep(std::time::Duration::from_millis(10)).await;
     let client = TestProtocolClient::connect(address).await.unwrap();
-    let res = client.addition(2, 5).await.unwrap();
-    // assert_eq!(client.addition(2, 5).await.unwrap(), 7);
-    // assert_eq!(
-    //     client.some_kind_of_question("Hello, world!".to_string())
-    //         .await
-    //         .unwrap(),
-    //     "Hello, world!".len() as i32
-    // );
-    // assert_eq!(
-    //     client.this_responds_with_a_string(42).await.unwrap(),
-    //     "The number is 42"
-    // );
-    // client.void().await.unwrap();
-    // server_task.abort();
+    assert_eq!(client.addition(2, 5).await.unwrap(), 7);
+    assert_eq!(
+        client
+            .some_kind_of_question("Hello, world!".to_string())
+            .await
+            .unwrap(),
+        "Hello, world!".len() as i32
+    );
+    assert_eq!(
+        client.this_responds_with_a_string(42).await.unwrap(),
+        "The number is 42"
+    );
+    client.void().await.unwrap();
+    server_task.abort();
 }