feat: graceful disconnection

src/config.rs

@@ -6,6 +6,7 @@ pub struct Config {
     pub web_port: u16,
     pub handshake_timeout_secs: u64,
     pub punch_hole_secs: u64,
+    pub disconnect_timeout_ms: u64,
 }
 
 impl Config {
@@ -17,6 +18,7 @@ impl Config {
             web_port: 8080,
             handshake_timeout_secs: 30,
             punch_hole_secs: 15,
+            disconnect_timeout_ms: 500,
         }
     }
 }

src/main.rs

@@ -5,7 +5,7 @@ mod web;
 
 use crate::{
     config::Config,
-    messaging::message_manager::MessageManager,
+    messaging::message_manager::{MessageManager, StreamMessage},
     web::{
         shared_state::{AppEvent, AppState, Command, SharedState, Status},
         web_server,
@@ -44,8 +44,7 @@ async fn main() -> Result<()> {
     let (event_tx, _event_rx) = broadcast::channel::<AppEvent>(32);
 
     // Initialize Shared State
-    // Note: We use the new constructor which automatically defaults internal fields.
-    let shared_state = Arc::new(RwLock::new(AppState::new(cmd_tx, event_tx)));
+    let shared_state = Arc::new(RwLock::new(AppState::new(cmd_tx.clone(), event_tx)));
 
     // Spawn Web Server
     let state_clone = Arc::clone(&shared_state);
@@ -56,13 +55,33 @@ async fn main() -> Result<()> {
         }
     });
 
-    // Start the Controller (Main Logic)
-    // We await this as it runs the main event loop
+    // Spawn signal handler for graceful shutdown
+    let cmd_tx_clone = cmd_tx.clone();
+    let disconnect_timeout = config.disconnect_timeout_ms;
+    tokio::spawn(async move {
+        match tokio::signal::ctrl_c().await {
+            Ok(()) => {
+                info!("Received Ctrl+C signal, initiating graceful shutdown");
+                // Send disconnect command
+                if let Err(e) = cmd_tx_clone.send(Command::Disconnect).await {
+                    warn!("Failed to send disconnect command on shutdown: {}", e);
+                }
+                // Time for disconnect to complete
+                tokio::time::sleep(Duration::from_millis(disconnect_timeout)).await;
+                std::process::exit(0);
+            }
+            Err(e) => {
+                error!("Failed to listen for Ctrl+C: {}", e);
+            }
+        }
+    });
+
+    // Start the main controller
     if let Err(e) = start_controller(&config, &shared_state, cmd_rx).await {
         error!("Controller error: {:?}", e);
     }
 
-    // Wait for web server (optional, usually controller keeps app alive)
+    // Wait for web server
     let _ = web_server_handle.await;
 
     Ok(())
@@ -82,7 +101,6 @@ async fn start_controller(
     mut cmd_rx: mpsc::Receiver<Command>,
 ) -> Result<()> {
     // 1. Bind the UDP Socket
-    // We bind to 0.0.0.0 to listen on all interfaces.
     let socket = UdpSocket::bind(("0.0.0.0", config.client_port)).await?;
     let socket = Arc::new(socket);
 
@@ -105,13 +123,11 @@ async fn start_controller(
     }
 
     // 3. Resolve Public IP via STUN
-    // Note: We pass a reference to the socket. net::resolve_public_ip now expects &UdpSocket.
     match net::resolve_public_ip(&socket, &config.stun_server).await {
         Ok(public_addr) => {
             info!("Public IP resolved via STUN: {}", public_addr);
 
-            // Update state safely using the setter.
-            // This triggers an event update so the UI displays the IP immediately.
+            // Update state
             shared_state.write().await.set_public_ip(
                 public_addr,
                 Some("Public IP resolved".into()),
@@ -174,7 +190,7 @@ async fn start_controller(
 
                             Command::SendMessage(msg) => {
                                 if message_manager.is_connected() {
-                                    match message_manager.send_message(msg.as_bytes()).await {
+                                    match message_manager.send_text(msg.clone()).await {
                                         Ok(_) => {
                                             shared_state.read().await.add_message(msg, true);
                                         },
@@ -184,23 +200,40 @@ async fn start_controller(
                                     warn!("Cannot send message: not connected to peer");
                                 }
                             }
+
+                            Command::Disconnect => {
+                                debug!("Disconnect command received");
+                                if let Err(e) = message_manager.disconnect().await {
+                                    error!("Disconnect failed: {:?}", e);
+                                } else {
+                                    info!("Successfully disconnected from peer");
+                                }
+                            }
                         }
                     }
                     None => {
-                        info!("Command channel closed, shutting down");
+                        debug!("Command channel closed, shutting down");
                         break;
                     }
                 }
             }
 
-
             // 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();
-                        debug!("Received message from peer: {}", msg_str);
-                        shared_state.read().await.add_message(msg_str, false);
+                        match bincode::deserialize::<StreamMessage>(&recv_buf[..n]) {
+                            Ok(StreamMessage::Bye) => {
+                                info!("Received BYE from peer. Disconnecting.");
+                                let _ = message_manager.disconnect_on_bye_received().await;
+                            }
+                            Ok(StreamMessage::Text(content)) => {
+                                shared_state.read().await.add_message(content, false);
+                            }
+                            Err(e) => {
+                                warn!("Failed to deserialize KCP packet: {}", e);
+                            }
+                        }
                     },
                     Err(e) => {
                         error!("KCP stream error: {}", e);
@@ -210,8 +243,7 @@ async fn start_controller(
 
             // C. Handle Keep-Alive (Heartbeat)
             _ = keep_alive_interval.tick() => {
-                // Only need to keep the NAT open if we are NOT connected.
-                // If we are connected, the MessageManager (chat session) handles traffic.
+                // Keep the NAT open if not connected.
                 let status = shared_state.read().await.status;
 
                 if status == Status::Disconnected {

src/messaging/handshake.rs

@@ -10,7 +10,7 @@ use tracing::{debug, info, warn};
 
 /// Represents handshake message being sent or received.
 #[derive(Serialize, Deserialize, Debug, PartialEq)]
-enum HandshakeMsg {
+pub enum HandshakeMsg {
     Syn,
     SynAck,
     Bye,
@@ -116,12 +116,12 @@ pub async fn handshake(
 
                         }
                         HandshakeMsg::SynAck => {
-                            info!("Received SYN-ACK from {}. Sending ACK.", sender);
+                            info!("Received SYN-ACK from {}.", sender);
 
                             // Transition to Connected state
                             state.write().await.set_status(
                                 Status::Punching,
-                                Some(format!("Received SYN-ACK from {}. Sending ACK.", sender)),
+                                Some(format!("Received SYN-ACK from {}.", sender)),
                                 Some(secs_left),
                             );
 

src/messaging/message_manager.rs

@@ -1,15 +1,16 @@
 use super::{
     super::web::shared_state::{SharedState, Status},
-    handshake,
+    handshake::{self, HandshakeMsg},
 };
 use anyhow::{Result, bail};
+use serde::{Deserialize, Serialize};
 use std::{net::SocketAddr, sync::Arc};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     net::UdpSocket,
 };
 use tokio_kcp::{KcpConfig, KcpNoDelayConfig, KcpStream};
-use tracing::{error, info};
+use tracing::{error, info, warn};
 
 /// Manages the lifecycle of a P2P connection, handling the transition from raw UDP to reliable KCP.
 ///
@@ -29,6 +30,15 @@ pub struct MessageManager {
     kcp_stream: Option<KcpStream>,
 }
 
+/// Represents a message being sent/received to/from a peer.
+#[derive(Serialize, Deserialize, Debug)]
+pub enum StreamMessage {
+    /// Regular chat content
+    Text(String),
+    /// Signal to close connection
+    Bye,
+}
+
 impl MessageManager {
     /// Creates a new `MessageManager` in a disconnected state.
     ///
@@ -106,7 +116,7 @@ impl MessageManager {
         if let Some(peer_addr) = self.peer_addr {
             info!("Upgrading connection to KCP with {}", peer_addr);
 
-            // Configure KCP for low-latency (Turbo Mode)
+            // Configure KCP for low-latency
             let config = KcpConfig {
                 nodelay: KcpNoDelayConfig {
                     nodelay: true,
@@ -133,12 +143,22 @@ impl MessageManager {
         }
     }
 
+    /// Sends a text message wrapped in the StreamMessage protocol
+    ///
+    /// # Arguments
+    ///
+    /// * `text` - Message to send.
+    pub async fn send_text(&mut self, text: String) -> Result<()> {
+        let payload = bincode::serialize(&StreamMessage::Text(text))?;
+        self.send_raw(&payload).await
+    }
+
     /// 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<()> {
+    async fn send_raw(&mut self, payload: &[u8]) -> Result<()> {
         if let Some(stream) = &mut self.kcp_stream {
             stream.write_all(payload).await?;
             stream.flush().await?;
@@ -212,6 +232,92 @@ impl MessageManager {
         }
     }
 
+    /// Gracefully disconnects from the peer by sending a Bye message and cleaning up resources.
+    ///
+    /// This method:
+    /// 1. Sends a Bye message to the peer (over KCP if connected, UDP as fallback)
+    /// 2. Closes the KCP stream if active
+    /// 3. Resets the connection state
+    /// 4. Updates shared state to Disconnected
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(())` - Disconnection successful
+    /// * `Err` - If sending the Bye message fails (cleanup still proceeds)
+    pub async fn disconnect(&mut self) -> Result<()> {
+        self.disconnect_internal(true).await
+    }
+
+    /// Disconnects from peer without sending Bye (used when receiving Bye from peer).
+    ///
+    /// This method performs cleanup without notifying the peer, since they already
+    /// initiated the disconnect.
+    ///
+    /// # Returns
+    ///
+    /// * `Ok(())` - Disconnection successful
+    pub async fn disconnect_on_bye_received(&mut self) -> Result<()> {
+        self.disconnect_internal(false).await
+    }
+
+    /// Internal disconnect implementation with option to send Bye message.
+    ///
+    /// # Arguments
+    ///
+    /// * `send_bye` - If true, sends Bye message to peer before cleanup
+    async fn disconnect_internal(&mut self, send_bye: bool) -> Result<()> {
+        info!("Initiating graceful disconnect (send_bye: {})", send_bye);
+
+        // Send Bye message to peer only if requested
+        if send_bye && let Some(peer_addr) = self.peer_addr {
+            let mut sent_via_kcp = false;
+
+            // Try to send via KCP first if available
+            if self.kcp_stream.is_some() {
+                let bye_packet = bincode::serialize(&StreamMessage::Bye)?;
+                match self.send_raw(&bye_packet).await {
+                    Ok(_) => {
+                        info!("Sent Bye message to peer via KCP");
+                        sent_via_kcp = true;
+                    }
+                    Err(e) => {
+                        warn!("Failed to send Bye via KCP: {}. Will try UDP fallback.", e);
+                    }
+                }
+            }
+
+            // 2. Fallback: UDP Raw (HandshakeMsg::Bye)
+            if !sent_via_kcp {
+                let udp_bye = bincode::serialize(&HandshakeMsg::Bye)?;
+                match self.client_socket.send_to(&udp_bye, peer_addr).await {
+                    Ok(_) => info!("Sent HandshakeMsg::Bye via UDP"),
+                    Err(e) => warn!("Failed to send Bye via UDP: {}", e),
+                }
+            }
+        }
+
+        // Close KCP stream if active
+        if let Err(e) = self.close_kcp().await {
+            warn!("Error closing KCP stream during disconnect: {}", e);
+        }
+
+        // Reset connection state
+        self.peer_addr = None;
+
+        // Clear chat history
+        self.state.read().await.clear_chat();
+
+        // Update shared state
+        self.state.write().await.set_status(
+            Status::Disconnected,
+            Some("Disconnected from peer".into()),
+            None,
+        );
+
+        info!("Disconnect complete");
+        Ok(())
+    }
+
     /// Closes the active KCP stream gracefully.
     ///
     /// This method:
@@ -225,7 +331,7 @@ impl MessageManager {
         if let Some(mut stream) = self.kcp_stream.take() {
             info!("Initiating KCP stream shutdown...");
 
-            // Attempt graceful shutdown. We log errors but don't fail the function
+            // Attempt graceful shutdown. Log errors but not fail the function
             if let Err(e) = stream.shutdown().await {
                 error!("Error during KCP shutdown: {}", e);
             } else {
@@ -289,7 +395,7 @@ 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;
+        let result = manager.send_text("hello".into()).await;
         assert!(result.is_err());
     }
 
@@ -346,4 +452,36 @@ mod tests {
             "Original socket should remain valid after cloned socket is dropped"
         );
     }
+
+    #[tokio::test]
+    async fn test_disconnect_without_connection() {
+        let mut manager = create_test_manager().await;
+
+        // Disconnect without being connected should work (idempotent)
+        let result = manager.disconnect().await;
+        assert!(result.is_ok());
+
+        // Verify state was updated to Disconnected
+        let state_guard = manager.state.read().await;
+        assert_eq!(state_guard.status, Status::Disconnected);
+    }
+
+    #[tokio::test]
+    async fn test_disconnect_with_peer_addr() {
+        let mut manager = create_test_manager().await;
+
+        // Set a peer address (simulating a connection)
+        manager.peer_addr = Some("127.0.0.1:9999".parse().unwrap());
+
+        // Disconnect
+        let result = manager.disconnect().await;
+        assert!(result.is_ok());
+
+        // Verify peer_addr is cleared
+        assert!(manager.peer_addr.is_none());
+
+        // Verify state was updated to Disconnected
+        let state_guard = manager.state.read().await;
+        assert_eq!(state_guard.status, Status::Disconnected);
+    }
 }