testing fixes

gateway/src/node/client_handling/websocket/connection_handler/fresh.rs

@@ -922,47 +922,54 @@ impl<R, S> FreshHandler<R, S> {
         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(