feat: use KCP for reliable message transmission

src/main.rs

@@ -117,10 +117,7 @@ async fn start_controller(
         }
     };
 
-    let message_manager = Arc::new(RwLock::new(MessageManager::new(
-        Arc::clone(&socket),
-        Arc::clone(shared_state),
-    )));
+    let mut message_manager = MessageManager::new(Arc::clone(&socket), Arc::clone(shared_state));
 
     // 4. Command Loop with Heartbeat
     info!("Waiting for commands (NAT Keep-Alive active)...");
@@ -129,6 +126,8 @@ async fn start_controller(
         tokio::time::interval(Duration::from_secs(config.punch_hole_secs));
     keep_alive_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
 
+    let mut recv_buf = [0u8; 4096];
+
     loop {
         tokio::select! {
             // A. Handle Commands from Web
@@ -145,16 +144,29 @@ async fn start_controller(
                                 if let Some(peer_addr) = peer_ip_opt {
                                     debug!("Initiating connection to peer: {}", peer_addr);
 
-                                    if let Err(e) = Arc::clone(&message_manager).write().await.handshake(peer_addr, config.handshake_timeout_secs).await {
+                                    if let Err(e) = message_manager.handshake(peer_addr, config.handshake_timeout_secs).await {
                                         error!("Connection attempt failed: {:?}", e);
                                     } else {
                                         debug!("Connection established successfully. MessageManager active.");
-                                        Arc::clone(&message_manager).write().await.upgrade_to_kcp().await?;
+                                        message_manager.upgrade_to_kcp().await?;
                                     }
                                 } else {
                                     warn!("Connect command received but no peer IP is set in state.");
                                 }
                             }
+
+                            Command::SendMessage(msg) => {
+                                if message_manager.is_connected() {
+                                    match message_manager.send_message(msg.as_bytes()).await {
+                                        Ok(_) => {
+                                            shared_state.read().await.add_message(msg, true);
+                                        },
+                                        Err(e) => error!("Failed to send KCP message: {}", e),
+                                    }
+                                } else {
+                                    warn!("Cannot send message: Not connected.");
+                                }
+                            }
                         }
                     }
                     None => {
@@ -163,15 +175,31 @@ async fn start_controller(
                     }
                 }
             }
-            // B. Handle Keep-Alive (Heartbeat)
+
+
+            // B. Handle Incoming KCP Messages (Only if connected)
+            res = message_manager.receive_message(&mut recv_buf), if message_manager.is_connected() => {
+                match res {
+                    Ok(n) => {
+                        let msg_str = String::from_utf8_lossy(&recv_buf[..n]).to_string();
+                        info!("Received message: {}", msg_str);
+                        shared_state.read().await.add_message(msg_str, false);
+                    },
+                    Err(e) => {
+                        error!("Error reading from KCP stream: {}", e);
+                    }
+                }
+            }
+
+            // C. Handle Keep-Alive (Heartbeat)
             _ = keep_alive_interval.tick() => {
-                // We only need to keep the NAT open if we are NOT connected.
+                // Only need to keep the NAT open if we are NOT connected.
                 // If we are connected, the MessageManager (chat session) handles traffic.
                 let status = shared_state.read().await.status;
 
                 if status == Status::Disconnected {
                     debug!("Sending Keep-Alive packet to STUN server...");
-                    // Re-resolving IP sends a STUN packet, which refreshes the NAT mapping.
+                    // Re resolving IP sends a STUN packet, which refreshes the NAT mapping.
                     match net::resolve_public_ip(&socket, &config.stun_server).await {
                         Ok(addr) => {
                             let mut state = shared_state.write().await;

src/messaging/message_manager.rs

@@ -4,7 +4,10 @@ use super::{
 };
 use anyhow::{Result, bail};
 use std::{net::SocketAddr, sync::Arc};
-use tokio::{io::AsyncWriteExt, net::UdpSocket};
+use tokio::{
+    io::{AsyncReadExt, AsyncWriteExt},
+    net::UdpSocket,
+};
 use tokio_kcp::{KcpConfig, KcpNoDelayConfig, KcpStream};
 use tracing::{error, info};
 
@@ -92,6 +95,8 @@ impl MessageManager {
     /// - Update Interval: 10ms
     /// - Resend: 2 (fast retransmission)
     /// - No Congestion Control (NC): enabled
+    /// - Windows: 1024 packets (allows higher throughput)
+    /// - MTU: 1400 (safe default for UDP)
     ///
     /// # Errors
     ///
@@ -109,6 +114,8 @@ impl MessageManager {
                     resend: 2,
                     nc: true,
                 },
+                wnd_size: (1024, 1024),
+                mtu: 1400,
                 ..Default::default()
             };
 
@@ -126,6 +133,44 @@ impl MessageManager {
         }
     }
 
+    /// Sends a binary message over the established KCP stream.
+    ///
+    /// # Arguments
+    ///
+    /// * `payload` - The bytes to send.
+    pub async fn send_message(&mut self, payload: &[u8]) -> Result<()> {
+        if let Some(stream) = &mut self.kcp_stream {
+            stream.write_all(payload).await?;
+            stream.flush().await?;
+            Ok(())
+        } else {
+            bail!("KCP stream not established");
+        }
+    }
+
+    /// Reads a message from the KCP stream into the provided buffer.
+    ///
+    /// # Arguments
+    ///
+    /// * `buf` - The buffer to write received data into.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(usize)` - The number of bytes read.
+    pub async fn receive_message(&mut self, buf: &mut [u8]) -> Result<usize> {
+        if let Some(stream) = &mut self.kcp_stream {
+            let n = stream.read(buf).await?;
+            Ok(n)
+        } else {
+            bail!("KCP stream not established");
+        }
+    }
+
+    /// Returns true if the KCP stream is currently active.
+    pub fn is_connected(&self) -> bool {
+        self.kcp_stream.is_some()
+    }
+
     /// Helper to clone the underlying UDP socket safely.
     ///
     /// `tokio-kcp` requires ownership of a `UdpSocket`, but we only have an `Arc<UdpSocket>`.
@@ -144,14 +189,13 @@ impl MessageManager {
             let fd = self.client_socket.as_raw_fd();
 
             // Create a std::net::UdpSocket from the raw fd.
-            // SAFETY: We must ensure we don't drop this variable normally,
-            // otherwise it will close the fd that our Arc<UdpSocket> relies on.
+            // must not drop this variable normally, it will close the fd Arc<UdpSocket> relies on.
             let std_sock = unsafe { std::net::UdpSocket::from_raw_fd(fd) };
 
             // try_clone() creates a new file descriptor (dup) that refers to the same socket.
             let new_std_sock = std_sock.try_clone();
 
-            // CRITICAL: Forget the original wrapper so the destructor doesn't fire and close the fd.
+            // Forget the original wrapper so the destructor doesn't fire and close the fd.
             std::mem::forget(std_sock);
 
             let new_std_sock = new_std_sock?;
@@ -182,7 +226,6 @@ impl MessageManager {
             info!("Initiating KCP stream shutdown...");
 
             // Attempt graceful shutdown. We log errors but don't fail the function
-            // because our goal is to clean up the local handle regardless of network state.
             if let Err(e) = stream.shutdown().await {
                 error!("Error during KCP shutdown: {}", e);
             } else {
@@ -227,6 +270,7 @@ mod tests {
         let manager = create_test_manager().await;
         assert!(manager.peer_addr.is_none());
         assert!(manager.kcp_stream.is_none());
+        assert!(!manager.is_connected());
     }
 
     #[tokio::test]
@@ -242,6 +286,21 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_send_fails_without_kcp() {
+        let mut manager = create_test_manager().await;
+        let result = manager.send_message(b"hello").await;
+        assert!(result.is_err());
+    }
+
+    #[tokio::test]
+    async fn test_receive_fails_without_kcp() {
+        let mut manager = create_test_manager().await;
+        let mut buf = [0u8; 10];
+        let result = manager.receive_message(&mut buf).await;
+        assert!(result.is_err());
+    }
+
     #[tokio::test]
     async fn test_close_kcp_safe_on_none() {
         let mut manager = create_test_manager().await;

src/web/shared_state.rs

@@ -92,7 +92,7 @@ impl AppState {
         timeout: Option<u64>,
     ) {
         self.public_ip = Some(addr);
-        self.broadcast_event(message, timeout);
+        self.broadcast_status_change(message, timeout);
     }
 
     /// Updates the NAT type and notifies listeners.
@@ -104,28 +104,26 @@ impl AppState {
         timeout: Option<u64>,
     ) {
         self.nat_type = nat_type;
-        self.broadcast_event(message, timeout);
+        self.broadcast_status_change(message, timeout);
     }
 
     /// Updates the connection status and notifies listeners.
     pub fn set_status(&mut self, status: Status, message: Option<String>, timeout: Option<u64>) {
         self.status = status;
-        self.broadcast_event(message, timeout);
+        self.broadcast_status_change(message, timeout);
     }
 
     /// Updates the peer's IP address and notifies listeners.
     pub fn set_peer_ip(&mut self, addr: SocketAddr, message: Option<String>, timeout: Option<u64>) {
         self.peer_ip = Some(addr);
-        self.broadcast_event(message, timeout);
+        self.broadcast_status_change(message, timeout);
     }
 
     /// Broadcasts the current state to all active listeners (e.g., Web UI).
     ///
     /// This constructs an `AppEvent` based on the current `status` and sends it
     /// via the `event_tx` channel.
-    ///
-    /// Note: If no receivers are active (e.g., no browser connected), the error is ignored.
-    fn broadcast_event(&self, message: Option<String>, timeout: Option<u64>) {
+    fn broadcast_status_change(&self, message: Option<String>, timeout: Option<u64>) {
         let event = match self.status {
             // When disconnected, we send the full state so the UI can sync up.
             Status::Disconnected => AppEvent::Disconnected {
@@ -136,9 +134,16 @@ impl AppState {
             // When connected, we send status messages.
             Status::Connected => AppEvent::Connected { message },
         };
+        self.broadcast_event(event);
+    }
+
+    /// Explicitly broadcasts a new chat message to the UI.
+    pub fn add_message(&self, content: String, from_me: bool) {
+        let _ = self.event_tx.send(AppEvent::Message { content, from_me });
+    }
 
-        // We ignore the error here because it's valid for there to be no active
-        // subscribers (e.g., the user hasn't opened the web page yet).
+    /// General helper for other updates if needed
+    fn broadcast_event(&self, event: AppEvent) {
         let _ = self.event_tx.send(event);
     }
 }
@@ -168,7 +173,9 @@ pub enum AppEvent {
     /// The application is idle or disconnected.
     ///
     /// Payload includes the full `AppState` to ensure the UI is fully synchronized.
-    Disconnected { state: AppState },
+    Disconnected {
+        state: AppState,
+    },
 
     /// The application is actively trying to punch a hole through the NAT.
     Punching {
@@ -183,6 +190,11 @@ pub enum AppEvent {
         /// Informational message from the peer or system.
         message: Option<String>,
     },
+
+    Message {
+        content: String,
+        from_me: bool,
+    },
 }
 
 /// Represents the high level connection state of the P2P node.
@@ -204,4 +216,6 @@ pub enum Status {
 pub enum Command {
     /// Instructs the controller to initiate a connection attempt to the configured peer.
     ConnectPeer,
+
+    SendMessage(String),
 }

src/web/web_server.rs

@@ -54,6 +54,7 @@ pub fn router(shared_state: SharedState) -> Router {
         // API Routes
         .route("/api/state", get(get_ip))
         .route("/api/connect", post(connect_peer))
+        .route("/api/message", post(send_message))
         .route("/api/events", get(sse_handler))
         // Static File Serving (Fallback)
         .fallback_service(ServeDir::new("static").append_index_html_on_directories(true))
@@ -122,6 +123,38 @@ async fn connect_peer(
     Ok(StatusCode::OK)
 }
 
+#[derive(Debug, Deserialize)]
+struct SendMessageRequest {
+    message: String,
+}
+
+/// Handler for `POST /api/message`.
+async fn send_message(
+    State(state): State<SharedState>,
+    Json(input): Json<SendMessageRequest>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    if input.message.trim().is_empty() {
+        return Err((StatusCode::BAD_REQUEST, "Message cannot be empty".into()));
+    }
+
+    // Check if connected
+    if state.read().await.status != Status::Connected {
+        return Err((StatusCode::BAD_REQUEST, "Not connected to a peer".into()));
+    }
+
+    // Send command to controller
+    let cmd_tx = state.read().await.cmd_tx().clone();
+    if let Err(e) = cmd_tx.send(Command::SendMessage(input.message)).await {
+        error!("Failed to send Message command: {}", e);
+        return Err((
+            StatusCode::INTERNAL_SERVER_ERROR,
+            "Internal Controller Error".to_string(),
+        ));
+    }
+
+    Ok(StatusCode::OK)
+}
+
 /// Handler for `GET /api/events`.
 /// Upgrades the connection to a Server Sent Events (SSE) stream.
 async fn sse_handler(