From a4c4d78b21a4fbcb87d3e0dda5f12372fa3e1932 Mon Sep 17 00:00:00 2001 From: pushkar-gr Date: Wed, 24 Dec 2025 18:39:11 +0530 Subject: [PATCH 1/2] feat: send_message and recieve_message methods to deliver message using KCP --- src/messaging/message_manager.rs | 54 +++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/src/messaging/message_manager.rs b/src/messaging/message_manager.rs index d1c649b..75c06e4 100644 --- a/src/messaging/message_manager.rs +++ b/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}; @@ -126,6 +129,40 @@ 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 { + if let Some(stream) = &mut self.kcp_stream { + // This will wait (await) until data is available. + let n = stream.read(buf).await?; + Ok(n) + } else { + bail!("KCP stream not established"); + } + } + /// Helper to clone the underlying UDP socket safely. /// /// `tokio-kcp` requires ownership of a `UdpSocket`, but we only have an `Arc`. @@ -242,6 +279,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; From e00e003cce7a38122463ecfe1646996fb4964edc Mon Sep 17 00:00:00 2001 From: pushkar-gr Date: Sun, 28 Dec 2025 10:08:44 +0530 Subject: [PATCH 2/2] feat: api to send message (closes #12) handles retransmitions and verifies delivery --- src/main.rs | 46 +++++++++++++++++++++++++------- src/messaging/message_manager.rs | 17 ++++++++---- src/web/shared_state.rs | 34 ++++++++++++++++------- src/web/web_server.rs | 33 +++++++++++++++++++++++ 4 files changed, 106 insertions(+), 24 deletions(-) diff --git a/src/main.rs b/src/main.rs index c0b547c..9b49460 100644 --- a/src/main.rs +++ b/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; diff --git a/src/messaging/message_manager.rs b/src/messaging/message_manager.rs index 75c06e4..0928fd4 100644 --- a/src/messaging/message_manager.rs +++ b/src/messaging/message_manager.rs @@ -95,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 /// @@ -112,6 +114,8 @@ impl MessageManager { resend: 2, nc: true, }, + wnd_size: (1024, 1024), + mtu: 1400, ..Default::default() }; @@ -155,7 +159,6 @@ impl MessageManager { /// * `Ok(usize)` - The number of bytes read. pub async fn receive_message(&mut self, buf: &mut [u8]) -> Result { if let Some(stream) = &mut self.kcp_stream { - // This will wait (await) until data is available. let n = stream.read(buf).await?; Ok(n) } else { @@ -163,6 +166,11 @@ impl MessageManager { } } + /// 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`. @@ -181,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 relies on. + // must not drop this variable normally, it will close the fd Arc 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?; @@ -219,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 { @@ -264,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] diff --git a/src/web/shared_state.rs b/src/web/shared_state.rs index ebc5fa8..d0f4e23 100644 --- a/src/web/shared_state.rs +++ b/src/web/shared_state.rs @@ -92,7 +92,7 @@ impl AppState { timeout: Option, ) { 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, ) { 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, timeout: Option) { 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, timeout: Option) { 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, timeout: Option) { + fn broadcast_status_change(&self, message: Option, timeout: Option) { 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, }, + + 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), } diff --git a/src/web/web_server.rs b/src/web/web_server.rs index cc041e5..a33f5bf 100644 --- a/src/web/web_server.rs +++ b/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, + Json(input): Json, +) -> Result { + 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(