diff --git a/nym-authenticator-client/src/mixnet_listener.rs b/nym-authenticator-client/src/mixnet_listener.rs index 9cb6828e83..5984f6f870 100644 --- a/nym-authenticator-client/src/mixnet_listener.rs +++ b/nym-authenticator-client/src/mixnet_listener.rs @@ -50,7 +50,7 @@ impl AuthClientMixnetListener { } } - async fn run(mut self) -> Self { + async fn run(mut self) { let mixnet_cancel_token = self.mixnet_client.cancellation_token(); self.shutdown_token.run_until_cancelled(async { loop { @@ -95,12 +95,8 @@ impl AuthClientMixnetListener { 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() { + tracing::debug!("AuthClientMixnetListener: Disconnect mixnet client"); + if !mixnet_cancel_token.is_cancelled() { self.mixnet_client.disconnect().await; } } @@ -128,7 +124,7 @@ pub struct AuthClientMixnetListenerHandle { message_sender: MixnetMessageInputSender, cancellation_token: CancellationToken, mixnet_cancellation_token: CancellationToken, - handle: JoinHandle, + handle: JoinHandle<()>, } impl AuthClientMixnetListenerHandle { @@ -148,13 +144,8 @@ impl AuthClientMixnetListenerHandle { // 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}"); - } + if let Err(e) = self.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..cac1225a8c 100644 --- a/nym-gateway-probe/src/lib.rs +++ b/nym-gateway-probe/src/lib.rs @@ -765,7 +765,7 @@ async fn connect_exit( ); // The IPR supports cancellation, but it's unused in the gateway probe let cancel_token = CancellationToken::new(); - let mut ipr_client = IprClientConnect::new(mixnet_client, cancel_token).await; + let mut ipr_client = IprClientConnect::new(mixnet_client, cancel_token); let maybe_ip_pair = ipr_client.connect(exit_router_address).await; let mixnet_client = ipr_client.into_mixnet_client(); diff --git a/nym-ip-packet-client/src/connect.rs b/nym-ip-packet-client/src/connect.rs index 5a37cf6d82..bf679f67cf 100644 --- a/nym-ip-packet-client/src/connect.rs +++ b/nym-ip-packet-client/src/connect.rs @@ -43,7 +43,7 @@ pub struct IprClientConnect { } impl IprClientConnect { - pub async fn new(mixnet_client: MixnetClient, cancel_token: CancellationToken) -> Self { + pub fn new(mixnet_client: MixnetClient, cancel_token: CancellationToken) -> Self { Self { mixnet_client, connected: ConnectionState::Disconnected, diff --git a/nym-registration-client/src/lib.rs b/nym-registration-client/src/lib.rs index 0ba6fea90e..ece0e6fd64 100644 --- a/nym-registration-client/src/lib.rs +++ b/nym-registration-client/src/lib.rs @@ -9,6 +9,7 @@ use nym_credentials_interface::TicketType; use nym_ip_packet_client::IprClientConnect; use nym_registration_common::AssignedAddresses; use nym_sdk::mixnet::{EventReceiver, MixnetClient, Recipient}; +use tracing::debug; use crate::config::RegistrationClientConfig; @@ -34,23 +35,49 @@ pub struct RegistrationClient { event_rx: EventReceiver, } +// Bundle of an actual error and the underlying mixnet client so it can be shutdown correctly if needed +struct RegistrationError { + mixnet_client: Option, + source: crate::RegistrationClientError, +} + impl RegistrationClient { - async fn register_mix_exit(self) -> Result { + async fn register_mix_exit(self) -> Result { let entry_mixnet_gateway_ip = self.config.entry.node.ip_address; let exit_mixnet_gateway_ip = self.config.exit.node.ip_address; - let ipr_address = self.config.exit.node.ipr_address.ok_or( - RegistrationClientError::NoIpPacketRouterAddress { - node_id: self.config.exit.node.identity.to_base58_string(), - }, - )?; + let Some(ipr_address) = self.config.exit.node.ipr_address else { + return Err(RegistrationError { + mixnet_client: Some(self.mixnet_client), + source: RegistrationClientError::NoIpPacketRouterAddress { + node_id: self.config.exit.node.identity.to_base58_string(), + }, + }); + }; + let mut ipr_client = - IprClientConnect::new(self.mixnet_client, self.cancel_token.clone()).await; - let interface_addresses = ipr_client - .connect(ipr_address) + IprClientConnect::new(self.mixnet_client, self.cancel_token.child_token()); + + let interface_addresses = match self + .cancel_token + .run_until_cancelled(ipr_client.connect(ipr_address)) .await - .map_err(RegistrationClientError::ConnectToIpPacketRouter)?; + { + Some(Ok(addr)) => addr, + Some(Err(e)) => { + return Err(RegistrationError { + mixnet_client: Some(ipr_client.into_mixnet_client()), + source: RegistrationClientError::ConnectToIpPacketRouter(e), + }); + } + None => { + return Err(RegistrationError { + mixnet_client: Some(ipr_client.into_mixnet_client()), + source: RegistrationClientError::Cancelled, + }); + } + }; Ok(RegistrationResult::Mixnet(Box::new( MixnetRegistrationResult { @@ -67,18 +94,24 @@ impl RegistrationClient { ))) } - async fn register_wg(self) -> Result { - let entry_auth_address = self.config.entry.node.authenticator_address.ok_or( - RegistrationClientError::AuthenticationNotPossible { - node_id: self.config.entry.node.identity.to_base58_string(), - }, - )?; + async fn register_wg(self) -> Result { + let Some(entry_auth_address) = self.config.entry.node.authenticator_address else { + return Err(RegistrationError { + mixnet_client: Some(self.mixnet_client), + source: RegistrationClientError::AuthenticationNotPossible { + node_id: self.config.entry.node.identity.to_base58_string(), + }, + }); + }; - let exit_auth_address = self.config.exit.node.authenticator_address.ok_or( - RegistrationClientError::AuthenticationNotPossible { - node_id: self.config.exit.node.identity.to_base58_string(), - }, - )?; + let Some(exit_auth_address) = self.config.exit.node.authenticator_address else { + return Err(RegistrationError { + mixnet_client: Some(self.mixnet_client), + source: RegistrationClientError::AuthenticationNotPossible { + node_id: self.config.exit.node.identity.to_base58_string(), + }, + }); + }; let entry_version = self.config.entry.node.version; tracing::debug!("Entry gateway version: {entry_version}"); @@ -87,8 +120,10 @@ impl RegistrationClient { // Start the auth client mixnet listener, which will listen for incoming messages from the // mixnet and rebroadcast them to the auth clients. + // From this point on, we don't need to care about the mixnet client anymore let mixnet_listener = - AuthClientMixnetListener::new(self.mixnet_client, self.cancel_token.clone()).start(); + AuthClientMixnetListener::new(self.mixnet_client, self.cancel_token.child_token()) + .start(); let mut entry_auth_client = AuthenticatorClient::new( mixnet_listener.subscribe(), @@ -115,24 +150,33 @@ impl RegistrationClient { let exit_fut = exit_auth_client .register_wireguard(&*self.bandwidth_controller, TicketType::V1WireguardExit); - let (entry, exit) = Box::pin(async { tokio::join!(entry_fut, exit_fut) }).await; + let (entry, exit) = Box::pin( + self.cancel_token + .run_until_cancelled(async { tokio::join!(entry_fut, exit_fut) }), + ) + .await + .ok_or(RegistrationError { + mixnet_client: None, + source: RegistrationClientError::Cancelled, + })?; - let entry = - entry.map_err( - |source| RegistrationClientError::EntryGatewayRegisterWireguard { - gateway_id: self.config.entry.node.identity.to_base58_string(), - authenticator_address: Box::new(entry_auth_address), - source: Box::new(source), - }, - )?; - let exit = - exit.map_err( - |source| RegistrationClientError::ExitGatewayRegisterWireguard { - gateway_id: self.config.exit.node.identity.to_base58_string(), - authenticator_address: Box::new(exit_auth_address), - source: Box::new(source), - }, - )?; + let entry = entry.map_err(|source| RegistrationError { + mixnet_client: None, + source: RegistrationClientError::EntryGatewayRegisterWireguard { + gateway_id: self.config.entry.node.identity.to_base58_string(), + authenticator_address: Box::new(entry_auth_address), + source: Box::new(source), + }, + })?; + + let exit = exit.map_err(|source| RegistrationError { + mixnet_client: None, + source: RegistrationClientError::EntryGatewayRegisterWireguard { + gateway_id: self.config.exit.node.identity.to_base58_string(), + authenticator_address: Box::new(exit_auth_address), + source: Box::new(source), + }, + })?; Ok(RegistrationResult::Wireguard(Box::new( WireguardRegistrationResult { @@ -147,16 +191,23 @@ impl RegistrationClient { } pub async fn register(self) -> Result { - self.cancel_token - .clone() - .run_until_cancelled(async { - if self.config.two_hops { - self.register_wg().await - } else { - self.register_mix_exit().await + let registration_result = if self.config.two_hops { + self.register_wg().await + } else { + self.register_mix_exit().await + }; + + // If we failed to register, and we were the owner of the mixnet client, shut it down + match registration_result { + Ok(result) => Ok(result), + Err(error) => { + debug!("Registration failed"); + if let Some(mixnet_client) = error.mixnet_client { + debug!("Shutting down mixnet client"); + mixnet_client.disconnect().await; } - }) - .await - .ok_or(RegistrationClientError::Cancelled)? + Err(error.source) + } + } } }