From 958cdfd22c09edd77003d5c2eb9b341974052ecb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20H=C3=A4ggblad?= Date: Fri, 18 Aug 2023 16:46:05 +0200 Subject: [PATCH] wip --- .../node/client_handling/active_clients.rs | 48 +++++++++++++++++-- .../connection_handler/authenticated.rs | 6 +++ .../websocket/connection_handler/fresh.rs | 24 +++++++++- 3 files changed, 73 insertions(+), 5 deletions(-) diff --git a/gateway/src/node/client_handling/active_clients.rs b/gateway/src/node/client_handling/active_clients.rs index d2de2cc5b5..8dd04f42c1 100644 --- a/gateway/src/node/client_handling/active_clients.rs +++ b/gateway/src/node/client_handling/active_clients.rs @@ -6,8 +6,17 @@ use dashmap::DashMap; use nym_sphinx::DestinationAddressBytes; use std::sync::Arc; +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RecentlyActive { + Yes, + YesButReset, + No, +} + #[derive(Clone)] -pub(crate) struct ActiveClientsStore(Arc>); +pub(crate) struct ActiveClientsStore( + Arc>, +); impl ActiveClientsStore { /// Creates new instance of `ActiveClientsStore` to store in-memory handles to all currently connected clients. @@ -23,7 +32,7 @@ impl ActiveClientsStore { /// * `client`: address of the client for which to obtain the handle. pub(crate) fn get(&self, client: DestinationAddressBytes) -> Option { let entry = self.0.get(&client)?; - let handle = entry.value(); + let handle = &entry.value().0; // if the entry is stale, remove it from the map // if handle.is_valid() { @@ -37,6 +46,38 @@ impl ActiveClientsStore { } } + pub(crate) fn get_is_recently_active(&self, client: DestinationAddressBytes) -> Option { + let entry = self.0.get(&client)?; + let is_recently_active = &entry.value().1; + + // if the entry is stale, remove it from the map + // if handle.is_valid() { + if !entry.value().0.is_closed() { + Some(is_recently_active.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; + } + } + + pub(crate) fn reset_activity(&self, client: DestinationAddressBytes) { + if let Some(mut entry) = self.0.get_mut(&client) { + if entry.value_mut().1 == RecentlyActive::Yes { + entry.value_mut().1 = RecentlyActive::YesButReset; + } else if entry.value_mut().1 == RecentlyActive::YesButReset { + entry.value_mut().1 = RecentlyActive::No; + } + } + } + /// Indicates particular client has disconnected from the gateway and its handle should get removed. /// /// # Arguments @@ -53,11 +94,12 @@ 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); + self.0.insert(client, (handle, RecentlyActive::Yes)); } /// Get number of active clients in store pub(crate) fn size(&self) -> usize { self.0.len() } + } 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 f9753cbdf8..61613172d6 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -135,6 +135,12 @@ where } } + fn register_activity(&self) { + self.inner + .active_clients_store + .register_activity(self.client.address) + } + /// Explicitly removes handle from the global store. fn disconnect(self) { self.inner 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 b1dd9458b1..afcc01aa12 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -1,7 +1,7 @@ // Copyright 2021 - Nym Technologies SA // SPDX-License-Identifier: Apache-2.0 -use crate::node::client_handling::active_clients::ActiveClientsStore; +use crate::node::client_handling::active_clients::{ActiveClientsStore, RecentlyActive}; use crate::node::client_handling::websocket::connection_handler::coconut::CoconutVerifier; use crate::node::client_handling::websocket::connection_handler::{ AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream, @@ -420,8 +420,28 @@ 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); + } + + 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) + // } else { // return Err(InitialAuthenticationError::DuplicateConnection); - self.active_clients_store.disconnect(address) + // } } let shared_keys = self