diff --git a/service-providers/ip-packet-router/src/connected_client_handler.rs b/service-providers/ip-packet-router/src/connected_client_handler.rs index 253765d470..f838c561c3 100644 --- a/service-providers/ip-packet-router/src/connected_client_handler.rs +++ b/service-providers/ip-packet-router/src/connected_client_handler.rs @@ -1,6 +1,8 @@ // Copyright 2024 - Nym Technologies SA // SPDX-License-Identifier: GPL-3.0-only +use std::time::Duration; + use nym_ip_packet_requests::response::IpPacketResponse; use nym_sdk::mixnet::{MixnetMessageSender, Recipient}; @@ -9,6 +11,8 @@ use crate::{ util::create_message::create_input_message, }; +const ACTIVITY_TIMEOUT_SEC: u64 = 15 * 60; + // Data flow // Out: mixnet_listener -> decode -> handle_packet -> write_to_tun // In: tun_listener -> [connected_client_handler -> encode] -> mixnet_sender @@ -21,18 +25,22 @@ pub(crate) struct ConnectedClientHandler { forward_from_tun_rx: tokio::sync::mpsc::UnboundedReceiver>, mixnet_client_sender: nym_sdk::mixnet::MixnetClientSender, close_rx: tokio::sync::oneshot::Receiver<()>, + finished_tx: Option>, + activity_timeout: tokio::time::Interval, } impl ConnectedClientHandler { - pub(crate) fn launch( + pub(crate) fn start( reply_to: Recipient, reply_to_hops: Option, mixnet_client_sender: nym_sdk::mixnet::MixnetClientSender, ) -> ( tokio::sync::mpsc::UnboundedSender>, tokio::sync::oneshot::Sender<()>, + tokio::sync::oneshot::Receiver<()>, ) { let (close_tx, close_rx) = tokio::sync::oneshot::channel(); + let (finished_tx, finished_rx) = tokio::sync::oneshot::channel(); let (forward_from_tun_tx, forward_from_tun_rx) = tokio::sync::mpsc::unbounded_channel(); let connected_client_handler = ConnectedClientHandler { @@ -41,6 +49,8 @@ impl ConnectedClientHandler { forward_from_tun_rx, mixnet_client_sender, close_rx, + finished_tx: Some(finished_tx), + activity_timeout: tokio::time::interval(Duration::from_secs(ACTIVITY_TIMEOUT_SEC)), }; tokio::spawn(async move { @@ -49,10 +59,11 @@ impl ConnectedClientHandler { } }); - (forward_from_tun_tx, close_tx) + (forward_from_tun_tx, close_tx, finished_rx) } async fn handle_packet(&mut self, packet: Vec) -> Result<()> { + self.activity_timeout.reset(); let response_packet = IpPacketResponse::new_ip_packet(packet.into()) .to_bytes() .map_err(|err| IpPacketRouterError::FailedToSerializeResponsePacket { source: err })?; @@ -71,6 +82,10 @@ impl ConnectedClientHandler { log::trace!("ConnectedClientHandler: received shutdown"); break; }, + _ = self.activity_timeout.tick() => { + log::debug!("ConnectedClientHandler: activity timeout reached for {}", self.nym_address); + break; + } packet = self.forward_from_tun_rx.recv() => match packet { Some(packet) => { if let Err(err) = self.handle_packet(packet).await { @@ -81,7 +96,7 @@ impl ConnectedClientHandler { log::debug!("connected client handler: tun channel closed"); break; } - } + }, } } @@ -89,3 +104,12 @@ impl ConnectedClientHandler { Ok(()) } } + +impl Drop for ConnectedClientHandler { + fn drop(&mut self) { + log::trace!("ConnectedClientHandler: dropping"); + if let Some(finished_tx) = self.finished_tx.take() { + let _ = finished_tx.send(()); + } + } +} diff --git a/service-providers/ip-packet-router/src/mixnet_listener.rs b/service-providers/ip-packet-router/src/mixnet_listener.rs index e4b7fd1698..9ffd43bbb4 100644 --- a/service-providers/ip-packet-router/src/mixnet_listener.rs +++ b/service-providers/ip-packet-router/src/mixnet_listener.rs @@ -14,6 +14,7 @@ use nym_ip_packet_requests::{ use nym_sdk::mixnet::{MixnetMessageSender, Recipient}; use nym_sphinx::receiver::ReconstructedMessage; use nym_task::TaskHandle; +use tap::TapFallible; #[cfg(target_os = "linux")] use tokio::io::AsyncWriteExt; @@ -85,6 +86,7 @@ impl ConnectedClients { mix_hops: Option, forward_from_tun_tx: tokio::sync::mpsc::UnboundedSender>, close_tx: tokio::sync::oneshot::Sender<()>, + finished_rx: tokio::sync::oneshot::Receiver<()>, ) { // The map of connected clients that the mixnet listener keeps track of. It monitors // activity and disconnects clients that have been inactive for too long. @@ -95,6 +97,7 @@ impl ConnectedClients { mix_hops, last_activity: std::time::Instant::now(), close_tx: Some(close_tx), + finished_rx, }, ); // Send the connected client info to the tun listener, which will use it to forward packets @@ -102,11 +105,12 @@ impl ConnectedClients { self.connected_client_tx .send(ConnectedClientEvent::Connect(Box::new(ConnectEvent { ip, - nym_address, - mix_hops, forward_from_tun_tx, }))) - .unwrap(); + .tap_err(|err| { + log::error!("Failed to send connected client event: {err}"); + }) + .ok(); } fn update_activity(&mut self, ip: &IpAddr) -> Result<()> { @@ -118,6 +122,30 @@ impl ConnectedClients { } } + fn disconnect_stopped_client_handlers(&mut self) { + let stopped_clients: Vec = self + .clients + .iter_mut() + .filter_map(|(ip, client)| { + if client.finished_rx.try_recv().is_ok() { + Some(*ip) + } else { + None + } + }) + .collect(); + for ip in stopped_clients { + log::info!("Removing stopped client: {ip}"); + self.clients.remove(&ip); + self.connected_client_tx + .send(ConnectedClientEvent::Disconnect(DisconnectEvent(ip))) + .tap_err(|err| { + log::error!("Failed to send disconnect event: {err}"); + }) + .ok(); + } + } + fn disconnect_inactive_clients(&mut self) { let now = std::time::Instant::now(); let inactive_clients: Vec = self @@ -136,7 +164,10 @@ impl ConnectedClients { self.clients.remove(&ip); self.connected_client_tx .send(ConnectedClientEvent::Disconnect(DisconnectEvent(ip))) - .unwrap(); + .tap_err(|err| { + log::error!("Failed to send disconnect event: {err}"); + }) + .ok(); } } @@ -159,6 +190,9 @@ pub(crate) struct ConnectedClient { // Send to connected clients listener to stop. This is option only because we need to take // ownership of it when the client is dropped. pub(crate) close_tx: Option>, + + // Receive event when the client listener for that client stopped + pub(crate) finished_rx: tokio::sync::oneshot::Receiver<()>, } impl ConnectedClient { @@ -242,8 +276,8 @@ impl MixnetListener { log::info!("Connecting a new client"); // Spawn the ConnectedClientHandler for the new client - let (forward_from_tun_tx, close_tx) = - connected_client_handler::ConnectedClientHandler::launch( + let (forward_from_tun_tx, close_tx, finished_rx) = + connected_client_handler::ConnectedClientHandler::start( reply_to, reply_to_hops, self.mixnet_client.split_sender(), @@ -256,6 +290,7 @@ impl MixnetListener { reply_to_hops, forward_from_tun_tx, close_tx, + finished_rx, ); Ok(Some(IpPacketResponse::new_static_connect_success( request_id, reply_to, @@ -324,8 +359,8 @@ impl MixnetListener { }; // Spawn the ConnectedClientHandler for the new client - let (forward_from_tun_tx, close_tx) = - connected_client_handler::ConnectedClientHandler::launch( + let (forward_from_tun_tx, close_tx, finished_rx) = + connected_client_handler::ConnectedClientHandler::start( reply_to, reply_to_hops, self.mixnet_client.split_sender(), @@ -338,6 +373,7 @@ impl MixnetListener { reply_to_hops, forward_from_tun_tx, close_tx, + finished_rx, ); Ok(Some(IpPacketResponse::new_dynamic_connect_success( request_id, reply_to, new_ip, @@ -501,6 +537,7 @@ impl MixnetListener { log::debug!("IpPacketRouter [main loop]: received shutdown"); }, _ = disconnect_timer.tick() => { + self.connected_clients.disconnect_stopped_client_handlers(); self.connected_clients.disconnect_inactive_clients(); }, msg = self.mixnet_client.next() => { @@ -541,7 +578,5 @@ pub(crate) struct DisconnectEvent(pub(crate) IpAddr); pub(crate) struct ConnectEvent { pub(crate) ip: IpAddr, - pub(crate) nym_address: Recipient, - pub(crate) mix_hops: Option, pub(crate) forward_from_tun_tx: tokio::sync::mpsc::UnboundedSender>, } diff --git a/service-providers/ip-packet-router/src/tun_listener.rs b/service-providers/ip-packet-router/src/tun_listener.rs index 27d197c777..f299fb45f5 100644 --- a/service-providers/ip-packet-router/src/tun_listener.rs +++ b/service-providers/ip-packet-router/src/tun_listener.rs @@ -1,6 +1,5 @@ use std::{collections::HashMap, net::IpAddr}; -use nym_sdk::mixnet::Recipient; use nym_task::TaskClient; #[cfg(target_os = "linux")] use tokio::io::AsyncReadExt; @@ -11,11 +10,10 @@ use crate::{ util::parse_ip::parse_dst_addr, }; +// The TUN listener keeps a local map of the connected clients that has its state updated by the +// mixnet listener. Basically it's just so that we don't have to have mutexes around shared state. +// It's even ok if this is slightly out of date pub(crate) struct ConnectedClientMirror { - pub(crate) nym_address: Recipient, - pub(crate) mix_hops: Option, - pub(crate) last_activity: std::time::Instant, - // Forward packets we read from the TUN device to the connected clients listener pub(crate) forward_from_tun_tx: tokio::sync::mpsc::UnboundedSender>, } @@ -46,17 +44,12 @@ impl ConnectedClientsListener { mixnet_listener::ConnectedClientEvent::Connect(connected_event) => { let mixnet_listener::ConnectEvent { ip, - nym_address, - mix_hops, forward_from_tun_tx, } = *connected_event; log::trace!("Connect client: {ip}"); self.clients.insert( ip, ConnectedClientMirror { - nym_address, - mix_hops, - last_activity: std::time::Instant::now(), forward_from_tun_tx, }, ); @@ -75,7 +68,6 @@ impl ConnectedClientsListener { #[cfg(target_os = "linux")] pub(crate) struct TunListener { pub(crate) tun_reader: tokio::io::ReadHalf, - // pub(crate) mixnet_client_sender: nym_sdk::mixnet::MixnetClientSender, pub(crate) task_client: TaskClient, pub(crate) connected_clients: ConnectedClientsListener, } @@ -89,16 +81,19 @@ impl TunListener { }; if let Some(ConnectedClientMirror { - nym_address, - mix_hops, - last_activity, forward_from_tun_tx, }) = self.connected_clients.get(&dst_addr) { let packet = buf[..len].to_vec(); - forward_from_tun_tx.send(packet).unwrap(); + if forward_from_tun_tx.send(packet).is_err() { + log::warn!("Failed to forward packet to connected client {dst_addr}: disconnecting it from tun listener"); + self.connected_clients + .update(mixnet_listener::ConnectedClientEvent::Disconnect( + mixnet_listener::DisconnectEvent(dst_addr), + )); + } } else { - log::info!("No registered nym-address for packet - dropping"); + log::info!("No registered client for packet with destination {dst_addr} - dropping"); } Ok(())