diff --git a/clients/desktop/src/client/topology_control.rs b/clients/desktop/src/client/topology_control.rs index 1f91f39ed9..eb041d22bc 100644 --- a/clients/desktop/src/client/topology_control.rs +++ b/clients/desktop/src/client/topology_control.rs @@ -157,7 +157,7 @@ impl TopologyRefresher { } async fn get_current_compatible_topology(&self) -> T { - // note: this call makes it neccessary that `T::new()`does *not* have 'static lifetime + // note: this call makes it necessary that `T::new()`does *not* have 'static lifetime let full_topology = T::new(self.directory_server.clone()).await; // just filter by version and assume the validators will remove all bad behaving // nodes with the staking diff --git a/clients/desktop/src/websocket/handler.rs b/clients/desktop/src/websocket/handler.rs index 8c9bf8a48d..8753bed071 100644 --- a/clients/desktop/src/websocket/handler.rs +++ b/clients/desktop/src/websocket/handler.rs @@ -309,9 +309,13 @@ impl Handler { // consume self to make sure `drop` is called after this is done pub(crate) async fn handle_connection(mut self, socket: TcpStream) { - let ws_stream = accept_async(socket) - .await - .expect("error while performing the websocket handshake"); + let ws_stream = match accept_async(socket).await { + Ok(ws_stream) => ws_stream, + Err(err) => { + warn!("error while performing the websocket handshake - {:?}", err); + return; + } + }; self.socket = Some(ws_stream); let (reconstructed_sender, reconstructed_receiver) = mpsc::unbounded(); diff --git a/gateway/src/node/client_handling/websocket/connection_handler.rs b/gateway/src/node/client_handling/websocket/connection_handler.rs index 3ecffe5f2e..75c83fda98 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler.rs @@ -74,19 +74,24 @@ where } } - async fn perform_websocket_handshake(&mut self) { + async fn perform_websocket_handshake(&mut self) -> Result<(), WsError> { self.socket_connection = match std::mem::replace(&mut self.socket_connection, SocketStream::Invalid) { SocketStream::RawTCP(conn) => { // TODO: perhaps in the future, rather than panic here (and uncleanly shut tcp stream) // return a result with an error? - let ws_stream = tokio_tungstenite::accept_async(conn) - .await - .expect("Failed to perform websocket handshake"); + let ws_stream = match tokio_tungstenite::accept_async(conn).await { + Ok(ws_stream) => ws_stream, + // note that socket will remain in `Invalid` state here, but that's + // absolutely fine because due to returned error the handler + // should terminate immediately + Err(err) => return Err(err), + }; SocketStream::UpgradedWebSocket(ws_stream) } other => other, - } + }; + Ok(()) } async fn next_websocket_request(&mut self) -> Option> { @@ -366,7 +371,13 @@ where } pub(crate) async fn start_handling(&mut self) { - self.perform_websocket_handshake().await; + if let Err(e) = self.perform_websocket_handshake().await { + warn!( + "Failed to complete WebSocket handshake - {:?}. Stopping the handler", + e + ); + return; + } trace!("Managed to perform websocket handshake!"); let mix_receiver = self.wait_for_initial_authentication().await; trace!("Performed initial authentication");