diff --git a/service-providers/ip-packet-router/src/ip_packet_router.rs b/service-providers/ip-packet-router/src/ip_packet_router.rs index 907424fed2..bcc0bf0601 100644 --- a/service-providers/ip-packet-router/src/ip_packet_router.rs +++ b/service-providers/ip-packet-router/src/ip_packet_router.rs @@ -136,15 +136,14 @@ impl IpPacketRouter { // Channel used by the IpPacketRouter to signal connected and disconnected clients to the // TunListener - let (connected_clients, connected_client_rx) = mixnet_listener::ConnectedClients::new(); + let (connected_clients, connected_clients_rx) = mixnet_listener::ConnectedClients::new(); // let (connected_client_tx, connected_client_rx) = tokio::sync::mpsc::unbounded_channel(); let tun_listener = tun_listener::TunListener { tun_reader, mixnet_client_sender: mixnet_client.split_sender(), task_client: task_handle.get_handle(), - connected_clients: Default::default(), - connected_client_rx, + connected_clients: connected_clients_rx, }; tun_listener.start(); diff --git a/service-providers/ip-packet-router/src/mixnet_listener.rs b/service-providers/ip-packet-router/src/mixnet_listener.rs index 827749f6c0..66911843c0 100644 --- a/service-providers/ip-packet-router/src/mixnet_listener.rs +++ b/service-providers/ip-packet-router/src/mixnet_listener.rs @@ -30,9 +30,6 @@ pub(crate) struct MixnetListener { pub(crate) tun_writer: tokio::io::WriteHalf, pub(crate) mixnet_client: nym_sdk::mixnet::MixnetClient, pub(crate) task_handle: TaskHandle, - - // pub(crate) connected_clients: HashMap, - // pub(crate) connected_client_tx: tokio::sync::mpsc::UnboundedSender, pub(crate) connected_clients: ConnectedClients, } @@ -41,18 +38,54 @@ pub(crate) struct ConnectedClients { connected_client_tx: tokio::sync::mpsc::UnboundedSender, } +pub(crate) struct ConnectedClientsListener { + clients: HashMap, + pub(crate) connected_client_rx: tokio::sync::mpsc::UnboundedReceiver, +} + +impl ConnectedClientsListener { + pub(crate) fn get(&self, ip: &IpAddr) -> Option<&ConnectedClient> { + self.clients.get(ip) + } + + pub(crate) fn update(&mut self, event: ConnectedClientEvent) { + match event { + ConnectedClientEvent::Connect(connected_event) => { + let ConnectEvent { + ip, + nym_address, + mix_hops, + } = *connected_event; + log::trace!("Connect client: {ip}"); + self.clients.insert( + ip, + ConnectedClient { + nym_address, + mix_hops, + last_activity: std::time::Instant::now(), + }, + ); + } + ConnectedClientEvent::Disconnect(DisconnectEvent(ip)) => { + log::trace!("Disconnect client: {ip}"); + self.clients.remove(&ip); + } + } + } +} + impl ConnectedClients { - pub(crate) fn new() -> ( - Self, - tokio::sync::mpsc::UnboundedReceiver, - ) { + pub(crate) fn new() -> (Self, ConnectedClientsListener) { let (connected_client_tx, connected_client_rx) = tokio::sync::mpsc::unbounded_channel(); ( Self { clients: Default::default(), connected_client_tx, }, - connected_client_rx, + ConnectedClientsListener { + clients: Default::default(), + connected_client_rx, + }, ) } diff --git a/service-providers/ip-packet-router/src/tun_listener.rs b/service-providers/ip-packet-router/src/tun_listener.rs index 3b51c23789..5dcc59d205 100644 --- a/service-providers/ip-packet-router/src/tun_listener.rs +++ b/service-providers/ip-packet-router/src/tun_listener.rs @@ -1,15 +1,12 @@ -use std::{collections::HashMap, net::IpAddr}; - use nym_ip_packet_requests::IpPacketResponse; use nym_sdk::mixnet::MixnetMessageSender; use nym_task::TaskClient; #[cfg(target_os = "linux")] use tokio::io::AsyncReadExt; -use tokio::sync::mpsc::UnboundedReceiver; use crate::{ error::{IpPacketRouterError, Result}, - mixnet_listener::{self, ConnectEvent}, + mixnet_listener::{self}, util::{create_message::create_input_message, parse_ip::parse_dst_addr}, }; @@ -19,10 +16,7 @@ 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, - - // A mirror of the one in IpPacketRouter - pub(crate) connected_clients: HashMap, - pub(crate) connected_client_rx: UnboundedReceiver, + pub(crate) connected_clients: mixnet_listener::ConnectedClientsListener, } #[cfg(target_os = "linux")] @@ -58,36 +52,6 @@ impl TunListener { Ok(()) } - async fn handle_connected_client_event( - &mut self, - event: mixnet_listener::ConnectedClientEvent, - ) { - match event { - mixnet_listener::ConnectedClientEvent::Connect(connected_event) => { - let ConnectEvent { - ip, - nym_address, - mix_hops, - } = *connected_event; - log::trace!("Connect client: {ip}"); - self.connected_clients.insert( - ip, - mixnet_listener::ConnectedClient { - nym_address, - mix_hops, - last_activity: std::time::Instant::now(), - }, - ); - } - mixnet_listener::ConnectedClientEvent::Disconnect( - mixnet_listener::DisconnectEvent(ip), - ) => { - log::trace!("Disconnect client: {ip}"); - self.connected_clients.remove(&ip); - } - } - } - async fn run(mut self) -> Result<()> { let mut buf = [0u8; 65535]; while !self.task_client.is_shutdown() { @@ -95,8 +59,9 @@ impl TunListener { _ = self.task_client.recv() => { log::trace!("TunListener: received shutdown"); }, - event = self.connected_client_rx.recv() => match event { - Some(event) => self.handle_connected_client_event(event).await, + // TODO: ConnectedClientsListener::update should poll the channel instead + event = self.connected_clients.connected_client_rx.recv() => match event { + Some(event) => self.connected_clients.update(event), None => { log::error!("TunListener: connected client channel closed"); break;