Add RPC that checks whether the app is performing "post-upgrade work", i.e. fetching a device from the API

6 files changed, 91 insertions, 22 deletions

diff --git a/mullvad-daemon/src/device.rs b/mullvad-daemon/src/device.rs
index c68cae1845..23ce331502 100644
--- a/mullvad-daemon/src/device.rs
+++ b/mullvad-daemon/src/device.rs
@@ -77,8 +77,8 @@ pub enum Error {
     ParseDeviceCache(#[error(source)] serde_json::Error),
     #[error(display = "Unexpected HTTP request error")]
     OtherRestError(#[error(source)] rest::Error),
-    #[error(display = "The device update task is not running")]
-    DeviceUpdaterCancelled(#[error(source)] oneshot::Canceled),
+    #[error(display = "The task was aborted")]
+    Cancelled,
     #[error(display = "The account manager is down")]
     AccountManagerDown,
 }
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 031573c25d..9aca859d9e 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -293,6 +293,8 @@ pub enum DaemonCommand {
     GetWireguardKey(ResponseTx<Option<PublicKey>, Error>),
     /// Get information about the currently running and latest app versions
     GetVersionInfo(oneshot::Sender<Option<AppVersionInfo>>),
+    /// Return whether the daemon is performing post-upgrade tasks
+    IsPerformingPostUpgrade(oneshot::Sender<bool>),
     /// Get current version of the app
     GetCurrentVersion(oneshot::Sender<AppVersion>),
     /// Remove settings and clear the cache
@@ -359,7 +361,7 @@ pub(crate) enum InternalDaemonEvent {
     /// Sent when a device is updated in any way (key rotation, login, logout, etc.).
     DeviceEvent(InnerDeviceEvent),
     /// Handles updates from versions without devices.
-    DeviceMigrationEvent(DeviceData),
+    DeviceMigrationEvent(Result<DeviceData, device::Error>),
     /// The split tunnel paths or state were updated.
     #[cfg(target_os = "windows")]
     ExcludedPathsEvent(ExcludedPathsUpdate, oneshot::Sender<Result<(), Error>>),
@@ -580,6 +582,7 @@ pub struct Daemon<L: EventListener> {
     tx: DaemonEventSender,
     reconnection_job: Option<AbortHandle>,
     event_listener: L,
+    migration_complete: migrations::MigrationComplete,
     settings: SettingsPersister,
     account_history: account_history::AccountHistory,
     device_checker: device::TunnelStateChangeHandler,
@@ -642,19 +645,20 @@ where
             .mullvad_rest_handle(proxy_provider, endpoint_updater.callback())
             .await;
 
-        if let Err(error) = migrations::migrate_all(
+        let migration_complete = migrations::migrate_all(
             &cache_dir,
             &settings_dir,
             api_handle.clone(),
             internal_event_tx.clone(),
         )
         .await
-        {
+        .unwrap_or_else(|error| {
             log::error!(
                 "{}",
                 error.display_chain_with_msg("Failed to migrate settings or cache")
             );
-        }
+            migrations::MigrationComplete::new(true)
+        });
         let settings = SettingsPersister::load(&settings_dir).await;
 
         let tunnel_parameters_generator = MullvadTunnelParametersGenerator {
@@ -781,6 +785,7 @@ where
             tx: internal_event_tx,
             reconnection_job: None,
             event_listener,
+            migration_complete,
             settings,
             account_history,
             device_checker: device::TunnelStateChangeHandler::new(account_manager.clone()),
@@ -1318,6 +1323,7 @@ where
             RotateWireguardKey(tx) => self.on_rotate_wireguard_key(tx).await,
             GetWireguardKey(tx) => self.on_get_wireguard_key(tx).await,
             GetVersionInfo(tx) => self.on_get_version_info(tx).await,
+            IsPerformingPostUpgrade(tx) => self.on_is_performing_post_upgrade(tx).await,
             GetCurrentVersion(tx) => self.on_get_current_version(tx),
             #[cfg(not(target_os = "android"))]
             FactoryReset(tx) => self.on_factory_reset(tx).await,
@@ -1456,16 +1462,22 @@ where
             .notify_device_event(DeviceEvent::from(event));
     }
 
-    async fn handle_device_migration_event(&mut self, data: DeviceData) {
+    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;
         }
-        if let Err(error) = self.account_manager.set(data).await {
+
+        let result = async { self.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));
         }
     }
 
@@ -1522,6 +1534,11 @@ where
         Self::oneshot_send(tx, self.tunnel_state.clone(), "current state");
     }
 
+    async fn on_is_performing_post_upgrade(&self, tx: oneshot::Sender<bool>) {
+        let performing_post_upgrade = !self.migration_complete.is_complete();
+        Self::oneshot_send(tx, performing_post_upgrade, "performing post upgrade");
+    }
+
     async fn on_get_current_location(&mut self, tx: oneshot::Sender<Option<GeoIpLocation>>) {
         use self::TunnelState::*;
 
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 1d645ae8d7..0e581055a9 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -156,6 +156,13 @@ impl ManagementService for ManagementServiceImpl {
             .map(Response::new)
     }
 
+    async fn is_performing_post_upgrade(&self, _: Request<()>) -> ServiceResult<bool> {
+        log::debug!("is_performing_post_upgrade");
+        let (tx, rx) = oneshot::channel();
+        self.send_command_to_daemon(DaemonCommand::IsPerformingPostUpgrade(tx))?;
+        Ok(Response::new(self.wait_for_result(rx).await?))
+    }
+
     // Relays and tunnel constraints
     //
 
diff --git a/mullvad-daemon/src/migrations/device.rs b/mullvad-daemon/src/migrations/device.rs
index cc7b0cd5c9..dd9ac8e94e 100644
--- a/mullvad-daemon/src/migrations/device.rs
+++ b/mullvad-daemon/src/migrations/device.rs
@@ -6,17 +6,22 @@
 //! This module is allowed to import a number of types, unlike other migration modules, as it
 //! does not modify any files directly and may safely fail.
 
-use super::v5::MigrationData;
+use super::{v5::MigrationData, MigrationComplete};
 use crate::{
     device::{self, DeviceService},
     DaemonEventSender, InternalDaemonEvent,
 };
 use mullvad_types::{account::AccountToken, device::DeviceData, wireguard::WireguardData};
+use std::time::Duration;
 use talpid_core::mpsc::Sender;
 use talpid_types::ErrorExt;
+use tokio::time::timeout;
+
+const TIMEOUT: Duration = Duration::from_secs(30);
 
 pub(crate) fn generate_device(
     migration_data: MigrationData,
+    mut migration_complete: MigrationComplete,
     rest_handle: mullvad_api::rest::MullvadRestHandle,
     daemon_tx: DaemonEventSender,
 ) {
@@ -45,9 +50,8 @@ pub(crate) fn generate_device(
                 cache_from_account(service, token).await
             }
         };
-        if let Ok(data) = result {
-            let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(data));
-        }
+        let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(result));
+        migration_complete.set_complete();
     });
 }
 
