diff --git a/service-providers/ip-packet-router/src/lib.rs b/service-providers/ip-packet-router/src/lib.rs index e7eb497b66..116a419511 100644 --- a/service-providers/ip-packet-router/src/lib.rs +++ b/service-providers/ip-packet-router/src/lib.rs @@ -4,7 +4,8 @@ use std::{ collections::HashMap, net::{IpAddr, SocketAddr}, path::Path, - time::Duration, sync::Arc, + sync::Arc, + time::Duration, }; use error::IpPacketRouterError; @@ -170,7 +171,7 @@ impl IpPacketRouterBuilder { tun_task_response_rx: Some(tun_task_response_rx), mixnet_client, task_handle, - connected_clients: Default::default(), + // connected_clients: Default::default(), }; log::info!("The address of this client is: {self_address}"); @@ -198,8 +199,7 @@ struct IpPacketRouter { tun_task_response_rx: Option, mixnet_client: nym_sdk::mixnet::MixnetClient, task_handle: TaskHandle, - - connected_clients: HashMap, + // connected_clients: HashMap, } struct ConnectedClient { @@ -212,6 +212,7 @@ impl IpPacketRouter { async fn on_static_connect_request( &mut self, connect_request: nym_ip_packet_requests::StaticConnectRequest, + connected_clients: &Arc>>, ) -> Result, IpPacketRouterError> { log::info!( "Received static connect request from {sender_address}", @@ -224,11 +225,16 @@ impl IpPacketRouter { // TODO: ignoring reply_to_hops and reply_to_avg_mix_delays for now // Check that the IP is available in the set of connected clients - let is_ip_taken = self.connected_clients.contains_key(&requested_ip); + // let is_ip_taken = self.connected_clients.contains_key(&requested_ip); + let is_ip_taken = connected_clients + .lock() + .unwrap() + .contains_key(&requested_ip); // Check that the nym address isn't already registered - let is_nym_address_taken = self - .connected_clients + let is_nym_address_taken = connected_clients + .lock() + .unwrap() .values() .any(|client| client.nym_address == reply_to); @@ -236,7 +242,7 @@ impl IpPacketRouter { (true, true) => { log::info!("Connecting an already connected client"); // Update the last activity time for the client - if let Some(client) = self.connected_clients.get_mut(&requested_ip) { + if let Some(client) = connected_clients.lock().unwrap().get_mut(&requested_ip) { client.last_activity = std::time::Instant::now(); } else { log::error!("Failed to update last activity time for client"); @@ -247,7 +253,7 @@ impl IpPacketRouter { } (false, false) => { log::info!("Connecting a new client"); - self.connected_clients.insert( + connected_clients.lock().unwrap().insert( requested_ip, ConnectedClient { nym_address: reply_to, @@ -280,6 +286,7 @@ impl IpPacketRouter { async fn on_dynamic_connect_request( &mut self, connect_request: nym_ip_packet_requests::DynamicConnectRequest, + connected_clients: &Arc>>, ) -> Result, IpPacketRouterError> { log::info!( "Received dynamic connect request from {sender_address}", @@ -293,7 +300,7 @@ impl IpPacketRouter { // Check if it's the same client connecting again, then we just reuse the same IP // TODO: this is problematic. Until we sign connect requests this means you can spam people // with return traffic - let existing_ip = self.connected_clients.iter().find_map(|(ip, client)| { + let existing_ip = connected_clients.lock().unwrap().iter().find_map(|(ip, client)| { if client.nym_address == reply_to { Some(*ip) } else { @@ -304,7 +311,7 @@ impl IpPacketRouter { if let Some(existing_ip) = existing_ip { log::info!("Found existing client for nym address"); // Update the last activity time for the client - if let Some(client) = self.connected_clients.get_mut(&existing_ip) { + if let Some(client) = connected_clients.lock().unwrap().get_mut(&existing_ip) { client.last_activity = std::time::Instant::now(); } else { log::error!("Failed to update last activity time for client"); @@ -316,7 +323,7 @@ impl IpPacketRouter { ))); } - let Some(new_ip) = generate_new_ip::find_new_ip(&self.connected_clients) else { + let Some(new_ip) = generate_new_ip::find_new_ip(&connected_clients.lock().unwrap()) else { log::info!("No available IP address"); return Ok(Some(IpPacketResponse::new_dynamic_connect_failure( request_id, @@ -325,7 +332,7 @@ impl IpPacketRouter { ))); }; - self.connected_clients.insert( + connected_clients.lock().unwrap().insert( new_ip, ConnectedClient { nym_address: reply_to, @@ -340,6 +347,7 @@ impl IpPacketRouter { async fn on_data_request( &mut self, data_request: nym_ip_packet_requests::DataRequest, + connected_clients: &Arc>>, ) -> Result, IpPacketRouterError> { log::info!("Received data request"); @@ -362,7 +370,7 @@ impl IpPacketRouter { // Check if there is a connected client for this src_addr. If there is, update the last activity time // for the client. If there isn't, drop the packet. - if let Some(client) = self.connected_clients.get_mut(&src_addr) { + if let Some(client) = connected_clients.lock().unwrap().get_mut(&src_addr) { client.last_activity = std::time::Instant::now(); } else { log::info!("Dropping packet: no connected client for {src_addr}"); @@ -393,6 +401,7 @@ impl IpPacketRouter { async fn on_reconstructed_message( &mut self, reconstructed: ReconstructedMessage, + connected_clients: &Arc>>, ) -> Result, IpPacketRouterError> { log::debug!( "Received message with sender_tag: {:?}", @@ -414,12 +423,13 @@ impl IpPacketRouter { match request.data { IpPacketRequestData::StaticConnect(connect_request) => { - self.on_static_connect_request(connect_request).await + self.on_static_connect_request(connect_request, connected_clients) + .await } IpPacketRequestData::DynamicConnect(connect_request) => { - self.on_dynamic_connect_request(connect_request).await + self.on_dynamic_connect_request(connect_request, connected_clients).await } - IpPacketRequestData::Data(data_request) => self.on_data_request(data_request).await, + IpPacketRequestData::Data(data_request) => self.on_data_request(data_request, connected_clients).await, } } @@ -430,7 +440,9 @@ impl IpPacketRouter { let mixnet_client_sender = self.mixnet_client.split_sender(); let mixnet_client_sender_clone = mixnet_client_sender.clone(); - let connected_clients = Arc::new(std::sync::Mutex::new(HashMap::::new())); + let connected_clients = Arc::new(std::sync::Mutex::new( + HashMap::::new(), + )); let connected_clients_clone = connected_clients.clone(); let tun_task_response_rx = self.tun_task_response_rx.take(); @@ -498,12 +510,12 @@ impl IpPacketRouter { .collect(); for ip in inactive_clients { log::info!("Disconnect inactive client: {ip}"); - self.connected_clients.remove(&ip); + connected_clients.lock().unwrap().remove(&ip); } }, msg = self.mixnet_client.next() => { if let Some(msg) = msg { - match self.on_reconstructed_message(msg).await { + match self.on_reconstructed_message(msg, &connected_clients).await { Ok(Some(response)) => { let Some(recipient) = response.recipient() else { log::error!("IpPacketRouter [main loop]: failed to get recipient from response");