From e0567dddf2fc425fc2497f8e5ffdb85abb2e72c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Thu, 8 Sep 2022 15:42:21 +0200 Subject: [PATCH] gateway-client: additional shutdown handling --- .../client-core/src/client/received_buffer.rs | 4 +- clients/client-core/src/init.rs | 1 + .../client-libs/gateway-client/src/client.rs | 28 ++++++++++---- .../client-libs/gateway-client/src/error.rs | 6 +++ .../gateway-client/src/packet_router.rs | 37 +++++++++++++------ .../gateway-client/src/socket_state.rs | 37 +++++++++++-------- common/task/src/shutdown.rs | 2 +- 7 files changed, 78 insertions(+), 37 deletions(-) diff --git a/clients/client-core/src/client/received_buffer.rs b/clients/client-core/src/client/received_buffer.rs index b116617f4e..567302121f 100644 --- a/clients/client-core/src/client/received_buffer.rs +++ b/clients/client-core/src/client/received_buffer.rs @@ -1,8 +1,8 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::client::SHUTDOWN_HAS_BEEN_SIGNALLED; use crate::client::reply_key_storage::ReplyKeyStorage; +use crate::client::SHUTDOWN_HAS_BEEN_SIGNALLED; use crypto::asymmetric::encryption; use crypto::symmetric::stream_cipher; use crypto::Digest; @@ -15,8 +15,8 @@ use nymsphinx::anonymous_replies::{encryption_key::EncryptionKeyDigest, SurbEncr use nymsphinx::params::{ReplySurbEncryptionAlgorithm, ReplySurbKeyDigestAlgorithm}; use nymsphinx::receiver::{MessageReceiver, MessageRecoveryError, ReconstructedMessage}; use std::collections::HashSet; -use std::sync::Arc; use std::sync::atomic::Ordering; +use std::sync::Arc; use task::ShutdownListener; use tokio::task::JoinHandle; diff --git a/clients/client-core/src/init.rs b/clients/client-core/src/init.rs index 66ae845b04..a2e3ed22f7 100644 --- a/clients/client-core/src/init.rs +++ b/clients/client-core/src/init.rs @@ -90,6 +90,7 @@ async fn register_with_gateway( gateway.owner.clone(), our_identity.clone(), timeout, + None, ); gateway_client .establish_connection() diff --git a/common/client-libs/gateway-client/src/client.rs b/common/client-libs/gateway-client/src/client.rs index 658a314535..71eb795e65 100644 --- a/common/client-libs/gateway-client/src/client.rs +++ b/common/client-libs/gateway-client/src/client.rs @@ -67,8 +67,7 @@ pub struct GatewayClient { /// Delay between each subsequent reconnection attempt. reconnection_backoff: Duration, - /// Listen to shutdown messages. This is an option since we don't require it when for example - /// when doing initial authentication. + /// Listen to shutdown messages. shutdown: Option, } @@ -97,7 +96,7 @@ impl GatewayClient { local_identity, shared_key, connection: SocketState::NotConnected, - packet_router: PacketRouter::new(ack_sender, mixnet_message_sender), + packet_router: PacketRouter::new(ack_sender, mixnet_message_sender, shutdown.clone()), response_timeout_duration, bandwidth_controller, should_reconnect_on_failure: true, @@ -130,6 +129,7 @@ impl GatewayClient { gateway_owner: String, local_identity: Arc, response_timeout_duration: Duration, + shutdown: Option, ) -> Self { use futures::channel::mpsc; @@ -137,7 +137,7 @@ impl GatewayClient { // perfectly fine here, because it's not meant to be used let (ack_tx, _) = mpsc::unbounded(); let (mix_tx, _) = mpsc::unbounded(); - let packet_router = PacketRouter::new(ack_tx, mix_tx); + let packet_router = PacketRouter::new(ack_tx, mix_tx, shutdown.clone()); GatewayClient { authenticated: false, @@ -155,7 +155,7 @@ impl GatewayClient { should_reconnect_on_failure: false, reconnection_attempts: DEFAULT_RECONNECTION_ATTEMPTS, reconnection_backoff: DEFAULT_RECONNECTION_BACKOFF, - shutdown: None, + shutdown, } } @@ -287,8 +287,20 @@ impl GatewayClient { let mut fused_timeout = timeout.fuse(); let mut fused_stream = conn.fuse(); + // Bit of an ugly workaround for selecting on an `Option`. The unwrap is lazy so we use + // this bool inside the `select` to guard against unwrapping in the `None` case. + let shutdown_is_some = self.shutdown.is_some(); + let shutdown_recv_lazy = async { self.shutdown.clone().unwrap().recv().await }; + tokio::pin!(shutdown_recv_lazy); + loop { - futures::select! { + tokio::select! { + biased; + _ = &mut shutdown_recv_lazy, if shutdown_is_some => { + log::trace!("GatewayClient control response: Received shutdown"); + log::debug!("GatewayClient control response: Exiting"); + break Err(GatewayClientError::ConnectionClosedGatewayShutdown); + } _ = &mut fused_timeout => { break Err(GatewayClientError::Timeout); } @@ -299,7 +311,9 @@ impl GatewayClient { }; match ws_msg { Message::Binary(bin_msg) => { - self.packet_router.route_received(vec![bin_msg]); + if let Err(err) = self.packet_router.route_received(vec![bin_msg]) { + log::warn!("Route received failed: {:?}", err); + } } Message::Text(txt_msg) => { break ServerResponse::try_from(txt_msg).map_err(|_| GatewayClientError::MalformedResponse); diff --git a/common/client-libs/gateway-client/src/error.rs b/common/client-libs/gateway-client/src/error.rs index 0890dcd12a..bda9352a3b 100644 --- a/common/client-libs/gateway-client/src/error.rs +++ b/common/client-libs/gateway-client/src/error.rs @@ -64,6 +64,9 @@ pub enum GatewayClientError { #[error("Connection was abruptly closed")] ConnectionAbruptlyClosed, + #[error("Connection was abruptly closed as gateway was stopped")] + ConnectionClosedGatewayShutdown, + #[error("Received response was malformed")] MalformedResponse, @@ -93,6 +96,9 @@ pub enum GatewayClientError { #[error("Timed out")] Timeout, + + #[error("Failed to send mixnet message")] + MixnetMsgSenderFailedToSend, } impl GatewayClientError { diff --git a/common/client-libs/gateway-client/src/packet_router.rs b/common/client-libs/gateway-client/src/packet_router.rs index 7d650ad695..2605b291fb 100644 --- a/common/client-libs/gateway-client/src/packet_router.rs +++ b/common/client-libs/gateway-client/src/packet_router.rs @@ -8,6 +8,9 @@ use futures::channel::mpsc; use log::*; use nymsphinx::addressing::nodes::MAX_NODE_ADDRESS_UNPADDED_LEN; use nymsphinx::params::packet_sizes::PacketSize; +use task::ShutdownListener; + +use crate::error::GatewayClientError; pub type MixnetMessageSender = mpsc::UnboundedSender>>; pub type MixnetMessageReceiver = mpsc::UnboundedReceiver>>; @@ -19,20 +22,26 @@ pub type AcknowledgementReceiver = mpsc::UnboundedReceiver>>; pub struct PacketRouter { ack_sender: AcknowledgementSender, mixnet_message_sender: MixnetMessageSender, + shutdown: Option, } impl PacketRouter { pub fn new( ack_sender: AcknowledgementSender, mixnet_message_sender: MixnetMessageSender, + shutdown: Option, ) -> Self { PacketRouter { ack_sender, mixnet_message_sender, + shutdown, } } - pub fn route_received(&self, unwrapped_packets: Vec>) { + pub fn route_received( + &mut self, + unwrapped_packets: Vec>, + ) -> Result<(), GatewayClientError> { let mut received_messages = Vec::new(); let mut received_acks = Vec::new(); @@ -60,24 +69,28 @@ impl PacketRouter { } } - // due to how we are currently using it, those unwraps can't fail, but if we ever - // wanted to make `gateway-client` into some more generic library, we would probably need - // to catch that error or something. if !received_messages.is_empty() { trace!("routing 'real'"); - self.mixnet_message_sender - .unbounded_send(received_messages) - .unwrap(); + if let Err(err) = self.mixnet_message_sender.unbounded_send(received_messages) { + if let Some(shutdown) = &mut self.shutdown { + if shutdown.is_shutdown_poll() { + // This should ideally not happen, but it's ok + log::warn!("Failed to send mixnet message due to receiver task shutdown"); + return Err(GatewayClientError::MixnetMsgSenderFailedToSend); + } + } + // This should never happen during ordinary operation the way it's currently used. + // Abort to be on the safe side + panic!("Failed to send mixnet message: {:?}", err); + } } if !received_acks.is_empty() { trace!("routing acks"); - match self.ack_sender.unbounded_send(received_acks) { - Ok(_) => {} - Err(e) => { - error!("failed to send ack: {:?}", e); - } + if let Err(e) = self.ack_sender.unbounded_send(received_acks) { + error!("failed to send ack: {:?}", e); }; } + Ok(()) } } diff --git a/common/client-libs/gateway-client/src/socket_state.rs b/common/client-libs/gateway-client/src/socket_state.rs index 3326d17c45..e828f7dbbf 100644 --- a/common/client-libs/gateway-client/src/socket_state.rs +++ b/common/client-libs/gateway-client/src/socket_state.rs @@ -45,9 +45,9 @@ pub(crate) struct PartiallyDelegated { impl PartiallyDelegated { fn route_socket_message( ws_msg: Message, - packet_router: &PacketRouter, + packet_router: &mut PacketRouter, shared_key: &SharedKeys, - ) { + ) -> Result<(), GatewayClientError> { match ws_msg { Message::Binary(bin_msg) => { // this function decrypts the request and checks the MAC @@ -61,7 +61,7 @@ impl PartiallyDelegated { "message received from the gateway was malformed! - {:?}", err ); - return; + return Ok(()); } }; @@ -75,12 +75,15 @@ impl PartiallyDelegated { // This would also require NOT discarding any text responses here. // TODO: those can return the "send confirmations" - perhaps it should be somehow worked around? - Message::Text(text) => trace!( - "received a text message - probably a response to some previous query! - {}", - text - ), - _ => (), - }; + Message::Text(text) => { + trace!( + "received a text message - probably a response to some previous query! - {}", + text + ); + Ok(()) + } + _ => Ok(()), + } } pub(crate) fn split_and_listen_for_mixnet_messages( @@ -99,6 +102,7 @@ impl PartiallyDelegated { let mixnet_receiver_future = async move { let mut fused_receiver = notify_receiver.fuse(); let mut fused_stream = (&mut stream).fuse(); + let mut packet_router = packet_router; // Bit of an ugly workaround for selecting on an `Option`. The unwrap is lazy so we use // this bool inside the `select` to guard against unwrapping in the `None` case. @@ -108,6 +112,12 @@ impl PartiallyDelegated { let ret_err = loop { tokio::select! { + biased; + _ = &mut shutdown_recv_lazy, if shutdown_is_some => { + log::trace!("GatewayClient listener: Received shutdown"); + log::debug!("GatewayClient listener: Exiting"); + return; + } _ = &mut fused_receiver => { break Ok(()); } @@ -116,12 +126,9 @@ impl PartiallyDelegated { Err(err) => break Err(err), Ok(msg) => msg }; - Self::route_socket_message(ws_msg, &packet_router, shared_key.as_ref()); - } - _ = &mut shutdown_recv_lazy, if shutdown_is_some => { - log::trace!("GatewayClient listener: Received shutdown"); - log::debug!("GatewayClient listener: Exiting"); - return; + if let Err(err) = Self::route_socket_message(ws_msg, &mut packet_router, shared_key.as_ref()) { + log::warn!("Route socket message failed: {:?}", err); + } } }; }; diff --git a/common/task/src/shutdown.rs b/common/task/src/shutdown.rs index 92696e8a23..21ea12569d 100644 --- a/common/task/src/shutdown.rs +++ b/common/task/src/shutdown.rs @@ -74,7 +74,7 @@ pub struct ShutdownListener { } impl ShutdownListener { - pub fn new(notify: watch::Receiver<()>) -> ShutdownListener { + fn new(notify: watch::Receiver<()>) -> ShutdownListener { ShutdownListener { shutdown: false, notify,