GRIN/nym
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

gateway-client: additional shutdown handling

Browse Source
This commit is contained in:
Jon Häggblad
2022-09-08 15:42:21 +02:00
parent d109c53370
commit e0567dddf2
7 changed files with 78 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
clients/client-core/src/client/received_buffer.rs
+2 -2
View File
@@ -1,8 +1,8 @@
// Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
// 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;
clients/client-core/src/init.rs
+1
View File
@@ -90,6 +90,7 @@ async fn register_with_gateway(
gateway.owner.clone(),
our_identity.clone(),
timeout,
None,
);
gateway_client
.establish_connection()
common/client-libs/gateway-client/src/client.rs
+21 -7
View File
@@ -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<ShutdownListener>,
}
@@ -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<identity::KeyPair>,
response_timeout_duration: Duration,
shutdown: Option<ShutdownListener>,
) -> 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);
common/client-libs/gateway-client/src/error.rs
+6
View File
@@ -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 {
common/client-libs/gateway-client/src/packet_router.rs
+25 -12
View File
@@ -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<Vec<Vec<u8>>>;
pub type MixnetMessageReceiver = mpsc::UnboundedReceiver<Vec<Vec<u8>>>;
@@ -19,20 +22,26 @@ pub type AcknowledgementReceiver = mpsc::UnboundedReceiver<Vec<Vec<u8>>>;
pub struct PacketRouter {
ack_sender: AcknowledgementSender,
mixnet_message_sender: MixnetMessageSender,
shutdown: Option<ShutdownListener>,
}
impl PacketRouter {
pub fn new(
ack_sender: AcknowledgementSender,
mixnet_message_sender: MixnetMessageSender,
shutdown: Option<ShutdownListener>,
) -> Self {
PacketRouter {
ack_sender,
mixnet_message_sender,
shutdown,
}
}
pub fn route_received(&self, unwrapped_packets: Vec<Vec<u8>>) {
pub fn route_received(
&mut self,
unwrapped_packets: Vec<Vec<u8>>,
) -> 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(())
}
}
common/client-libs/gateway-client/src/socket_state.rs
+22 -15
View File
@@ -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);
}
}
};
};
common/task/src/shutdown.rs
+1 -1
View File
@@ -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,