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 ab308e7f8e..b7fe12fc3a 100644 --- a/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs +++ b/gateway/src/node/client_handling/websocket/connection_handler/fresh.rs @@ -922,47 +922,54 @@ impl FreshHandler { S: AsyncRead + AsyncWrite + Unpin + Send, R: CryptoRng + RngCore + Send, { - let initial_request = match self.wait_for_initial_message().await { - Ok(req) => req, - Err(err) => { - self.send_and_forget_error_response(err).await; + let shutdown = self.shutdown.clone(); + loop { + if shutdown.is_cancelled() { + trace!("received cancellation before authentication"); 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; + let req = self.wait_for_initial_message().await; + 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.shared_state.active_clients_store.insert_remote( + registration_details.address, + mix_sender, + is_active_request_sender, + registration_details.session_request_timestamp, + ); + + return AuthenticatedHandler::upgrade( + self, + registration_details, + mix_receiver, + is_active_request_receiver, + ) + .await + .inspect_err(|err| error!("failed to upgrade client handler: {err}")) + .ok(); } - }; - - 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.shared_state.active_clients_store.insert_remote( - registration_details.address, - mix_sender, - is_active_request_sender, - registration_details.session_request_timestamp, - ); - - AuthenticatedHandler::upgrade( - self, - registration_details, - mix_receiver, - is_active_request_receiver, - ) - .await - .inspect_err(|err| error!("failed to upgrade client handler: {err}")) - .ok(); } - - None } pub(crate) async fn wait_for_initial_message(