@@ -56,9 +60,12 @@ async fn cache_from_wireguard_key(
     token: AccountToken,
     wg_data: WireguardData,
 ) -> Result<DeviceData, device::Error> {
-    let devices = service
-        .list_devices_with_backoff(token.clone())
+    let devices = timeout(TIMEOUT, service.list_devices_with_backoff(token.clone()))
         .await
+        .map_err(|_error| {
+            log::error!("Failed to enumerate devices for account: timed out");
+            device::Error::Cancelled
+        })?
         .map_err(|error| {
             log::error!(
                 "{}",
@@ -84,9 +91,12 @@ async fn cache_from_account(
     service: DeviceService,
     token: AccountToken,
 ) -> Result<DeviceData, device::Error> {
-    service
-        .generate_for_account_with_backoff(token)
+    timeout(TIMEOUT, service.generate_for_account_with_backoff(token))
         .await
+        .map_err(|_error| {
+            log::error!("Failed to generate new device for account: timed out");
+            device::Error::Cancelled
+        })?
         .map_err(|error| {
             log::error!(
                 "{}",
diff --git a/mullvad-daemon/src/migrations/mod.rs b/mullvad-daemon/src/migrations/mod.rs
index 5392f9828c..a280d55af5 100644
--- a/mullvad-daemon/src/migrations/mod.rs
+++ b/mullvad-daemon/src/migrations/mod.rs
@@ -31,7 +31,13 @@
 //! 1. Implement the migration and add adequate tests.
 //! 1. Add to the changelog: "Settings format updated to `vY`"
 
-use std::path::Path;
+use std::{
+    path::Path,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 use tokio::{
     fs,
     io::{self, AsyncWriteExt},
@@ -87,12 +93,31 @@ pub enum Error {
 
 pub type Result<T> = std::result::Result<T, Error>;
 
+/// Returns whether there is any background work remaining
+/// after `migrate_all` has returned.
+#[derive(Clone)]
+pub(crate) struct MigrationComplete(Arc<AtomicBool>);
+
+impl MigrationComplete {
+    pub fn new(state: bool) -> Self {
+        Self(Arc::new(AtomicBool::new(state)))
+    }
+
+    pub fn is_complete(&self) -> bool {
+        self.0.load(Ordering::Relaxed)
+    }
+
+    fn set_complete(&mut self) {
+        self.0.store(true, Ordering::Relaxed);
+    }
+}
+
 pub(crate) async fn migrate_all(
     cache_dir: &Path,
     settings_dir: &Path,
     rest_handle: mullvad_api::rest::MullvadRestHandle,
     daemon_tx: crate::DaemonEventSender,
-) -> Result<()> {
+) -> Result<MigrationComplete> {
     #[cfg(windows)]
     windows::migrate_after_windows_update(settings_dir)
         .await
@@ -101,7 +126,7 @@ pub(crate) async fn migrate_all(
     let path = settings_dir.join(SETTINGS_FILE);
 
     if !path.is_file() {
-        return Ok(());
+        return Ok(MigrationComplete::new(true));
     }
 
     let settings_bytes = fs::read(&path).await.map_err(Error::ReadError)?;
@@ -124,14 +149,22 @@ pub(crate) async fn migrate_all(
     account_history::migrate_formats(settings_dir, &mut settings).await?;
 
     let migration_data = v5::migrate(&mut settings).await?;
+    let mut migration_complete = MigrationComplete::new(false);
 
     if let Some(migration_data) = migration_data {
-        device::generate_device(migration_data, rest_handle, daemon_tx);
+        device::generate_device(
+            migration_data,
+            migration_complete.clone(),
+            rest_handle,
+            daemon_tx,
+        );
+    } else {
+        migration_complete.set_complete();
     }
 
     if settings == old_settings {
         // Nothing changed
-        return Ok(());
+        return Ok(migration_complete);
     }
 
     let buffer = serde_json::to_string_pretty(&settings).map_err(Error::SerializeError)?;
@@ -155,7 +188,7 @@ pub(crate) async fn migrate_all(
 
     log::debug!("Migrated settings. Wrote settings to {}", path.display());
 
-    Ok(())
+    Ok(migration_complete)
 }
 
 #[cfg(windows)]
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index cb11e0b2e0..887cc97d03 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -23,6 +23,8 @@ service ManagementService {
 	rpc GetCurrentVersion(google.protobuf.Empty) returns (google.protobuf.StringValue) {}
 	rpc GetVersionInfo(google.protobuf.Empty) returns (AppVersionInfo) {}
 
+	rpc IsPerformingPostUpgrade(google.protobuf.Empty) returns (google.protobuf.BoolValue) {}
+
 	// Relays and tunnel constraints
 	rpc UpdateRelayLocations(google.protobuf.Empty) returns (google.protobuf.Empty) {}
 	rpc UpdateRelaySettings(RelaySettingsUpdate) returns (google.protobuf.Empty) {}