Block on login when requesting account data

2 files changed, 41 insertions, 30 deletions

diff --git a/mullvad-daemon/src/device/mod.rs b/mullvad-daemon/src/device/mod.rs
index b2bb669eb2..65a872e24b 100644
--- a/mullvad-daemon/src/device/mod.rs
+++ b/mullvad-daemon/src/device/mod.rs
@@ -2,7 +2,6 @@ use chrono::{DateTime, Utc};
 use futures::{
     channel::{mpsc, oneshot},
     stream::StreamExt,
-    FutureExt,
 };
 
 use mullvad_api::{availability::ApiAvailabilityHandle, rest};
@@ -125,6 +124,7 @@ enum AccountManagerCommand {
     Logout(ResponseTx<()>),
     SetData(DeviceData, ResponseTx<()>),
     GetData(ResponseTx<Option<DeviceData>>),
+    GetDataAfterLogin(ResponseTx<Option<DeviceData>>),
     RotateKey(ResponseTx<()>),
     SetRotationInterval(RotationInterval, ResponseTx<()>),
     ValidateDevice(ResponseTx<()>),
@@ -160,6 +160,11 @@ impl AccountManagerHandle {
             .await
     }
 
+    pub async fn data_after_login(&self) -> Result<Option<DeviceData>, Error> {
+        self.send_command(|tx| AccountManagerCommand::GetDataAfterLogin(tx))
+            .await
+    }
+
     pub async fn rotate_key(&self) -> Result<(), Error> {
         self.send_command(|tx| AccountManagerCommand::RotateKey(tx))
             .await
@@ -214,6 +219,7 @@ pub(crate) struct AccountManager {
     last_validation: Option<SystemTime>,
     validation_requests: Vec<ResponseTx<()>>,
     rotation_requests: Vec<ResponseTx<()>>,
+    data_requests: Vec<ResponseTx<Option<DeviceData>>>,
 }
 
 impl AccountManager {
@@ -240,6 +246,7 @@ impl AccountManager {
             last_validation: None,
             validation_requests: vec![],
             rotation_requests: vec![],
+            data_requests: vec![],
         };
 
         tokio::spawn(manager.run(cmd_rx));
@@ -282,6 +289,13 @@ impl AccountManager {
                         Some(AccountManagerCommand::GetData(tx)) => {
                             let _ = tx.send(Ok(self.data.clone()));
                         }
+                        Some(AccountManagerCommand::GetDataAfterLogin(tx)) => {
+                            if current_api_call.is_logging_in() {
+                                self.data_requests.push(tx);
+                            } else {
+                                let _ = tx.send(Ok(self.data.clone()));
+                            }
+                        }
                         Some(AccountManagerCommand::RotateKey(tx)) => {
                             if current_api_call.is_logging_in() {
                                 let _ = tx.send(Err(Error::AccountChange));
@@ -303,9 +317,7 @@ impl AccountManager {
                         Some(AccountManagerCommand::SetRotationInterval(interval, tx)) => {
                             self.rotation_interval = interval;
                             if current_api_call.is_running_timed_totation() {
-                                if let Some(timed_rotation) = self.spawn_timed_key_rotation() {
-                                    current_api_call.set_timed_rotation(Box::pin(timed_rotation))
-                                }
+                                current_api_call.clear();
                             }
                             let _ = tx.send(Ok(()));
                         }
@@ -315,6 +327,7 @@ impl AccountManager {
                         Some(AccountManagerCommand::ReceiveEvents(events_tx, tx)) => {
                             let _ = tx.send(Ok(self.listeners.push(events_tx)));
                         },
+
                         None => {
                             break;
                         }
@@ -383,6 +396,8 @@ impl AccountManager {
         tx: ResponseTx<()>,
     ) {
         let _ = tx.send(async { self.set(InnerDeviceEvent::Login(device_response?)).await }.await);
+        let data = self.data.clone();
+        Self::drain_requests(&mut self.data_requests, || Ok(data.clone()));
     }
 
     async fn consume_validation(
@@ -532,6 +547,7 @@ impl AccountManager {
     }
 
     async fn logout(&mut self, tx: ResponseTx<()>) {
+        Self::drain_requests(&mut self.data_requests, || Err(Error::AccountChange));
         let data = match self.data.take() {
             Some(it) => it,
             _ => return,
@@ -544,19 +560,11 @@ impl AccountManager {
         self.listeners
             .retain(|listener| listener.send(InnerDeviceEvent::Logout).is_ok());
 
-        let mut logout_call = Box::pin(self.logout_api_call(data).fuse());
+        let logout_call = tokio::spawn(Box::pin(self.logout_api_call(data)));
+
         tokio::spawn(async move {
-            let timeout = tokio::time::sleep(LOGOUT_TIMEOUT).fuse();
-            futures::pin_mut!(timeout);
-            futures::select! {
-                _timeout = timeout => {
-                    let _ = tx.send(Ok(()));
-                    logout_call.await
-                },
-                _logout = logout_call => {
-                    let _ = tx.send(Ok(()));
-                }
-            }
+            let _response = tokio::time::timeout(LOGOUT_TIMEOUT, logout_call).await;
+            let _ = tx.send(Ok(()));
         });
     }
 
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index b4e1763b6e..e2e18cc6a1 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -1463,22 +1463,25 @@ where
     }
 
     async fn handle_device_migration_event(&mut self, result: Result<DeviceData, device::Error>) {
-        if let Ok(Some(_)) = self.account_manager.data().await {
-            // Discard stale device
-            return;
-        }
+        let account_manager = self.account_manager.clone();
+        let event_listener = self.event_listener.clone();
+        tokio::spawn(async move {
+            if let Ok(Some(_)) = account_manager.data_after_login().await {
+                // Discard stale device
+                return;
+            }
 
-        let result = async { self.account_manager.set(result?).await }.await;
+            let result = async { account_manager.set(result?).await }.await;
 
-        if let Err(error) = result {
-            log::error!(
-                "{}",
-                error.display_chain_with_msg("Failed to move over account from old settings")
-            );
-            // Synthesize a logout event.
-            self.event_listener
-                .notify_device_event(DeviceEvent::revoke(false));
-        }
+            if let Err(error) = result {
+                log::error!(
+                    "{}",
+                    error.display_chain_with_msg("Failed to move over account from old settings")
+                );
+                // Synthesize a logout event.
+                event_listener.notify_device_event(DeviceEvent::revoke(false));
+            }
+        });
     }
 
     #[cfg(windows)]