From 3b30432e5f1c6860044c77fe60d2143eed235517 Mon Sep 17 00:00:00 2001
From: Sachymetsu <sachymetsu@tutamail.com>
Date: Thu, 2 Apr 2026 11:59:18 +0200
Subject: [PATCH] Improve handshake net code

---
 src/net.rs | 23 ++++++++++++++++++-----
 1 file changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/net.rs b/src/net.rs
index c3666be..85cbb8e 100644
--- a/src/net.rs
+++ b/src/net.rs
@@ -129,20 +129,33 @@ impl Plugin for NetPlugin {
 async fn transport_handshake(tcp: &Async<TcpStream>) -> color_eyre::Result<Client> {
     let (encap, handshake) = ClientBuilder::start();
 
-    tcp.write_with(|mut tcp| tcp.write_all(encap.serialize().as_bytes()))
-        .await?;
+    let msg = encap.serialize();
+    let msg = msg.as_bytes();
+    let msg_len = u16::try_from(msg.len()).unwrap().to_be_bytes();
+
+    tcp.write_with(|mut tcp| tcp.write_all(&msg_len)).await?;
+    tcp.write_with(|mut tcp| tcp.write_all(msg)).await?;
 
     let mut payload = vec![0u8; 1024];
+    let mut msg_len = [0u8; 2];
+
+    let msg_len = tcp
+        .read_with(|mut tcp| {
+            tcp.read_exact(&mut msg_len)?;
+            Ok(usize::from(u16::from_be_bytes(msg_len)))
+        })
+        .await?;
 
-    let read = tcp.read_with(|mut tcp| tcp.read(&mut payload)).await?;
+    tcp.read_with(|mut tcp| tcp.read_exact(&mut payload[..msg_len]))
+        .await?;
 
-    let transport = handshake.finish(&payload[..read], &NOISE_PSK)?;
+    let transport = handshake.finish(&payload[..msg_len], &NOISE_PSK)?;
 
     Ok(transport)
 }
 
 async fn recv(stream: &Async<TcpStream>) -> std::io::Result<Vec<u8>> {
-    let mut msg_len_buf = [0_u8; 2];
+    let mut msg_len_buf = [0u8; 2];
     let mut msg = vec![0u8; 4096];
 
     stream
-- 
2.51.2