Fix connected clients

This commit is contained in:
Jon Häggblad
2023-11-28 08:56:20 +01:00
parent 0581d40bbf
commit c2cb343135
+32 -20
View File
@@ -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<nym_tun::tun_task_channel::TunTaskResponseRx>,
mixnet_client: nym_sdk::mixnet::MixnetClient,
task_handle: TaskHandle,
connected_clients: HashMap<IpAddr, ConnectedClient>,
// connected_clients: HashMap<IpAddr, ConnectedClient>,
}
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<std::sync::Mutex<HashMap<IpAddr, ConnectedClient>>>,
) -> Result<Option<IpPacketResponse>, 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<std::sync::Mutex<HashMap<IpAddr, ConnectedClient>>>,
) -> Result<Option<IpPacketResponse>, 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<std::sync::Mutex<HashMap<IpAddr, ConnectedClient>>>,
) -> Result<Option<IpPacketResponse>, 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<std::sync::Mutex<HashMap<IpAddr, ConnectedClient>>>,
) -> Result<Option<IpPacketResponse>, 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::<IpAddr, ConnectedClient>::new()));
let connected_clients = Arc::new(std::sync::Mutex::new(
HashMap::<IpAddr, ConnectedClient>::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");