diff --git a/gateway/src/node/client_handling/active_clients.rs b/gateway/src/node/client_handling/active_clients.rs index 8dd04f42c1..89217aa1cb 100644 --- a/gateway/src/node/client_handling/active_clients.rs +++ b/gateway/src/node/client_handling/active_clients.rs @@ -4,6 +4,7 @@ use crate::node::client_handling::websocket::message_receiver::MixMessageSender; use dashmap::DashMap; use nym_sphinx::DestinationAddressBytes; +use futures::channel::{mpsc, oneshot}; use std::sync::Arc; #[derive(Clone, Debug, PartialEq, Eq)] @@ -15,7 +16,7 @@ pub enum RecentlyActive { #[derive(Clone)] pub(crate) struct ActiveClientsStore( - Arc>, + Arc>)>>, ); impl ActiveClientsStore { @@ -62,6 +63,22 @@ impl ActiveClientsStore { } } + pub(crate) fn get_is_active_sender(&self, client: DestinationAddressBytes) -> Option>> { + let entry = self.0.get(&client)?; + let is_active_sender = &entry.value().2; + + // if the entry is stale, remove it from the map + // if handle.is_valid() { + if !entry.value().0.is_closed() { + Some(is_active_sender.clone()) + } else { + // drop the reference to the map to prevent deadlocks + drop(entry); + self.0.remove(&client); + None + } + } + pub(crate) fn register_activity(&self, client: DestinationAddressBytes) { if let Some(mut entry) = self.0.get_mut(&client) { entry.value_mut().1 = RecentlyActive::Yes; @@ -93,8 +110,8 @@ 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, RecentlyActive::Yes)); + pub(crate) fn insert(&self, client: DestinationAddressBytes, handle: MixMessageSender, is_active_sender: mpsc::UnboundedSender>) { + self.0.insert(client, (handle, RecentlyActive::Yes, is_active_sender)); } /// Get number of active clients in store diff --git a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs index 61613172d6..a51f7b3ee9 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -317,6 +317,7 @@ where /// /// * `bin_msg`: raw message to handle. async fn handle_binary(&self, bin_msg: Vec) -> Message { + println!("handle_binary"); // this function decrypts the request and checks the MAC match BinaryRequest::try_from_encrypted_tagged_bytes(bin_msg, &self.client.shared_keys) { Err(e) => { @@ -341,6 +342,7 @@ where /// /// * `raw_request`: raw message to handle. async fn handle_text(&mut self, raw_request: String) -> Message { + println!("handle_text"); match ClientControlRequest::try_from(raw_request) { Err(e) => RequestHandlingError::InvalidTextRequest(e).into_error_message(), Ok(request) => match request { @@ -363,12 +365,29 @@ where /// /// * `raw_request`: raw received websocket message. async fn handle_request(&mut self, raw_request: Message) -> Option { + println!("handle_request"); // apparently tungstenite auto-handles ping/pong/close messages so for now let's ignore // them and let's test that claim. If that's not the case, just copy code from // desktop nym-client websocket as I've manually handled everything there 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::Ping(s) => { + println!("Received ping from client: {}", String::from_utf8_lossy(&s)); + None + } + Message::Pong(s) => { + // println!("1 Received pong from client: {}", String::from_utf8_lossy(&s)); + let ss = nym_sphinx::DestinationAddressBytes::try_from_byte_slice(&s).expect("JON"); + if true { + println!("Received pong from client: {}", ss); + if let Some(tx) = self.inner.is_active_pending_replies.remove(&ss) { + println!("Reporting back"); + tx.send(true).expect("JON"); + } + } + None + } _ => None, } } @@ -383,11 +402,36 @@ where { trace!("Started listening for ALL incoming requests..."); + let mut is_active_receiver = self + .inner + .is_active_receiver + .take() + .expect("we should always have a receiver"); + while !shutdown.is_shutdown() { tokio::select! { _ = shutdown.recv() => { log::trace!("client_handling::AuthenticatedHandler: received shutdown"); - } + }, + tx = is_active_receiver.next() => { + match tx { + None => { + println!("Got None in listen_for_requests for is_active_receiver"); + } + Some(tx) => { + println!("Received tx"); + + // WIP(JON) + // if we receive a request to ping the client, we should do so and respond with + // the result + println!("Got request to ping our connection"); + // self.inner.send_websocket_message(Message::Ping("jon".to_string().into_bytes())).await; + let payload = self.client.address.as_bytes().to_vec(); + self.inner.send_websocket_message(Message::Ping(payload)).await; + self.inner.is_active_pending_replies.insert(self.client.address, tx); + } + }; + }, socket_msg = self.inner.read_websocket_message() => { println!("socket_msg"); let socket_msg = match socket_msg { diff --git a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs index afcc01aa12..5b8dc0f504 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -8,6 +8,7 @@ use crate::node::client_handling::websocket::connection_handler::{ }; use crate::node::storage::error::StorageError; use crate::node::storage::Storage; +use futures::channel::oneshot; use futures::{channel::mpsc, SinkExt, StreamExt}; use log::*; use nym_crypto::asymmetric::identity; @@ -22,8 +23,10 @@ use nym_gateway_requests::{BinaryResponse, PROTOCOL_VERSION}; use nym_mixnet_client::forwarder::MixForwardingSender; use nym_sphinx::DestinationAddressBytes; use rand::{CryptoRng, Rng}; +use std::collections::HashMap; use std::convert::TryFrom; use std::sync::Arc; +use std::time::Duration; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError}; @@ -75,6 +78,8 @@ pub(crate) struct FreshHandler { pub(crate) socket_connection: SocketStream, pub(crate) storage: St, pub(crate) coconut_verifier: Arc, + pub(crate) is_active_receiver: Option>>, + pub(crate) is_active_pending_replies: HashMap>, } impl FreshHandler @@ -106,6 +111,8 @@ where local_identity, storage, coconut_verifier, + is_active_receiver: None, + is_active_pending_replies: HashMap::new(), } } @@ -421,21 +428,53 @@ where if self.active_clients_store.get(address).is_some() { println!("duplicate connection when authenticating"); - let msg = Message::Ping("duplicate connection".to_string().into_bytes()); - if let Err(err) = self.send_websocket_message(msg).await { - warn!( - "Failed to send message over websocket: {err}. \ - Assuming the connection is dead.", - ); - println!("disconnecting"); - self.active_clients_store.disconnect(address); + // WIP(JON) + if let Some(mut tx) = self.active_clients_store.get_is_active_sender(address) { + // Ask the other connection to ping if they are still active. + // Use a oneshot channel to return the result to us + let (reply_sender, reply_receiver) = oneshot::channel::(); + println!("Asking other connection to ping"); + tx.send(reply_sender).await; + + // Wait for the reply + // let r = reply_receiver.await.expect("JON"); + let r = tokio::time::timeout(Duration::from_millis(1000), reply_receiver).await; + match r { + Ok(r) => { + let r = r.expect("JON"); + if r == true { + println!("Get a positive reply"); + return Err(InitialAuthenticationError::DuplicateConnection); + } else { + println!("Get a negative reply"); + self.active_clients_store.disconnect(address); + } + } + Err(err) => { + println!("ping timeout"); + self.active_clients_store.disconnect(address); + } + } + + // let msg = Message::Ping("duplicate connection".to_string().into_bytes()); + // println!("Sending ping"); + // if let Err(err) = self.send_websocket_message(msg).await { + // warn!( + // "Failed to send message over websocket: {err}. \ + // Assuming the connection is dead.", + // ); + // println!("disconnecting"); + // self.active_clients_store.disconnect(address); + // } + // println!("waiting.."); + // let msg = self.read_websocket_message().await; + // dbg!(&msg); + // tokio::time::sleep(Duration::from_secs(1)).await; + } else { + + println!("fail without pinging"); + return Err(InitialAuthenticationError::DuplicateConnection); } - - println!("waiting.."); - let msg = self.read_websocket_message().await; - dbg!(&msg); - - return Err(InitialAuthenticationError::DuplicateConnection); // if self.active_clients_store.get_is_recently_active(address) == Some(RecentlyActive::No) { // println!("disconnecting"); // self.active_clients_store.disconnect(address) @@ -629,8 +668,16 @@ where } return if let Some(client_details) = auth_result.client_details { - self.active_clients_store - .insert(client_details.address, mix_sender); + // create mpsc channel + let (is_active_sender, is_active_receiver) = + mpsc::unbounded::>(); + self.is_active_receiver = Some(is_active_receiver); + + self.active_clients_store.insert( + client_details.address, + mix_sender, + is_active_sender, + ); Some(AuthenticatedHandler::upgrade( self, client_details,