diff --git a/common/client-libs/gateway-client/src/client/mod.rs b/common/client-libs/gateway-client/src/client/mod.rs index d899ec2928..e6e907c307 100644 --- a/common/client-libs/gateway-client/src/client/mod.rs +++ b/common/client-libs/gateway-client/src/client/mod.rs @@ -293,7 +293,7 @@ impl GatewayClient { // as we need to be able to write the request and read the subsequent response async fn send_websocket_message( &mut self, - msg: Message, + msg: impl Into, ) -> Result { let should_restart_mixnet_listener = if self.connection.is_partially_delegated() { self.recover_socket_connection().await?; @@ -307,7 +307,7 @@ impl GatewayClient { SocketState::NotConnected => return Err(GatewayClientError::ConnectionNotEstablished), _ => return Err(GatewayClientError::ConnectionInInvalidState), }; - conn.send(msg).await?; + conn.send(msg.into()).await?; let response = self.read_control_response().await; if should_restart_mixnet_listener { @@ -481,8 +481,7 @@ impl GatewayClient { encrypted_address, iv, self.cfg.bandwidth.require_tickets, - ) - .into(); + ); match self.send_websocket_message(msg).await? { ServerResponse::Authenticate { @@ -532,6 +531,21 @@ impl GatewayClient { } } + pub async fn get_gateway_protocol(&mut self) -> Result { + if !self.connection.is_established() { + return Err(GatewayClientError::ConnectionNotEstablished); + } + + match self + .send_websocket_message(ClientControlRequest::SupportedProtocol {}) + .await? + { + ServerResponse::SupportedProtocol { version } => Ok(version), + ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)), + _ => Err(GatewayClientError::UnexpectedResponse), + } + } + async fn claim_ecash_bandwidth( &mut self, credential: CredentialSpendingData, @@ -543,8 +557,7 @@ impl GatewayClient { credential, self.shared_key.as_ref().unwrap(), iv, - ) - .into(); + ); let bandwidth_remaining = match self.send_websocket_message(msg).await? { ServerResponse::Bandwidth { available_total } => Ok(available_total), ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)), @@ -562,7 +575,7 @@ impl GatewayClient { } async fn try_claim_testnet_bandwidth(&mut self) -> Result<(), GatewayClientError> { - let msg = ClientControlRequest::ClaimFreeTestnetBandwidth.into(); + let msg = ClientControlRequest::ClaimFreeTestnetBandwidth; let bandwidth_remaining = match self.send_websocket_message(msg).await? { ServerResponse::Bandwidth { available_total } => Ok(available_total), ServerResponse::Error { message } => Err(GatewayClientError::GatewayError(message)), diff --git a/common/gateway-requests/src/types.rs b/common/gateway-requests/src/types.rs index f9006300a1..07145dab8d 100644 --- a/common/gateway-requests/src/types.rs +++ b/common/gateway-requests/src/types.rs @@ -190,6 +190,7 @@ pub enum ClientControlRequest { iv: Vec, }, ClaimFreeTestnetBandwidth, + SupportedProtocol {}, } impl ClientControlRequest { @@ -229,6 +230,7 @@ impl ClientControlRequest { ClientControlRequest::ClaimFreeTestnetBandwidth => { "ClaimFreeTestnetBandwidth".to_string() } + ClientControlRequest::SupportedProtocol { .. } => "SupportedProtocol".to_string(), } } @@ -272,7 +274,15 @@ impl TryFrom for ClientControlRequest { type Error = serde_json::Error; fn try_from(msg: String) -> Result { - serde_json::from_str(&msg) + msg.parse() + } +} + +impl FromStr for ClientControlRequest { + type Err = serde_json::Error; + + fn from_str(s: &str) -> Result { + serde_json::from_str(s) } } @@ -304,6 +314,9 @@ pub enum ServerResponse { Send { remaining_bandwidth: i64, }, + SupportedProtocol { + version: u8, + }, // Generic error Error { message: String, @@ -324,6 +337,7 @@ impl ServerResponse { ServerResponse::Send { .. } => "Send".to_string(), ServerResponse::Error { .. } => "Error".to_string(), ServerResponse::TypedError { .. } => "TypedError".to_string(), + ServerResponse::SupportedProtocol { .. } => "SupportedProtocol".to_string(), } } pub fn new_error>(msg: S) -> Self { 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 12bbb16b51..18efea6588 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/authenticated.rs @@ -115,6 +115,12 @@ impl IntoWSMessage for Result { } } +impl IntoWSMessage for ServerResponse { + fn into_ws_message(self) -> Message { + self.into() + } +} + pub(crate) struct AuthenticatedHandler { inner: FreshHandler, bandwidth_storage_manager: BandwidthStorageManager, @@ -322,13 +328,28 @@ where .await .map_err(|e| e.into()) .into_ws_message(), - other => RequestHandlingError::IllegalRequest { - additional_context: format!( - "received illegal message of type {} in an authenticated client", - other.name() - ), + ClientControlRequest::SupportedProtocol { .. } => self + .inner + .handle_supported_protocol_request() + .into_ws_message(), + other @ ClientControlRequest::Authenticate { .. } => { + RequestHandlingError::IllegalRequest { + additional_context: format!( + "received illegal message of type {} in an authenticated client", + other.name() + ), + } + .into_error_message() + } + other @ ClientControlRequest::RegisterHandshakeInitRequest { .. } => { + RequestHandlingError::IllegalRequest { + additional_context: format!( + "received illegal message of type {} in an authenticated client", + other.name() + ), + } + .into_error_message() } - .into_error_message(), }, } } 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 ab78e855b0..a4beeb1cc6 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,17 @@ -// Copyright 2021-2023 - Nym Technologies SA +// Copyright 2021-2024 - Nym Technologies SA // SPDX-License-Identifier: GPL-3.0-only -use crate::error::RequestHandlingError; +use crate::node::client_handling::websocket::common_state::CommonHandlerState; +use crate::node::client_handling::websocket::connection_handler::INITIAL_MESSAGE_TIMEOUT; +use crate::node::client_handling::{ + active_clients::ActiveClientsStore, + websocket::{ + connection_handler::{ + AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream, + }, + message_receiver::{IsActive, IsActiveRequestSender}, + }, +}; use futures::{ channel::{mpsc, oneshot}, SinkExt, StreamExt, @@ -18,6 +28,7 @@ use nym_gateway_requests::{ types::{ClientControlRequest, ServerResponse}, BinaryResponse, CURRENT_PROTOCOL_VERSION, INITIAL_PROTOCOL_VERSION, }; +use nym_gateway_storage::{error::StorageError, Storage}; use nym_mixnet_client::forwarder::MixForwardingSender; use nym_sphinx::DestinationAddressBytes; use nym_task::TaskClient; @@ -26,21 +37,10 @@ use std::net::SocketAddr; use std::time::Duration; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::time::timeout; use tokio_tungstenite::tungstenite::{protocol::Message, Error as WsError}; use tracing::*; -use crate::node::client_handling::websocket::common_state::CommonHandlerState; -use crate::node::client_handling::{ - active_clients::ActiveClientsStore, - websocket::{ - connection_handler::{ - AuthenticatedHandler, ClientDetails, InitialAuthResult, SocketStream, - }, - message_receiver::{IsActive, IsActiveRequestSender}, - }, -}; -use nym_gateway_storage::{error::StorageError, Storage}; - #[derive(Debug, Error)] pub(crate) enum InitialAuthenticationError { #[error("Internal gateway storage error")] @@ -80,12 +80,6 @@ pub(crate) enum InitialAuthenticationError { #[error("Attempted to negotiate connection with client using incompatible protocol version. Ours is {current} and the client reports {client:?}")] IncompatibleProtocol { client: Option, current: u8 }, - #[error("failed to send authentication error response: {source}")] - ErrorResponseSendFailure { - #[source] - source: WsError, - }, - #[error("failed to send authentication response: {source}")] ResponseSendFailure { #[source] @@ -95,9 +89,6 @@ pub(crate) enum InitialAuthenticationError { #[error("possibly received a sphinx packet without prior authentication. Request is going to be ignored")] BinaryRequestWithoutAuthentication, - #[error("received a connection close message")] - CloseMessage, - #[error("the connection has unexpectedly closed")] ClosedConnection, @@ -107,21 +98,11 @@ pub(crate) enum InitialAuthenticationError { source: WsError, }, + #[error("timed out while waiting for initial data")] + Timeout, + #[error("could not establish client details")] EmptyClientDetails, - - #[error("failed to upgrade the client handler: {source}")] - HandlerUpgradeFailure { source: RequestHandlingError }, - - #[error("received shutdown")] - ReceivedShutdown, -} - -impl InitialAuthenticationError { - /// Converts this Error into an appropriate websocket Message. - fn to_error_message(&self) -> Message { - ServerResponse::new_error(self.to_string()).into() - } } pub(crate) struct FreshHandler { @@ -232,7 +213,10 @@ where /// # Arguments /// /// * `msg`: WebSocket message to write back to the client. - pub(crate) async fn send_websocket_message(&mut self, msg: Message) -> Result<(), WsError> + pub(crate) async fn send_websocket_message( + &mut self, + msg: impl Into, + ) -> Result<(), WsError> where S: AsyncRead + AsyncWrite + Unpin, { @@ -240,11 +224,31 @@ where // TODO: more closely investigate difference between `Sink::send` and `Sink::send_all` // it got something to do with batching and flushing - it might be important if it // turns out somehow we've got a bottleneck here - SocketStream::UpgradedWebSocket(ref mut ws_stream) => ws_stream.send(msg).await, + SocketStream::UpgradedWebSocket(ref mut ws_stream) => ws_stream.send(msg.into()).await, _ => panic!("impossible state - websocket handshake was somehow reverted"), } } + pub(crate) async fn send_error_response( + &mut self, + err: impl std::error::Error, + ) -> Result<(), WsError> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + self.send_websocket_message(ServerResponse::new_error(err.to_string())) + .await + } + + pub(crate) async fn send_and_forget_error_response(&mut self, err: impl std::error::Error) + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + if let Err(err) = self.send_error_response(err).await { + debug!("failed to send error response: {err}") + } + } + /// Sends unwrapped sphinx packets (payloads) back to the client. Note that each message is encrypted and tagged with /// the previously derived shared keys. /// @@ -691,167 +695,177 @@ where )) } - /// Handles data that resembles request to either start registration handshake or perform - /// authentication. - /// - /// # Arguments - /// - /// * `raw_request`: raw text request received from the websocket. - async fn handle_initial_authentication_request( - &mut self, - raw_request: String, - ) -> Result - where - S: AsyncRead + AsyncWrite + Unpin + Send, - { - if let Ok(request) = ClientControlRequest::try_from(raw_request) { - match request { - ClientControlRequest::Authenticate { - protocol_version, - address, - enc_address, - iv, - } => { - self.handle_authenticate(protocol_version, address, enc_address, iv) - .await - } - ClientControlRequest::RegisterHandshakeInitRequest { - protocol_version, - data, - } => self.handle_register(protocol_version, data).await, - // won't accept anything else (like bandwidth) without prior authentication - _ => Err(InitialAuthenticationError::InvalidRequest), - } - } else { - Err(InitialAuthenticationError::InvalidRequest) + pub(crate) fn handle_supported_protocol_request(&self) -> ServerResponse { + debug!("returning gateway protocol version"); + ServerResponse::SupportedProtocol { + version: CURRENT_PROTOCOL_VERSION, } } - /// Listens for only a subset of possible client requests, i.e. for those that can either - /// result in client getting registered or authenticated. All other requests, such as forwarding - /// sphinx packets considered an error and terminate the connection. - // TODO: somehow cleanup this method - pub(crate) async fn perform_initial_authentication( - mut self, - ) -> Result, InitialAuthenticationError> + async fn handle_reply_supported_protocol_request(&mut self) where S: AsyncRead + AsyncWrite + Unpin + Send, { - trace!("Started waiting for authenticate/register request..."); + if let Err(err) = self + .send_websocket_message(self.handle_supported_protocol_request()) + .await + { + debug!("failed to reply with protocol version: {err}") + } + } - let mut shutdown = self.shutdown.clone(); - loop { - tokio::select! { + pub(crate) async fn handle_initial_client_request( + &mut self, + request: ClientControlRequest, + ) -> Result, InitialAuthenticationError> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + // we can handle stateless client requests without prior authentication, like `ClientControlRequest::SupportedProtocol` + let auth_result = match request { + ClientControlRequest::Authenticate { + protocol_version, + address, + enc_address, + iv, + } => { + self.handle_authenticate(protocol_version, address, enc_address, iv) + .await + } + ClientControlRequest::RegisterHandshakeInitRequest { + protocol_version, + data, + } => self.handle_register(protocol_version, data).await, + ClientControlRequest::SupportedProtocol { .. } => { + self.handle_reply_supported_protocol_request().await; + return Ok(None); + } + _ => { + debug!("received an invalid client request"); + return Err(InitialAuthenticationError::InvalidRequest); + } + }; + + let auth_result = match auth_result { + Ok(res) => res, + Err(err) => { + match &err { + InitialAuthenticationError::StorageError(inner_storage) => { + debug!("authentication failure due to storage issue: {inner_storage}") + } + other => debug!("authentication failure: {other}"), + } + + self.send_and_forget_error_response(&err).await; + return Err(err); + } + }; + + // try to send auth response back to the client + if let Err(source) = self + .send_websocket_message(auth_result.server_response) + .await + { + debug!("failed to send authentication response: {source}"); + return Err(InitialAuthenticationError::ResponseSendFailure { source }); + } + + let Some(client_details) = auth_result.client_details else { + // honestly, it's been so long I don't remember under what conditions its possible (if at all) + // to have empty client details + warn!("could not establish client details"); + return Err(InitialAuthenticationError::EmptyClientDetails); + }; + Ok(Some(client_details)) + } + + pub(crate) async fn handle_until_authenticated_or_failure( + mut self, + shutdown: &mut TaskClient, + ) -> Option> + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + while !shutdown.is_shutdown() { + let req = tokio::select! { biased; _ = shutdown.recv() => { - trace!("received shutdown signal while performing initial authentication"); - return Err(InitialAuthenticationError::ReceivedShutdown); + return None }, - msg = self.read_websocket_message() => { - let Some(msg) = msg else { - break; - }; + req = self.wait_for_initial_message() => req, + }; - let msg = match msg { - Ok(msg) => msg, - Err(source) => { - debug!("failed to obtain message from websocket stream! stopping connection handler: {source}"); - return Err(InitialAuthenticationError::FailedToReadMessage { source }); - } - }; - - if msg.is_close() { - return Err(InitialAuthenticationError::CloseMessage); - } - - // ONLY handle 'Authenticate' or 'Register' requests, ignore everything else - match msg { - // we have explicitly checked for close message - Message::Close(_) => unreachable!(), - Message::Text(text_msg) => { - let (mix_sender, mix_receiver) = mpsc::unbounded(); - return match self.handle_initial_authentication_request(text_msg).await { - Err(err) => { - debug!("authentication failure: {err}"); - - // try to send error to the client - if let Err(source) = - self.send_websocket_message(err.to_error_message()).await - { - debug!("Failed to send authentication error response: {source}"); - return Err(InitialAuthenticationError::ErrorResponseSendFailure { - source, - }); - } - // return the underlying error - Err(err) - } - Ok(auth_result) => { - // try to send auth response back to the client - if let Err(source) = self - .send_websocket_message(auth_result.server_response.into()) - .await - { - debug!("Failed to send authentication response: {source}"); - return Err(InitialAuthenticationError::ResponseSendFailure { - source, - }); - } - - if let Some(client_details) = auth_result.client_details { - // 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_remote( - client_details.address, - mix_sender, - is_active_request_sender, - ); - AuthenticatedHandler::upgrade( - self, - client_details, - mix_receiver, - is_active_request_receiver, - ) - .await - .map_err(|source| { - InitialAuthenticationError::HandlerUpgradeFailure { source } - }) - } else { - // honestly, it's been so long I don't remember under what conditions its possible (if at all) - // to have empty client details - Err(InitialAuthenticationError::EmptyClientDetails) - } - } - }; - } - Message::Binary(_) => { - // perhaps logging level should be reduced here, let's leave it for now and see what happens - // if client is working correctly, this should have never happened - debug!("possibly received a sphinx packet without prior authentication. Request is going to be ignored"); - if let Err(source) = self - .send_websocket_message( - ServerResponse::new_error( - "binary request without prior authentication", - ) - .into(), - ) - .await - { - return Err(InitialAuthenticationError::ErrorResponseSendFailure { - source, - }); - } - return Err(InitialAuthenticationError::BinaryRequestWithoutAuthentication); - } - - _ => continue, - }; + let initial_request = match req { + Ok(req) => req, + Err(err) => { + self.send_and_forget_error_response(err).await; + return None; } + }; + + // see if we managed to register the client through this request + let maybe_auth_res = match self.handle_initial_client_request(initial_request).await { + Ok(maybe_auth_res) => maybe_auth_res, + Err(err) => { + debug!("initial client request handling error: {err}"); + self.send_and_forget_error_response(err).await; + return None; + } + }; + + if let Some(registration_details) = maybe_auth_res { + let (mix_sender, mix_receiver) = mpsc::unbounded(); + // 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_remote( + registration_details.address, + mix_sender, + is_active_request_sender, + ); + + return AuthenticatedHandler::upgrade( + self, + registration_details, + mix_receiver, + is_active_request_receiver, + ) + .await + .inspect_err(|err| error!("failed to upgrade client handler: {err}")) + .ok(); } } - Err(InitialAuthenticationError::ClosedConnection) + None + } + + pub(crate) async fn wait_for_initial_message( + &mut self, + ) -> Result + where + S: AsyncRead + AsyncWrite + Unpin + Send, + { + let msg = match timeout(INITIAL_MESSAGE_TIMEOUT, self.read_websocket_message()).await { + Ok(Some(Ok(msg))) => msg, + Ok(Some(Err(source))) => { + debug!("failed to obtain message from websocket stream! stopping connection handler: {source}"); + return Err(InitialAuthenticationError::FailedToReadMessage { source }); + } + Ok(None) => return Err(InitialAuthenticationError::ClosedConnection), + Err(_timeout) => return Err(InitialAuthenticationError::Timeout), + }; + + let text = match msg { + Message::Text(text) => text, + Message::Binary(_) => { + return Err(InitialAuthenticationError::BinaryRequestWithoutAuthentication); + } + _ => unreachable!( + "the underlying tunsgenite stream should be handling other message types" + ), + }; + + text.parse() + .map_err(|_| InitialAuthenticationError::InvalidRequest) } pub(crate) async fn start_handling(self) diff --git a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs index 6e1d8a9b5a..6251ab86bb 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/mod.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/mod.rs @@ -2,7 +2,6 @@ // SPDX-License-Identifier: GPL-3.0-only use crate::config::Config; -use fresh::InitialAuthenticationError; use nym_credential_verification::BandwidthFlushingBehaviourConfig; use nym_gateway_requests::registration::handshake::SharedKeys; use nym_gateway_requests::ServerResponse; @@ -12,7 +11,7 @@ use rand::{CryptoRng, Rng}; use std::time::Duration; use tokio::io::{AsyncRead, AsyncWrite}; use tokio_tungstenite::WebSocketStream; -use tracing::{instrument, trace, warn}; +use tracing::{debug, instrument, trace, warn}; use zeroize::{Zeroize, ZeroizeOnDrop}; pub(crate) use self::authenticated::AuthenticatedHandler; @@ -22,6 +21,7 @@ pub(crate) mod authenticated; mod fresh; const WEBSOCKET_HANDSHAKE_TIMEOUT: Duration = Duration::from_millis(1_500); +const INITIAL_MESSAGE_TIMEOUT: Duration = Duration::from_millis(10_000); // TODO: note for my future self to consider the following idea: // split the socket connection into sink and stream @@ -92,6 +92,12 @@ where S: AsyncRead + AsyncWrite + Unpin + Send, St: Storage + Clone + 'static, { + // don't accept any new requests if we have already received shutdown + if handle.shutdown.is_shutdown() { + debug!("stopping the handle as we have received a shutdown"); + return; + } + // If the connection handler abruptly stops, we shouldn't signal global shutdown handle.shutdown.disarm(); @@ -101,35 +107,29 @@ where ) .await { - Err(timeout_err) => { - warn!("websocket handshake timedout: {timeout_err}"); + Err(_elapsed) => { + warn!("websocket handshake timeout"); return; } Ok(Err(err)) => { - warn!("Failed to complete WebSocket handshake: {err}. Stopping the handler"); + debug!("failed to complete WebSocket handshake: {err}. Stopping the handler"); return; } _ => {} } - trace!("Managed to perform websocket handshake!"); + trace!("managed to perform websocket handshake!"); - let shutdown = handle.shutdown.clone(); - match handle.perform_initial_authentication().await { - // For storage error, we want to print the extended storage error, but without - // including it in the error that's returned to the clients - Err(InitialAuthenticationError::StorageError(err)) => { - warn!("authentication has failed: {err}"); - return; - } - Err(err) => { - warn!("authentication has failed: {err}"); - return; - } - Ok(auth_handle) => auth_handle.listen_for_requests(shutdown).await, + let mut shutdown = handle.shutdown.clone(); + + if let Some(auth_handle) = handle + .handle_until_authenticated_or_failure(&mut shutdown) + .await + { + auth_handle.listen_for_requests(shutdown).await } - trace!("The handler is done!"); + trace!("the handler is done!"); } impl<'a> From<&'a Config> for BandwidthFlushingBehaviourConfig {