Implement drop for client handlers too

This commit is contained in:
Jon Häggblad
2024-02-08 15:27:23 +01:00
parent 2f0074821c
commit d4ca2a7220
3 changed files with 83 additions and 29 deletions
@@ -1,6 +1,8 @@
// Copyright 2024 - Nym Technologies SA <contact@nymtech.net>
// 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<Vec<u8>>,
mixnet_client_sender: nym_sdk::mixnet::MixnetClientSender,
close_rx: tokio::sync::oneshot::Receiver<()>,
finished_tx: Option<tokio::sync::oneshot::Sender<()>>,
activity_timeout: tokio::time::Interval,
}
impl ConnectedClientHandler {
pub(crate) fn launch(
pub(crate) fn start(
reply_to: Recipient,
reply_to_hops: Option<u8>,
mixnet_client_sender: nym_sdk::mixnet::MixnetClientSender,
) -> (
tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
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<u8>) -> 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(());
}
}
}
@@ -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<u8>,
forward_from_tun_tx: tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
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<IpAddr> = 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<IpAddr> = 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<tokio::sync::oneshot::Sender<()>>,
// 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<u8>,
pub(crate) forward_from_tun_tx: tokio::sync::mpsc::UnboundedSender<Vec<u8>>,
}
@@ -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<u8>,
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<Vec<u8>>,
}
@@ -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<tokio_tun::Tun>,
// 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(())