gateway: disconnect inactive duplicate clients (#3796)

* gateway: disconnect inactive duplicate clients

* wip: see if we can switch to single ping at a time

* Finish reworking ping pong request flow

* Use workspace version of tokio

* Bundle active client channels into struct

* Fix typo

gateway/Cargo.toml

@@ -37,7 +37,7 @@ serde_json = { workspace = true }
 sqlx = { version = "0.5", features = [ "runtime-tokio-rustls", "sqlite", "macros", "migrate", ] }
 subtle-encoding = { version = "0.5", features = ["bech32-preview"] }
 thiserror = "1"
-tokio = { version = "1.24.1", features = [ "rt-multi-thread", "net", "signal", "fs", ] }
+tokio = { workspace = true, features = [ "rt-multi-thread", "net", "signal", "fs", "time" ] }
 tokio-stream = { version = "0.1.11", features = ["fs"] }
 tokio-tungstenite = "0.14"
 tokio-util = { version = "0.7.4", features = ["codec"] }

gateway/src/node/client_handling/active_clients.rs

@@ -1,13 +1,23 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::node::client_handling::websocket::message_receiver::MixMessageSender;
 use dashmap::DashMap;
 use nym_sphinx::DestinationAddressBytes;
 use std::sync::Arc;
 
+use super::websocket::message_receiver::{IsActiveRequestSender, MixMessageSender};
+
 #[derive(Clone)]
-pub(crate) struct ActiveClientsStore(Arc<DashMap<DestinationAddressBytes, MixMessageSender>>);
+pub(crate) struct ActiveClientsStore(Arc<DashMap<DestinationAddressBytes, ClientIncomingChannels>>);
+
+#[derive(Clone)]
+pub(crate) struct ClientIncomingChannels {
+    // Mix messages coming from the mixnet to the handler of a client.
+    pub mix_message_sender: MixMessageSender,
+
+    // Requests sent from the handler of one client to the handler of other clients.
+    pub is_active_request_sender: IsActiveRequestSender,
+}
 
 impl ActiveClientsStore {
     /// Creates new instance of `ActiveClientsStore` to store in-memory handles to all currently connected clients.
@@ -21,13 +31,13 @@ impl ActiveClientsStore {
     /// # Arguments
     ///
     /// * `client`: address of the client for which to obtain the handle.
-    pub(crate) fn get(&self, client: DestinationAddressBytes) -> Option<MixMessageSender> {
+    pub(crate) fn get(&self, client: DestinationAddressBytes) -> Option<ClientIncomingChannels> {
         let entry = self.0.get(&client)?;
         let handle = entry.value();
 
         // if the entry is stale, remove it from the map
         // if handle.is_valid() {
-        if !handle.is_closed() {
+        if !handle.mix_message_sender.is_closed() {
             Some(handle.clone())
         } else {
             // drop the reference to the map to prevent deadlocks
@@ -52,8 +62,19 @@ impl ActiveClientsStore {
     ///
     /// * `client`: address of the client for which to insert the handle.
     /// * `handle`: the sender channel for all mix packets to be pushed back onto the websocket
-    pub(crate) fn insert(&self, client: DestinationAddressBytes, handle: MixMessageSender) {
-        self.0.insert(client, handle);
+    pub(crate) fn insert(
+        &self,
+        client: DestinationAddressBytes,
+        handle: MixMessageSender,
+        is_active_request_sender: IsActiveRequestSender,
+    ) {
+        self.0.insert(
+            client,
+            ClientIncomingChannels {
+                mix_message_sender: handle,
+                is_active_request_sender,
+            },
+        );
     }
 
     /// Get number of active clients in store

gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs

@@ -1,28 +1,39 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::node::client_handling::websocket::connection_handler::{ClientDetails, FreshHandler};
-use crate::node::client_handling::websocket::message_receiver::MixMessageReceiver;
-use crate::node::storage::error::StorageError;
-use crate::node::storage::Storage;
-use futures::StreamExt;
+use futures::{
+    future::{FusedFuture, OptionFuture},
+    FutureExt, StreamExt,
+};
 use log::*;
-use nym_gateway_requests::iv::IVConversionError;
-use nym_gateway_requests::types::{BinaryRequest, ServerResponse};
-use nym_gateway_requests::{ClientControlRequest, GatewayRequestsError};
+use nym_gateway_requests::{
+    iv::{IVConversionError, IV},
+    types::{BinaryRequest, ServerResponse},
+    ClientControlRequest, GatewayRequestsError,
+};
 use nym_sphinx::forwarding::packet::MixPacket;
-use rand::{CryptoRng, Rng};
-use std::convert::TryFrom;
-use std::process;
-use thiserror::Error;
-use tokio::io::{AsyncRead, AsyncWrite};
-use tokio_tungstenite::tungstenite::protocol::Message;
-
-use crate::node::client_handling::bandwidth::Bandwidth;
-use crate::node::client_handling::FREE_TESTNET_BANDWIDTH_VALUE;
-use nym_gateway_requests::iv::IV;
 use nym_task::TaskClient;
 use nym_validator_client::coconut::CoconutApiError;
+use rand::{CryptoRng, Rng};
+use thiserror::Error;
+use tokio::io::{AsyncRead, AsyncWrite};
+use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError};
+
+use std::{convert::TryFrom, process, time::Duration};
+
+use crate::node::{
+    client_handling::{
+        bandwidth::Bandwidth,
+        websocket::{
+            connection_handler::{ClientDetails, FreshHandler},
+            message_receiver::{
+                IsActive, IsActiveRequestReceiver, IsActiveResultSender, MixMessageReceiver,
+            },
+        },
+        FREE_TESTNET_BANDWIDTH_VALUE,
+    },
+    storage::{error::StorageError, Storage},
+};
 
 #[derive(Debug, Error)]
 pub(crate) enum RequestHandlingError {
@@ -97,6 +108,11 @@ pub(crate) struct AuthenticatedHandler<R, S, St> {
     inner: FreshHandler<R, S, St>,
     client: ClientDetails,
     mix_receiver: MixMessageReceiver,
+    // Occasionally the handler is requested to ping the connected client for confirm that it's
+    // active, such as when a duplicate connection is detected. This hashmap stores the oneshot
+    // senders that are used to return the result of the ping to the handler requesting the ping.
+    is_active_request_receiver: IsActiveRequestReceiver,
+    is_active_ping_pending_reply: Option<(u64, IsActiveResultSender)>,
 }
 
 // explicitly remove handle from the global store upon being dropped
@@ -127,11 +143,14 @@ where
         fresh: FreshHandler<R, S, St>,
         client: ClientDetails,
         mix_receiver: MixMessageReceiver,
+        is_active_request_receiver: IsActiveRequestReceiver,
     ) -> Self {
         AuthenticatedHandler {
             inner: fresh,
             client,
             mix_receiver,
+            is_active_request_receiver,
+            is_active_ping_pending_reply: None,
         }
     }
 
@@ -351,6 +370,29 @@ where
         }
     }
 
+    /// Handles pong message received from the client.
+    /// If the client is still active, the handler that requested the ping will receive a reply.
+    async fn handle_pong(&mut self, msg: Vec<u8>) {
+        if let Ok(msg) = msg.try_into() {
+            let msg = u64::from_be_bytes(msg);
+            trace!("Received pong from client: {}", msg);
+            if let Some((tag, _)) = &self.is_active_ping_pending_reply {
+                if tag == &msg {
+                    debug!("Reporting back to the handler that the client is still active");
+                    let tx = self.is_active_ping_pending_reply.take().unwrap().1;
+                    if let Err(err) = tx.send(IsActive::Active) {
+                        warn!("Failed to send pong reply back to the requesting handler: {err:?}");
+                    }
+                } else {
+                    warn!(
+                        "Received pong reply from the client with unexpected tag: {}",
+                        msg
+                    );
+                }
+            }
+        }
+    }
+
     /// Attempts to handle websocket message received from the connected client.
     ///
     /// # Arguments
@@ -363,10 +405,64 @@ where
         match raw_request {
             Message::Binary(bin_msg) => Some(self.handle_binary(bin_msg).await),
             Message::Text(text_msg) => Some(self.handle_text(text_msg).await),
+            Message::Pong(msg) => {
+                self.handle_pong(msg).await;
+                None
+            }
             _ => None,
         }
     }
 
+    /// Send a ping to the connected client and return a tag identifying the ping.
+    async fn send_ping(&mut self) -> Result<u64, WsError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        let tag: u64 = rand::thread_rng().gen();
+        debug!("Got request to ping our connection: {}", tag);
+        self.inner
+            .send_websocket_message(Message::Ping(tag.to_be_bytes().to_vec()))
+            .await?;
+        Ok(tag)
+    }
+
+    /// Handles the ping timeout by responding back to the handler that requested the ping.
+    async fn handle_ping_timeout(&mut self) {
+        debug!("Ping timeout expired!");
+        if let Some((_tag, reply_tx)) = self.is_active_ping_pending_reply.take() {
+            if let Err(err) = reply_tx.send(IsActive::NotActive) {
+                warn!("Failed to respond back to the handler requesting the ping: {err:?}");
+            }
+        }
+    }
+
+    async fn handle_is_active_request(
+        &mut self,
+        reply_tx: IsActiveResultSender,
+    ) -> Result<(), WsError>
+    where
+        S: AsyncRead + AsyncWrite + Unpin,
+    {
+        if self.is_active_ping_pending_reply.is_some() {
+            warn!("Received request to ping the client, but a ping is already in progress!");
+            if let Err(err) = reply_tx.send(IsActive::BusyPinging) {
+                warn!("Failed to respond back to the handler requesting the ping: {err:?}");
+            }
+            return Ok(());
+        }
+
+        match self.send_ping().await {
+            Ok(tag) => {
+                self.is_active_ping_pending_reply = Some((tag, reply_tx));
+                Ok(())
+            }
+            Err(err) => {
+                warn!("Failed to send ping to client: {err}. Assuming the connection is dead.");
+                Err(err)
+            }
+        }
+    }
+
     /// Simultaneously listens for incoming client requests, which realistically should only be
     /// binary requests to forward sphinx packets or increase bandwidth
     /// and for sphinx packets received from the mix network that should be sent back to the client.
@@ -377,11 +473,32 @@ where
     {
         trace!("Started listening for ALL incoming requests...");
 
+        // Ping timeout future used to check if the client responded to our ping request
+        let mut ping_timeout: OptionFuture<_> = None.into();
+
         while !shutdown.is_shutdown() {
             tokio::select! {
                 _ = shutdown.recv() => {
                     log::trace!("client_handling::AuthenticatedHandler: received shutdown");
-                }
+                },
+                // Received a request to ping the client to check if it's still active
+                tx = self.is_active_request_receiver.next() => {
+                    match tx {
+                        None => break,
+                        Some(reply_tx) => {
+                            if self.handle_is_active_request(reply_tx).await.is_err() {
+                                break;
+                            }
+                            // NOTE: fuse here due to .is_terminated() check below
+                            ping_timeout = Some(Box::pin(tokio::time::sleep(Duration::from_millis(1000)).fuse())).into();
+                        }
+                    };
+                },
+                // The ping timeout expired, meaning the client didn't respond to our ping request
+                _ = &mut ping_timeout, if !ping_timeout.is_terminated() => {
+                   ping_timeout = None.into();
+                   self.handle_ping_timeout().await;
+                },
                 socket_msg = self.inner.read_websocket_message() => {
                     let socket_msg = match socket_msg {
                         None => break,
@@ -406,7 +523,13 @@ where
                     }
                 },
                 mix_messages = self.mix_receiver.next() => {
-                    let mix_messages = mix_messages.expect("sender was unexpectedly closed! this shouldn't have ever happened!");
+                    let mix_messages = match mix_messages {
+                        None => {
+                            warn!("mix receiver was closed! Assuming the connection is dead.");
+                            break;
+                        }
+                        Some(mix_messages) => mix_messages,
+                    };
                     if let Err(err) = self.inner.push_packets_to_client(&self.client.shared_keys, mix_messages).await {
                         warn!("failed to send the unwrapped sphinx packets back to the client - {err}, assuming the connection is dead");
                         break;

gateway/src/node/client_handling/websocket/connection_handler/fresh.rs

@@ -1,33 +1,43 @@
 // Copyright 2021 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::node::client_handling::active_clients::ActiveClientsStore;
-use crate::node::client_handling::websocket::connection_handler::coconut::CoconutVerifier;
-use crate::node::client_handling::websocket::connection_handler::{
-    AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream,
+use futures::{
+    channel::{mpsc, oneshot},
+    SinkExt, StreamExt,
 };
-use crate::node::storage::error::StorageError;
-use crate::node::storage::Storage;
-use futures::{channel::mpsc, SinkExt, StreamExt};
 use log::*;
 use nym_crypto::asymmetric::identity;
 use nym_gateway_requests::authentication::encrypted_address::{
     EncryptedAddressBytes, EncryptedAddressConversionError,
 };
-use nym_gateway_requests::iv::{IVConversionError, IV};
-use nym_gateway_requests::registration::handshake::error::HandshakeError;
-use nym_gateway_requests::registration::handshake::{gateway_handshake, SharedKeys};
-use nym_gateway_requests::types::{ClientControlRequest, ServerResponse};
-use nym_gateway_requests::{BinaryResponse, PROTOCOL_VERSION};
+use nym_gateway_requests::{
+    iv::{IVConversionError, IV},
+    registration::handshake::{error::HandshakeError, gateway_handshake, SharedKeys},
+    types::{ClientControlRequest, ServerResponse},
+    BinaryResponse, PROTOCOL_VERSION,
+};
 use nym_mixnet_client::forwarder::MixForwardingSender;
 use nym_sphinx::DestinationAddressBytes;
 use rand::{CryptoRng, Rng};
-use std::convert::TryFrom;
-use std::sync::Arc;
+use std::{convert::TryFrom, sync::Arc, time::Duration};
 use thiserror::Error;
 use tokio::io::{AsyncRead, AsyncWrite};
 use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError};
 
+use crate::node::{
+    client_handling::{
+        active_clients::ActiveClientsStore,
+        websocket::{
+            connection_handler::{
+                coconut::CoconutVerifier, AuthenticatedHandler, ClientDetails, InitialAuthResult,
+                SocketStream,
+            },
+            message_receiver::{IsActive, IsActiveRequestSender},
+        },
+    },
+    storage::{error::StorageError, Storage},
+};
+
 #[derive(Debug, Error)]
 pub(crate) enum InitialAuthenticationError {
     #[error("Internal gateway storage error")]
@@ -394,6 +404,59 @@ where
         }
     }
 
+    async fn handle_duplicate_client(
+        &mut self,
+        address: DestinationAddressBytes,
+        mut is_active_request_tx: IsActiveRequestSender,
+    ) -> Result<(), InitialAuthenticationError> {
+        // Ask the other connection to ping if they are still active.
+        // Use a oneshot channel to return the result to us
+        let (ping_result_sender, ping_result_receiver) = oneshot::channel();
+        log::debug!("Asking other connection handler to ping the connected client to see if it is still active");
+        if let Err(err) = is_active_request_tx.send(ping_result_sender).await {
+            warn!("Failed to send ping request to other handler: {err}");
+        }
+
+        // Wait for the reply
+        match tokio::time::timeout(Duration::from_millis(2000), ping_result_receiver).await {
+            Ok(Ok(res)) => {
+                match res {
+                    IsActive::NotActive => {
+                        // The other handler reported that the client is not active, so we can
+                        // disconnect the other client and continue with this connection.
+                        log::debug!("Other handler reports it is not active");
+                        self.active_clients_store.disconnect(address);
+                    }
+                    IsActive::Active => {
+                        // The other handled reported a positive reply, so we have to assume it's
+                        // still active and disconnect this connection.
+                        log::info!("Other handler reports it is active");
+                        return Err(InitialAuthenticationError::DuplicateConnection);
+                    }
+                    IsActive::BusyPinging => {
+                        // The other handler is already busy pinging the client, so we have to
+                        // assume it's still active and disconnect this connection.
+                        log::debug!("Other handler reports it is already busy pinging the client");
+                        return Err(InitialAuthenticationError::DuplicateConnection);
+                    }
+                }
+            }
+            Ok(Err(_)) => {
+                // Other channel failed to reply (the channel sender probably dropped)
+                log::info!("Other connection failed to reply, disconnecting it in favour of this new connection");
+                self.active_clients_store.disconnect(address);
+            }
+            Err(_) => {
+                // Timeout waiting for reply
+                log::warn!(
+                    "Other connection timed out, disconnecting it in favour of this new connection"
+                );
+                self.active_clients_store.disconnect(address);
+            }
+        }
+        Ok(())
+    }
+
     /// Tries to handle the received authentication request by checking correctness of the received data.
     ///
     /// # Arguments
@@ -418,8 +481,11 @@ where
         let encrypted_address = EncryptedAddressBytes::try_from_base58_string(enc_address)?;
         let iv = IV::try_from_base58_string(iv)?;
 
-        if self.active_clients_store.get(address).is_some() {
-            return Err(InitialAuthenticationError::DuplicateConnection);
+        // Check for duplicate clients
+        if let Some(client_tx) = self.active_clients_store.get(address) {
+            log::warn!("Detected duplicate connection for client: {}", address);
+            self.handle_duplicate_client(address, client_tx.is_active_request_sender)
+                .await?;
         }
 
         let shared_keys = self
@@ -606,12 +672,19 @@ where
                             }
 
                             return if let Some(client_details) = auth_result.client_details {
-                                self.active_clients_store
-                                    .insert(client_details.address, mix_sender);
+                                // Channel for handlers to ask other handlers if they are still active.
+                                let (is_active_request_sender, is_active_request_receiver) =
+                                    mpsc::unbounded();
+                                self.active_clients_store.insert(
+                                    client_details.address,
+                                    mix_sender,
+                                    is_active_request_sender,
+                                );
                                 Some(AuthenticatedHandler::upgrade(
                                     self,
                                     client_details,
                                     mix_receiver,
+                                    is_active_request_receiver,
                                 ))
                             } else {
                                 None

gateway/src/node/client_handling/websocket/connection_handler/mod.rs

@@ -100,7 +100,7 @@ pub(crate) async fn handle_connection<R, S, St>(
         .await
     {
         None => {
-            trace!("received shutdown signal while performing initial authetnication");
+            trace!("received shutdown signal while performing initial authentication");
             return;
         }
         Some(None) => {

gateway/src/node/client_handling/websocket/message_receiver.rs

@@ -1,7 +1,20 @@
 // Copyright 2020 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use futures::channel::mpsc;
+use futures::channel::{mpsc, oneshot};
 
 pub(crate) type MixMessageSender = mpsc::UnboundedSender<Vec<Vec<u8>>>;
 pub(crate) type MixMessageReceiver = mpsc::UnboundedReceiver<Vec<Vec<u8>>>;
+
+// Channels used for one handler to requester another handler to check that the client is still
+// active. The result is then passed back to the requesting handler in the oneshot channel.
+pub(crate) type IsActiveRequestSender = mpsc::UnboundedSender<IsActiveResultSender>;
+pub(crate) type IsActiveRequestReceiver = mpsc::UnboundedReceiver<IsActiveResultSender>;
+
+#[derive(Debug)]
+pub(crate) enum IsActive {
+    Active,
+    NotActive,
+    BusyPinging,
+}
+pub(crate) type IsActiveResultSender = oneshot::Sender<IsActive>;

gateway/src/node/mixnet_handling/receiver/connection_handler.rs

@@ -70,9 +70,9 @@ impl<St: Storage> ConnectionHandler<St> {
     }
 
     fn update_clients_store_cache_entry(&mut self, client_address: DestinationAddressBytes) {
-        if let Some(client_sender) = self.active_clients_store.get(client_address) {
+        if let Some(client_senders) = self.active_clients_store.get(client_address) {
             self.clients_store_cache
-                .insert(client_address, client_sender);
+                .insert(client_address, client_senders.mix_message_sender);
         }
     }