From adf22e330ad32e280dfaa8e3d698f91e96d0e127 Mon Sep 17 00:00:00 2001 From: Andrej Mihajlov Date: Wed, 15 Oct 2025 09:27:44 +0200 Subject: [PATCH] Remove superficial AuthClientMixnetListener, replace with spawn() --- nym-authenticator-client/src/lib.rs | 3 +- .../src/mixnet_listener.rs | 167 +++++++----------- nym-gateway-probe/src/lib.rs | 8 +- nym-registration-client/src/lib.rs | 9 +- nym-registration-client/src/types.rs | 4 +- 5 files changed, 75 insertions(+), 116 deletions(-) diff --git a/nym-authenticator-client/src/lib.rs b/nym-authenticator-client/src/lib.rs index 8fe7edb511..e2b2248559 100644 --- a/nym-authenticator-client/src/lib.rs +++ b/nym-authenticator-client/src/lib.rs @@ -23,10 +23,9 @@ use nym_wireguard_types::PeerPublicKey; mod error; mod helpers; -mod mixnet_listener; +pub mod mixnet_listener; pub use crate::error::{Error, Result}; -pub use crate::mixnet_listener::{AuthClientMixnetListener, AuthClientMixnetListenerHandle}; pub struct AuthenticatorClient { mixnet_listener: MixnetMessageBroadcastReceiver, diff --git a/nym-authenticator-client/src/mixnet_listener.rs b/nym-authenticator-client/src/mixnet_listener.rs index 9cb6828e83..cf4815ba19 100644 --- a/nym-authenticator-client/src/mixnet_listener.rs +++ b/nym-authenticator-client/src/mixnet_listener.rs @@ -16,119 +16,81 @@ pub type MixnetMessageBroadcastReceiver = broadcast::Receiver; pub type MixnetMessageInputReceiver = mpsc::Receiver; // This could be another type, to abstract the mixnet message creation to here -// The AuthClientsMixnetListener listens to mixnet messages and rebroadcasts them to the +// Spawn a task that listens to mixnet messages and rebroadcasts them to the // AuthClients, or whoever else is interested. // It also manages the message input for the mixnet so it can keep the sole ownership of the MixnetClient // // NOTE: this is potentially bit wasteful. Ideally we should have proper channels where the // recipient only gets messages they're interested in. -pub struct AuthClientMixnetListener { - // The mixnet client that we're listening to - mixnet_client: MixnetClient, - - // Broadcast channel for the messages that we re-broadcast to the AuthClients - message_broadcast: MixnetMessageBroadcastSender, - - // Channel for message to send to the mixnet - input_message_tx: MixnetMessageInputSender, // we keep on to make sure it's open - input_message_rx: MixnetMessageInputReceiver, - - // Listen to cancel from the outside world +pub fn spawn( + mut mixnet_client: MixnetClient, shutdown_token: CancellationToken, -} +) -> AuthClientMixnetListenerHandle { + // Broadcast channel for the messages that we re-broadcast to the AuthClients + let (message_broadcast, _) = broadcast::channel(100); + // Channel for message to send to the mixnet + let (input_message_tx, mut input_message_rx) = mpsc::channel(100); -impl AuthClientMixnetListener { - pub fn new(mixnet_client: MixnetClient, shutdown_token: CancellationToken) -> Self { - let (message_broadcast, _) = broadcast::channel(100); - let (input_message_tx, input_message_rx) = mpsc::channel(100); - Self { - mixnet_client, - message_broadcast, - input_message_tx, - input_message_rx, - shutdown_token, - } - } + let cloned_message_broadcast = message_broadcast.clone(); + let cloned_message_sender = input_message_tx.clone(); + let child_shutdown_token = shutdown_token.child_token(); - async fn run(mut self) -> Self { - let mixnet_cancel_token = self.mixnet_client.cancellation_token(); - self.shutdown_token.run_until_cancelled(async { - loop { - tokio::select! { - biased; - _ = mixnet_cancel_token.cancelled() => { - tracing::debug!("AuthClientMixnetListener: mixnet client was shutdown"); - break; + let join_handle = tokio::spawn(async move { + loop { + tokio::select! { + biased; + _ = child_shutdown_token.cancelled() => { + tracing::debug!("AuthClientMixnetListener: received shutdown"); + // Disconnect mixnet which should send forget_me or remember_me if needed. + mixnet_client.disconnect().await; + break; + } + + // Sending loop + input_msg = input_message_rx.recv() => { + match input_msg { + None => { + tracing::error!("All senders were dropped. It shouldn't happen as we're holding one"); + break; + }, + Some(mix_msg) => { + if let Err(err) = mixnet_client.send(mix_msg).await { + tracing::error!("Failed to send mixnet message: {err}"); + } + }, } + } + // Receiving loop + msg = mixnet_client.next() => { + match msg { + None => { + tracing::error!("Mixnet client stream ended unexpectedly"); + break; + }, + Some(event) => { + if let Err(err) = message_broadcast.send(Arc::new(event)) { + tracing::error!("Failed to broadcast mixnet message: {err}"); + } + }, - // Sending loop - input_msg = self.input_message_rx.recv() => { - match input_msg { - None => { - tracing::error!("All senders were dropped. It shouldn't happen as we're holding one"); - break; - }, - Some(mix_msg) => { - if let Err(err) = self.mixnet_client.send(mix_msg).await { - tracing::error!("Failed to send mixnet message: {err}"); - } - }, - } - } - // Receiving loop - msg = self.mixnet_client.next() => { - match msg { - None => { - tracing::error!("Mixnet client stream ended unexpectedly"); - break; - }, - Some(event) => { - if let Err(err) = self.message_broadcast.send(Arc::new(event)) { - tracing::error!("Failed to broadcast mixnet message: {err}"); - } - }, - - } } } } - tracing::debug!("AuthClientMixnetListener is shutting down"); - }).await; - - self - } - - // Disconnects the mixnet client and effectively drop itself, since it doesn't work without one, and reconnecting isn't supported - pub async fn disconnect_mixnet_client(self) { - if !self.mixnet_client.cancellation_token().is_cancelled() { - self.mixnet_client.disconnect().await; } - } + tracing::debug!("AuthClientMixnetListener is shutting down"); + }); - pub fn start(self) -> AuthClientMixnetListenerHandle { - let message_broadcast = self.message_broadcast.clone(); - let message_sender = self.input_message_tx.clone(); - // Allows stopping only this, e.g. if we don't need it in the new bandwidth controller - let cancellation_token = self.shutdown_token.clone(); - let mixnet_cancellation_token = self.mixnet_client.cancellation_token(); - let handle = tokio::spawn(self.run()); - - AuthClientMixnetListenerHandle { - message_broadcast, - message_sender, - cancellation_token, - mixnet_cancellation_token, - handle, - } + AuthClientMixnetListenerHandle { + message_broadcast: cloned_message_broadcast, + message_sender: cloned_message_sender, + join_handle, } } pub struct AuthClientMixnetListenerHandle { message_broadcast: MixnetMessageBroadcastSender, message_sender: MixnetMessageInputSender, - cancellation_token: CancellationToken, - mixnet_cancellation_token: CancellationToken, - handle: JoinHandle, + join_handle: JoinHandle<()>, } impl AuthClientMixnetListenerHandle { @@ -140,21 +102,12 @@ impl AuthClientMixnetListenerHandle { self.message_broadcast.subscribe() } - pub fn mixnet_cancel_token(&self) -> CancellationToken { - self.mixnet_cancellation_token.clone() - } - - pub async fn stop(self) { - // If shutdown was externally called, that call is a no-op - // If we're only stopping this, it is very much needed - self.cancellation_token.cancel(); - match self.handle.await { - Ok(auth_client_mixnet_listener) => { - auth_client_mixnet_listener.disconnect_mixnet_client().await; - } - Err(e) => { - tracing::error!("Error waiting for auth clients mixnet listener to stop: {e}"); - } + /// Join on listener handle to wait until it stops + /// + /// Important: Use cancellation token to initiate the shutdown + pub async fn join(self) { + if let Err(e) = self.join_handle.await { + tracing::error!("Error waiting for auth clients mixnet listener to stop: {e}"); } } } diff --git a/nym-gateway-probe/src/lib.rs b/nym-gateway-probe/src/lib.rs index 2224c911e0..e66a1c9dd1 100644 --- a/nym-gateway-probe/src/lib.rs +++ b/nym-gateway-probe/src/lib.rs @@ -13,7 +13,7 @@ use base64::{Engine as _, engine::general_purpose}; use bytes::BytesMut; use clap::Args; use futures::StreamExt; -use nym_authenticator_client::{AuthClientMixnetListener, AuthenticatorClient}; +use nym_authenticator_client::{AuthenticatorClient, mixnet_listener}; use nym_authenticator_requests::{ AuthenticatorVersion, client_message::ClientMessage, response::AuthenticatorResponse, v2, v3, v4, v5, @@ -400,8 +400,9 @@ impl Probe { (node_info.authenticator_address, node_info.ip_address) { // Start the mixnet listener that the auth clients use to receive messages. + let mixnet_listener_cancel_token = CancellationToken::new(); let mixnet_listener_task = - AuthClientMixnetListener::new(mixnet_client, CancellationToken::new()).start(); + mixnet_listener::spawn(mixnet_client, mixnet_listener_cancel_token.child_token()); let auth_client = AuthenticatorClient::new( mixnet_listener_task.subscribe(), @@ -441,7 +442,8 @@ impl Probe { .await .unwrap_or_default(); - mixnet_listener_task.stop().await; + mixnet_listener_cancel_token.cancel(); + mixnet_listener_task.join().await; outcome } else { diff --git a/nym-registration-client/src/lib.rs b/nym-registration-client/src/lib.rs index c9794fa9a0..3fddc076ae 100644 --- a/nym-registration-client/src/lib.rs +++ b/nym-registration-client/src/lib.rs @@ -3,7 +3,7 @@ use tokio_util::sync::CancellationToken; -use nym_authenticator_client::{AuthClientMixnetListener, AuthenticatorClient}; +use nym_authenticator_client::{AuthenticatorClient, mixnet_listener}; use nym_bandwidth_controller::BandwidthTicketProvider; use nym_credentials_interface::TicketType; use nym_ip_packet_client::IprClientConnect; @@ -85,8 +85,11 @@ impl RegistrationClient { // Start the auth client mixnet listener, which will listen for incoming messages from the // mixnet and rebroadcast them to the auth clients. - let mixnet_listener = - AuthClientMixnetListener::new(self.mixnet_client, self.cancel_token.clone()).start(); + let mixnet_listener = mixnet_listener::spawn( + self.mixnet_client, + // todo: are we sure we want to clone the cancel token here? + self.cancel_token.clone(), + ); let mut entry_auth_client = AuthenticatorClient::new( mixnet_listener.subscribe(), diff --git a/nym-registration-client/src/types.rs b/nym-registration-client/src/types.rs index d7f7ad80be..14588c3b3b 100644 --- a/nym-registration-client/src/types.rs +++ b/nym-registration-client/src/types.rs @@ -1,7 +1,9 @@ // Copyright 2025 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use nym_authenticator_client::{AuthClientMixnetListenerHandle, AuthenticatorClient}; +use nym_authenticator_client::{ + AuthenticatorClient, mixnet_listener::AuthClientMixnetListenerHandle, +}; use nym_bandwidth_controller::BandwidthTicketProvider; use nym_registration_common::{AssignedAddresses, GatewayData}; use nym_sdk::mixnet::MixnetClient;