From 84d54cfbaeba16652f34cd12d741093aab49b582 Mon Sep 17 00:00:00 2001
From: Davide Baldo <davide@baldo.me>
Date: Thu, 16 Nov 2023 16:42:17 +0100
Subject: [PATCH] fix(rust): fixed enrollment when not successfully completed
 on first try

---
 .../ockam_app_lib/src/enroll/enroll_user.rs   | 50 ++++++++++++-------
 1 file changed, 32 insertions(+), 18 deletions(-)

diff --git a/implementations/rust/ockam/ockam_app_lib/src/enroll/enroll_user.rs b/implementations/rust/ockam/ockam_app_lib/src/enroll/enroll_user.rs
index bc4a2f118a8..8df747a4167 100644
--- a/implementations/rust/ockam/ockam_app_lib/src/enroll/enroll_user.rs
+++ b/implementations/rust/ockam/ockam_app_lib/src/enroll/enroll_user.rs
@@ -24,17 +24,26 @@ impl AppState {
     pub async fn enroll_user(&self) -> Result<()> {
         let result = self.enroll_with_token().await;
 
-        if let Err(err) = result {
-            error!(?err, "Failed to enroll user");
-            self.update_orchestrator_status(OrchestratorStatus::Disconnected);
-            self.publish_state().await;
-            self.notify(Notification {
-                kind: Kind::Error,
-                title: "Failed to enroll user".to_string(),
-                message: format!("{}", err),
-            });
-            return Err(err);
+        let enrolled = match result {
+            Ok(enrolled) => enrolled,
+            Err(err) => {
+                error!(?err, "Failed to enroll user");
+                self.update_orchestrator_status(OrchestratorStatus::Disconnected);
+                self.publish_state().await;
+                self.notify(Notification {
+                    kind: Kind::Error,
+                    title: "Failed to enroll user".to_string(),
+                    message: format!("{}", err),
+                });
+                return Err(err);
+            }
+        };
+
+        if !enrolled {
+            // no error to report but the enrollment was interrupted
+            return Ok(());
         }
+
         // Reset the node manager to include the project's setup, needed to create the relay.
         // This is necessary because the project data is used in the worker initialization,
         // which can't be rerun manually once the worker is started.
@@ -61,10 +70,10 @@ impl AppState {
         Ok(())
     }
 
-    async fn enroll_with_token(&self) -> Result<()> {
+    async fn enroll_with_token(&self) -> Result<bool> {
         if self.is_enrolled().await.unwrap_or_default() {
             debug!("User is already enrolled");
-            return Ok(());
+            return Ok(false);
         }
 
         self.update_orchestrator_status(OrchestratorStatus::WaitingForToken);
@@ -77,10 +86,6 @@ impl AppState {
         // retrieve the user information
         let user_info = oidc_service.get_user_info(&token).await?;
         info!(?user_info, "User info retrieved successfully");
-        let cli_state = self.state().await;
-        cli_state
-            .users_info
-            .overwrite(&user_info.email, user_info.clone())?;
 
         if !user_info.email_verified {
             self.notify(Notification {
@@ -91,9 +96,18 @@ impl AppState {
                      Please review your inbox and follow the provided steps \
                      to complete the verification process"
                     .to_string(),
-            })
+            });
+            self.update_orchestrator_status(OrchestratorStatus::Disconnected);
+            self.publish_state().await;
+            return Ok(false);
         }
 
+        let cli_state = self.state().await;
+        cli_state
+            .users_info
+            .overwrite(&user_info.email, user_info.clone())?;
+        cli_state.users_info.set_default(&user_info.email)?;
+
         // enroll the current user using that token on the controller
         {
             let controller = self.controller().await.into_diagnostic()?;
@@ -120,7 +134,7 @@ impl AppState {
             message: "You can now use the Ockam app".to_string(),
         });
 
-        Ok(())
+        Ok(true)
     }
 
     async fn retrieve_space(&self) -> Result<Space> {