diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d6d57cd6b..770201d659 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,7 @@ Post 1.0.0 release, the changelog format is based on [Keep a Changelog](https:// ### Fixed - gateway-client: fix decrypting stored messages on reconnect ([#1786]) +- socks5-client: fix shutting down all tasks if anyone of them panics or errors out ([#1805]) [#1678]: https://github.com/nymtech/nym/pull/1678 [#1708]: https://github.com/nymtech/nym/pull/1708 @@ -28,6 +29,7 @@ Post 1.0.0 release, the changelog format is based on [Keep a Changelog](https:// [#1747]: https://github.com/nymtech/nym/pull/1747 [#1783]: https://github.com/nymtech/nym/pull/1783 [#1786]: https://github.com/nymtech/nym/pull/1786 +[#1805]: https://github.com/nymtech/nym/pull/1805 ## [v1.1.0](https://github.com/nymtech/nym/tree/v1.1.0) (2022-11-09) diff --git a/Cargo.lock b/Cargo.lock index 4303004d95..0ad2f180df 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5693,7 +5693,9 @@ checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" name = "task" version = "0.1.0" dependencies = [ + "futures", "log", + "thiserror", "tokio", ] @@ -5833,18 +5835,18 @@ checksum = "b1141d4d61095b28419e22cb0bbf02755f5e54e0526f97f1e3d1d160e60885fb" [[package]] name = "thiserror" -version = "1.0.35" +version = "1.0.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c53f98874615aea268107765aa1ed8f6116782501d18e53d08b471733bea6c85" +checksum = "10deb33631e3c9018b9baf9dcbbc4f737320d2b576bac10f6aefa048fa407e3e" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.35" +version = "1.0.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8b463991b4eab2d801e724172285ec4195c650e8ec79b149e6c2a8e6dd3f783" +checksum = "982d17546b47146b28f7c22e3d08465f6b8903d0ea13c1660d9d84a6e7adcdbb" dependencies = [ "proc-macro2", "quote", diff --git a/clients/client-core/src/client/cover_traffic_stream.rs b/clients/client-core/src/client/cover_traffic_stream.rs index ee3b7851cb..228234663a 100644 --- a/clients/client-core/src/client/cover_traffic_stream.rs +++ b/clients/client-core/src/client/cover_traffic_stream.rs @@ -228,7 +228,9 @@ impl LoopCoverTrafficStream { } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("LoopCoverTrafficStream: Exiting"); }) } diff --git a/clients/client-core/src/client/mix_traffic.rs b/clients/client-core/src/client/mix_traffic.rs index d5702cc0a8..06846316cc 100644 --- a/clients/client-core/src/client/mix_traffic.rs +++ b/clients/client-core/src/client/mix_traffic.rs @@ -69,6 +69,8 @@ impl MixTrafficController { #[cfg(not(target_arch = "wasm32"))] pub fn start_with_shutdown(mut self, mut shutdown: task::ShutdownListener) { + use std::time::Duration; + spawn_future(async move { debug!("Started MixTrafficController with graceful shutdown support"); @@ -88,7 +90,9 @@ impl MixTrafficController { } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("MixTrafficController: Exiting"); }) } diff --git a/clients/client-core/src/client/mod.rs b/clients/client-core/src/client/mod.rs index bce1a089dc..309be529f1 100644 --- a/clients/client-core/src/client/mod.rs +++ b/clients/client-core/src/client/mod.rs @@ -1,5 +1,3 @@ -use std::sync::atomic::AtomicBool; - pub mod cover_traffic_stream; pub mod inbound_messages; pub mod key_manager; @@ -9,10 +7,3 @@ pub mod received_buffer; #[cfg(feature = "reply-surb")] pub mod reply_key_storage; pub mod topology_control; - -// This is *NOT* used to signal shutdown. -// It's critical that we don't have any tasks finishing early, this is an additional safety check -// that tasks exiting are doing so because shutdown has been signalled, and no other reason. -// In particular for tasks that rely on their associated channel being closed to signal shutdown, -// and don't have access to a shutdown listener channel. -pub static SHUTDOWN_HAS_BEEN_SIGNALLED: AtomicBool = AtomicBool::new(false); diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs index 9b285d8444..dff476ad3c 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/acknowledgement_listener.rs @@ -72,6 +72,8 @@ impl AcknowledgementListener { #[cfg(not(target_arch = "wasm32"))] pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + use std::time::Duration; + debug!("Started AcknowledgementListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -88,7 +90,9 @@ impl AcknowledgementListener { } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("AcknowledgementListener: Exiting"); } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs index 2b21639230..822cea0a2f 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/action_controller.rs @@ -272,7 +272,9 @@ impl ActionController { } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("ActionController: Exiting"); } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs index 0738d0aa80..c58967104c 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs @@ -196,6 +196,8 @@ where #[cfg(not(target_arch = "wasm32"))] pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + use std::time::Duration; + debug!("Started InputMessageListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -214,7 +216,9 @@ where } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("InputMessageListener: Exiting"); } diff --git a/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs b/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs index 4b0bc3cc99..d0fd805d22 100644 --- a/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs +++ b/clients/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs @@ -127,6 +127,8 @@ where #[cfg(not(target_arch = "wasm32"))] pub(super) async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + use std::time::Duration; + debug!("Started RetransmissionRequestListener with graceful shutdown support"); while !shutdown.is_shutdown() { @@ -143,7 +145,9 @@ where } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("RetransmissionRequestListener: Exiting"); } diff --git a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs index 503203922c..b10e0f8b96 100644 --- a/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs +++ b/clients/client-core/src/client/real_messages_control/real_traffic_stream.rs @@ -514,7 +514,9 @@ where } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("OutQueueControl: Exiting"); } diff --git a/clients/client-core/src/client/received_buffer.rs b/clients/client-core/src/client/received_buffer.rs index fa61b9b9a0..7e5a0d7db8 100644 --- a/clients/client-core/src/client/received_buffer.rs +++ b/clients/client-core/src/client/received_buffer.rs @@ -322,6 +322,8 @@ impl RequestReceiver { #[cfg(not(target_arch = "wasm32"))] async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + use std::time::Duration; + debug!("Started RequestReceiver with graceful shutdown support"); while !shutdown.is_shutdown() { tokio::select! { @@ -340,7 +342,9 @@ impl RequestReceiver { }, } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("RequestReceiver: Exiting"); } @@ -372,6 +376,8 @@ impl FragmentedMessageReceiver { #[cfg(not(target_arch = "wasm32"))] async fn run_with_shutdown(&mut self, mut shutdown: task::ShutdownListener) { + use std::time::Duration; + debug!("Started FragmentedMessageReceiver with graceful shutdown support"); while !shutdown.is_shutdown() { tokio::select! { @@ -389,7 +395,9 @@ impl FragmentedMessageReceiver { } } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("FragmentedMessageReceiver: Exiting"); } diff --git a/clients/client-core/src/client/topology_control.rs b/clients/client-core/src/client/topology_control.rs index afedc7ad61..b685cfb8b4 100644 --- a/clients/client-core/src/client/topology_control.rs +++ b/clients/client-core/src/client/topology_control.rs @@ -311,7 +311,9 @@ impl TopologyRefresher { }, } } - assert!(shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("TopologyRefresher: Exiting"); }) } diff --git a/clients/client-core/src/error.rs b/clients/client-core/src/error.rs index 24e846087f..0313db076b 100644 --- a/clients/client-core/src/error.rs +++ b/clients/client-core/src/error.rs @@ -26,4 +26,7 @@ pub enum ClientCoreError { CouldNotLoadExistingGatewayConfiguration(std::io::Error), #[error("The current network topology seem to be insufficient to route any packets through")] InsufficientNetworkTopology, + + #[error("Unexpected exit")] + UnexpectedExit, } diff --git a/clients/socks5/src/client/mod.rs b/clients/socks5/src/client/mod.rs index c6b734a749..399254c0d1 100644 --- a/clients/socks5/src/client/mod.rs +++ b/clients/socks5/src/client/mod.rs @@ -1,7 +1,7 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use std::sync::atomic::Ordering; +use std::error::Error; use crate::client::config::Config; use crate::error::Socks5ClientError; @@ -38,7 +38,8 @@ use log::*; use nymsphinx::addressing::clients::Recipient; use nymsphinx::addressing::nodes::NodeIdentity; use tap::TapFallible; -use task::{wait_for_signal, ShutdownListener, ShutdownNotifier}; +use task::signal::wait_for_signal_and_error; +use task::{ShutdownListener, ShutdownNotifier}; pub mod config; @@ -299,7 +300,7 @@ impl NymClient { msg_input: InputMessageSender, client_connection_tx: ConnectionCommandSender, lane_queue_lengths: LaneQueueLengths, - shutdown: ShutdownListener, + mut shutdown: ShutdownListener, ) { info!("Starting socks5 listener..."); let auth_methods = vec![AuthenticationMethods::NoAuth as u8]; @@ -312,38 +313,57 @@ impl NymClient { self.config.get_provider_mix_address(), self.as_mix_recipient(), lane_queue_lengths, - shutdown, + shutdown.clone(), ); tokio::spawn(async move { - sphinx_socks + // Ideally we should have a fully fledged task manager to check for errors in all + // tasks. + // However, pragmatically, we start out by at least reporting errors for some of the + // tasks that interact with the outside world and can fail in normal operation, such as + // network issues. + // TODO: replace this by a generic solution, such as a task manager that stores all + // JoinHandles of all spawned tasks. + if let Err(res) = sphinx_socks .serve(msg_input, buffer_requester, client_connection_tx) .await + { + shutdown.send_we_stopped(Box::new(res)); + } }); } /// blocking version of `start` method. Will run forever (or until SIGINT is sent) - pub async fn run_forever(&mut self) -> Result<(), Socks5ClientError> { - let mut shutdown = self.start().await?; - wait_for_signal().await; + pub async fn run_forever(&mut self) -> Result<(), Box> { + let mut shutdown = self + .start() + .await + .map_err(|err| Box::new(err) as Box)?; + + let res = wait_for_signal_and_error(&mut shutdown).await; log::info!("Sending shutdown"); - client_core::client::SHUTDOWN_HAS_BEEN_SIGNALLED.store(true, Ordering::Relaxed); shutdown.signal_shutdown().ok(); log::info!("Waiting for tasks to finish... (Press ctrl-c to force)"); shutdown.wait_for_shutdown().await; log::info!("Stopping nym-socks5-client"); - Ok(()) + res } // Variant of `run_forever` that listends for remote control messages pub async fn run_and_listen( &mut self, mut receiver: Socks5ControlMessageReceiver, - ) -> Result<(), Socks5ClientError> { - let mut shutdown = self.start().await?; - tokio::select! { + ) -> Result<(), Box> { + // Start the main task + let mut shutdown = self + .start() + .await + .map_err(|err| Box::new(err) as Box)?; + + let res = tokio::select! { + biased; message = receiver.next() => { log::debug!("Received message: {:?}", message); match message { @@ -354,21 +374,26 @@ impl NymClient { log::info!("Channel closed, stopping"); } } + Ok(()) + } + Some(msg) = shutdown.wait_for_error() => { + log::info!("Task error: {:?}", msg); + Err(msg) } _ = tokio::signal::ctrl_c() => { log::info!("Received SIGINT"); + Ok(()) }, - } + }; log::info!("Sending shutdown"); - client_core::client::SHUTDOWN_HAS_BEEN_SIGNALLED.store(true, Ordering::Relaxed); shutdown.signal_shutdown().ok(); log::info!("Waiting for tasks to finish... (Press ctrl-c to force)"); shutdown.wait_for_shutdown().await; log::info!("Stopping nym-socks5-client"); - Ok(()) + res } pub async fn start(&mut self) -> Result { diff --git a/clients/socks5/src/commands/mod.rs b/clients/socks5/src/commands/mod.rs index b069c13542..e1d42e0fb1 100644 --- a/clients/socks5/src/commands/mod.rs +++ b/clients/socks5/src/commands/mod.rs @@ -1,8 +1,9 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::error::Error; + use crate::client::config::Config; -use crate::error::Socks5ClientError; use clap::CommandFactory; use clap::{Parser, Subcommand}; use completions::{fig_generate, ArgShell}; @@ -87,7 +88,7 @@ pub(crate) struct OverrideConfig { enabled_credentials_mode: bool, } -pub(crate) async fn execute(args: &Cli) -> Result<(), Socks5ClientError> { +pub(crate) async fn execute(args: &Cli) -> Result<(), Box> { let bin_name = "nym-socks5-client"; match &args.command { diff --git a/clients/socks5/src/commands/run.rs b/clients/socks5/src/commands/run.rs index 406ae72d94..d74e6678fe 100644 --- a/clients/socks5/src/commands/run.rs +++ b/clients/socks5/src/commands/run.rs @@ -86,14 +86,16 @@ fn version_check(cfg: &Config) -> bool { } } -pub(crate) async fn execute(args: &Run) -> Result<(), Socks5ClientError> { +pub(crate) async fn execute(args: &Run) -> Result<(), Box> { let id = &args.id; let mut config = match Config::load_from_file(Some(id)) { Ok(cfg) => cfg, Err(err) => { error!("Failed to load config for {}. Are you sure you have run `init` before? (Error was: {})", id, err); - return Err(Socks5ClientError::FailedToLoadConfig(id.to_string())); + return Err(Box::new(Socks5ClientError::FailedToLoadConfig( + id.to_string(), + ))); } }; @@ -102,7 +104,7 @@ pub(crate) async fn execute(args: &Run) -> Result<(), Socks5ClientError> { if !version_check(&config) { error!("failed the local version check"); - return Err(Socks5ClientError::FailedLocalVersionCheck); + return Err(Box::new(Socks5ClientError::FailedLocalVersionCheck)); } NymClient::new(config).run_forever().await diff --git a/clients/socks5/src/error.rs b/clients/socks5/src/error.rs index 4ebca0e2e5..6b67e6a339 100644 --- a/clients/socks5/src/error.rs +++ b/clients/socks5/src/error.rs @@ -3,6 +3,8 @@ use crypto::asymmetric::identity::Ed25519RecoveryError; use gateway_client::error::GatewayClientError; use validator_client::ValidatorClientError; +use crate::socks::types::SocksProxyError; + #[derive(thiserror::Error, Debug)] pub enum Socks5ClientError { #[error("I/O error: {0}")] @@ -18,8 +20,13 @@ pub enum Socks5ClientError { #[error("Reply key storage error: {0}")] ReplyKeyStorageError(#[from] ReplyKeyStorageError), + #[error("SOCKS proxy error")] + SocksProxyError(SocksProxyError), + #[error("Failed to load config for: {0}")] FailedToLoadConfig(String), #[error("Failed local version check, client and config mismatch")] FailedLocalVersionCheck, + #[error("Fail to bind address")] + FailToBindAddress, } diff --git a/clients/socks5/src/main.rs b/clients/socks5/src/main.rs index d7a0fdf01f..d41e5e4bd8 100644 --- a/clients/socks5/src/main.rs +++ b/clients/socks5/src/main.rs @@ -1,8 +1,9 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 +use std::error::Error; + use clap::{crate_version, Parser}; -use error::Socks5ClientError; use logging::setup_logging; use network_defaults::setup_env; @@ -12,7 +13,7 @@ pub mod error; pub mod socks; #[tokio::main] -async fn main() -> Result<(), Socks5ClientError> { +async fn main() -> Result<(), Box> { setup_logging(); println!("{}", banner()); diff --git a/clients/socks5/src/socks/client.rs b/clients/socks5/src/socks/client.rs index df73d9ee7f..3e3c24347c 100644 --- a/clients/socks5/src/socks/client.rs +++ b/clients/socks5/src/socks/client.rs @@ -202,6 +202,7 @@ impl SocksClient { pub async fn shutdown(&mut self) -> Result<(), SocksProxyError> { info!("client is shutting down its TCP stream"); self.stream.shutdown().await?; + self.shutdown_listener.mark_as_success(); Ok(()) } @@ -318,6 +319,7 @@ impl SocksClient { SocksCommand::UdpAssociate => unimplemented!(), // not handled }; + self.shutdown_listener.mark_as_success(); Ok(()) } diff --git a/clients/socks5/src/socks/mixnet_responses.rs b/clients/socks5/src/socks/mixnet_responses.rs index 760ec36d24..fa9ad1abe8 100644 --- a/clients/socks5/src/socks/mixnet_responses.rs +++ b/clients/socks5/src/socks/mixnet_responses.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use futures::channel::mpsc; use futures::StreamExt; use log::*; @@ -103,7 +105,9 @@ impl MixnetResponseListener { } } } - assert!(self.shutdown.is_shutdown_poll()); + tokio::time::timeout(Duration::from_secs(5), self.shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); log::debug!("MixnetResponseListener: Exiting"); } } diff --git a/clients/socks5/src/socks/server.rs b/clients/socks5/src/socks/server.rs index f61d6052a0..551f82b788 100644 --- a/clients/socks5/src/socks/server.rs +++ b/clients/socks5/src/socks/server.rs @@ -1,9 +1,8 @@ +use crate::error::Socks5ClientError; + use super::authentication::Authenticator; use super::client::SocksClient; -use super::{ - mixnet_responses::MixnetResponseListener, - types::{ResponseCode, SocksProxyError}, -}; +use super::{mixnet_responses::MixnetResponseListener, types::ResponseCode}; use client_connections::{ConnectionCommandSender, LaneQueueLengths}; use client_core::client::{ inbound_messages::InputMessageSender, received_buffer::ReceivedBufferRequestSender, @@ -12,6 +11,7 @@ use log::*; use nymsphinx::addressing::clients::Recipient; use proxy_helpers::connection_controller::{BroadcastActiveConnections, Controller}; use std::net::SocketAddr; +use tap::TapFallible; use task::ShutdownListener; use tokio::net::TcpListener; @@ -56,8 +56,10 @@ impl SphinxSocksServer { input_sender: InputMessageSender, buffer_requester: ReceivedBufferRequestSender, client_connection_tx: ConnectionCommandSender, - ) -> Result<(), SocksProxyError> { - let listener = TcpListener::bind(self.listening_address).await.unwrap(); + ) -> Result<(), Socks5ClientError> { + let listener = TcpListener::bind(self.listening_address) + .await + .tap_err(|err| log::error!("Failed to bind to address: {err}"))?; info!("Serving Connections..."); // controller for managing all active connections diff --git a/common/client-libs/gateway-client/src/client.rs b/common/client-libs/gateway-client/src/client.rs index bac92ec092..4b34f2239b 100644 --- a/common/client-libs/gateway-client/src/client.rs +++ b/common/client-libs/gateway-client/src/client.rs @@ -307,6 +307,8 @@ impl GatewayClient { let m_shutdown = self.shutdown.clone(); async { if let Some(mut s) = m_shutdown { + // TODO: fix this by marking as success _after_ the select + s.mark_as_success(); s.recv().await } else { std::future::pending::<()>().await diff --git a/common/socks5/proxy-helpers/src/connection_controller.rs b/common/socks5/proxy-helpers/src/connection_controller.rs index 12fdba7dbb..f0f552d83c 100644 --- a/common/socks5/proxy-helpers/src/connection_controller.rs +++ b/common/socks5/proxy-helpers/src/connection_controller.rs @@ -254,6 +254,9 @@ impl Controller { }, } } + tokio::time::timeout(Duration::from_secs(5), self.shutdown.recv()) + .await + .expect("Task stopped without shutdown called"); assert!(self.shutdown.is_shutdown_poll()); log::debug!("SOCKS5 Controller: Exiting"); } diff --git a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs index 12caff2071..f2ebc2cd9b 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/inbound.rs @@ -208,5 +208,6 @@ where trace!("{} - inbound closed", connection_id); shutdown_notify.notify_one(); + shutdown_listener.mark_as_success(); reader } diff --git a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs index fbd8cd21b3..defcbf4f5f 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/mod.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/mod.rs @@ -134,6 +134,7 @@ where } pub fn into_inner(mut self) -> (TcpStream, ConnectionReceiver) { + self.shutdown_listener.mark_as_success(); ( self.socket.take().unwrap(), self.mix_receiver.take().unwrap(), diff --git a/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs b/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs index 8fad665785..78db73afb6 100644 --- a/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs +++ b/common/socks5/proxy-helpers/src/proxy_runner/outbound.rs @@ -90,5 +90,6 @@ pub(super) async fn run_outbound( trace!("{} - outbound closed", connection_id); shutdown_notify.notify_one(); + shutdown_listener.mark_as_success(); (writer, mix_receiver) } diff --git a/common/task/Cargo.toml b/common/task/Cargo.toml index 4c0c088d67..f074cd2abd 100644 --- a/common/task/Cargo.toml +++ b/common/task/Cargo.toml @@ -6,7 +6,9 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] +futures = "0.3" log = "0.4" +thiserror = "1.0.37" tokio = { version = "1.21.2", features = ["macros", "signal", "time", "sync"] } [dev-dependencies] diff --git a/common/task/src/shutdown.rs b/common/task/src/shutdown.rs index 21ea12569d..5d4f2c77bb 100644 --- a/common/task/src/shutdown.rs +++ b/common/task/src/shutdown.rs @@ -1,27 +1,62 @@ // Copyright 2022 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use std::time::Duration; +use std::{error::Error, time::Duration}; -use tokio::sync::watch::{self, error::SendError}; +use futures::FutureExt; +use tokio::{ + sync::{ + mpsc, + watch::{self, error::SendError}, + }, + time::sleep, +}; const DEFAULT_SHUTDOWN_TIMER_SECS: u64 = 5; +pub(crate) type SentError = Box; +type ErrorSender = mpsc::UnboundedSender; +type ErrorReceiver = mpsc::UnboundedReceiver; + +#[derive(thiserror::Error, Debug)] +enum TaskError { + #[error("Task halted unexpectedly")] + UnexpectedHalt, +} + /// Used to notify other tasks to gracefully shutdown #[derive(Debug)] pub struct ShutdownNotifier { + // These channels have the dual purpose of signalling it's time to shutdown, but also to keep + // track of which tasks we are still waiting for. notify_tx: watch::Sender<()>, notify_rx: Option>, shutdown_timer_secs: u64, + + // If any task failed, it needs to report separately + task_return_error_tx: ErrorSender, + task_return_error_rx: Option, + + // Also signal when the notifier is dropped, in case the task exits unexpectedly. + // Why are we not reusing the return error channel? Well, let me tell you kids, it's because I + // didn't manage to reliably get the explicitly sent error (and not the error sent during drop) + task_drop_tx: ErrorSender, + task_drop_rx: Option, } impl Default for ShutdownNotifier { fn default() -> Self { let (notify_tx, notify_rx) = watch::channel(()); + let (task_halt_tx, task_halt_rx) = mpsc::unbounded_channel(); + let (task_drop_tx, task_drop_rx) = mpsc::unbounded_channel(); Self { notify_tx, notify_rx: Some(notify_rx), shutdown_timer_secs: DEFAULT_SHUTDOWN_TIMER_SECS, + task_return_error_tx: task_halt_tx, + task_return_error_rx: Some(task_halt_rx), + task_drop_tx, + task_drop_rx: Some(task_drop_rx), } } } @@ -40,6 +75,8 @@ impl ShutdownNotifier { .as_ref() .expect("Unable to subscribe to shutdown notifier that is already shutdown") .clone(), + self.task_return_error_tx.clone(), + self.task_drop_tx.clone(), ) } @@ -47,7 +84,31 @@ impl ShutdownNotifier { self.notify_tx.send(()) } + pub async fn wait_for_error(&mut self) -> Option { + let mut error_rx = self + .task_return_error_rx + .take() + .expect("Unable to wait for error: attempt to wait twice?"); + let mut drop_rx = self + .task_drop_rx + .take() + .expect("Unable to wait for error: attempt to wait twice?"); + + // During an error we are likely like to be swamped with drop notifications as well, this + // is a crude way to give priority to real errors (if there are any). + let drop_rx = drop_rx.recv().then(|msg| async move { + sleep(Duration::from_millis(50)).await; + msg + }); + + tokio::select! { + msg = error_rx.recv() => msg, + msg = drop_rx => msg + } + } + pub async fn wait_for_shutdown(&mut self) { + log::info!("Waiting for shutdown"); if let Some(notify_rx) = self.notify_rx.take() { drop(notify_rx); } @@ -56,7 +117,7 @@ impl ShutdownNotifier { _ = self.notify_tx.closed() => { log::info!("All registered tasks succesfully shutdown"); }, - _ = tokio::signal::ctrl_c() => { + _ = tokio::signal::ctrl_c() => { log::info!("Forcing shutdown"); } _ = tokio::time::sleep(Duration::from_secs(self.shutdown_timer_secs)) => { @@ -69,15 +130,36 @@ impl ShutdownNotifier { /// Listen for shutdown notifications #[derive(Clone, Debug)] pub struct ShutdownListener { + // If a shutdown notification has been registered shutdown: bool, + + // Listen for shutdown notifications, as well as a mechanism to report back that we have + // finished (the receiver is closed). notify: watch::Receiver<()>, + + // Send back error if we stopped + return_error: ErrorSender, + + // Also notify if we dropped without shutdown being registered + drop_error: ErrorSender, + + // Sometimes it's necessary to clone and drop the shutdown listener during normal operation, + // for those situations we need to explicitly not drop (and trigger shutdown). + set_not_drop: bool, } impl ShutdownListener { - fn new(notify: watch::Receiver<()>) -> ShutdownListener { + fn new( + notify: watch::Receiver<()>, + return_error: ErrorSender, + drop_error: ErrorSender, + ) -> ShutdownListener { ShutdownListener { shutdown: false, notify, + return_error, + drop_error, + set_not_drop: false, } } @@ -111,6 +193,29 @@ impl ShutdownListener { } } } + + pub fn send_we_stopped(&mut self, err: SentError) { + log::trace!("Notifying we stopped: {:?}", err); + if self.return_error.send(err).is_err() { + log::error!("Failed to send back error message"); + } + } + + pub fn mark_as_success(&mut self) { + self.set_not_drop = true; + } +} + +impl Drop for ShutdownListener { + fn drop(&mut self) { + if !self.set_not_drop && !self.is_shutdown_poll() { + log::trace!("Notifying stop on unexpected drop"); + // If we can't send, well then there is not much to do + self.drop_error + .send(Box::new(TaskError::UnexpectedHalt)) + .ok(); + } + } } #[cfg(test)] diff --git a/common/task/src/signal.rs b/common/task/src/signal.rs index 40ca122848..a580c5e293 100644 --- a/common/task/src/signal.rs +++ b/common/task/src/signal.rs @@ -17,6 +17,35 @@ pub async fn wait_for_signal() { } } +#[cfg(unix)] +pub async fn wait_for_signal_and_error( + shutdown: &mut crate::ShutdownNotifier, +) -> Result<(), crate::shutdown::SentError> { + use tokio::signal::unix::{signal, SignalKind}; + + let mut sigterm = signal(SignalKind::terminate()).expect("Failed to setup SIGTERM channel"); + let mut sigquit = signal(SignalKind::quit()).expect("Failed to setup SIGQUIT channel"); + + tokio::select! { + _ = tokio::signal::ctrl_c() => { + log::info!("Received SIGINT"); + Ok(()) + }, + _ = sigterm.recv() => { + log::info!("Received SIGTERM"); + Ok(()) + } + _ = sigquit.recv() => { + log::info!("Received SIGQUIT"); + Ok(()) + } + Some(msg) = shutdown.wait_for_error() => { + log::info!("Task error: {:?}", msg); + Err(msg) + } + } +} + #[cfg(not(unix))] pub async fn wait_for_signal() { tokio::select! { diff --git a/nym-connect/Cargo.lock b/nym-connect/Cargo.lock index cab899eacf..22331cf94d 100644 --- a/nym-connect/Cargo.lock +++ b/nym-connect/Cargo.lock @@ -5459,7 +5459,9 @@ dependencies = [ name = "task" version = "0.1.0" dependencies = [ + "futures", "log", + "thiserror", "tokio", ] @@ -5784,18 +5786,18 @@ checksum = "8eaa81235c7058867fa8c0e7314f33dcce9c215f535d1913822a2b3f5e289f3c" [[package]] name = "thiserror" -version = "1.0.35" +version = "1.0.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c53f98874615aea268107765aa1ed8f6116782501d18e53d08b471733bea6c85" +checksum = "10deb33631e3c9018b9baf9dcbbc4f737320d2b576bac10f6aefa048fa407e3e" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "1.0.35" +version = "1.0.37" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8b463991b4eab2d801e724172285ec4195c650e8ec79b149e6c2a8e6dd3f783" +checksum = "982d17546b47146b28f7c22e3d08465f6b8903d0ea13c1660d9d84a6e7adcdbb" dependencies = [ "proc-macro2", "quote", diff --git a/nym-connect/Cargo.toml b/nym-connect/Cargo.toml index b9fb14e41b..3ee3f20be3 100644 --- a/nym-connect/Cargo.toml +++ b/nym-connect/Cargo.toml @@ -1,2 +1,2 @@ [workspace] -members = ["src-tauri"] \ No newline at end of file +members = ["src-tauri"] diff --git a/nym-connect/src-tauri/src/tasks.rs b/nym-connect/src-tauri/src/tasks.rs index 751f7c5566..85da77dfee 100644 --- a/nym-connect/src-tauri/src/tasks.rs +++ b/nym-connect/src-tauri/src/tasks.rs @@ -50,7 +50,7 @@ pub fn start_nym_socks5_client( .block_on(async move { socks5_client.run_and_listen(socks5_ctrl_rx).await }); if let Err(err) = result { - log::error!("SOCKS5 proxy failed to start: {err}"); + log::error!("SOCKS5 proxy failed: {err}"); socks5_status_tx .send(Socks5StatusMessage::FailedToStart) .expect("Failed to send status message back to main task");