Migrate old settings to device cache

6 files changed, 197 insertions, 31 deletions

diff --git a/mullvad-daemon/src/device.rs b/mullvad-daemon/src/device.rs
index 42c8ee23bf..0e934626ef 100644
--- a/mullvad-daemon/src/device.rs
+++ b/mullvad-daemon/src/device.rs
@@ -313,7 +313,7 @@ pub struct DeviceService {
 }
 
 impl DeviceService {
-    fn new(handle: rest::MullvadRestHandle, api_availability: ApiAvailabilityHandle) -> Self {
+    pub fn new(handle: rest::MullvadRestHandle, api_availability: ApiAvailabilityHandle) -> Self {
         Self {
             proxy: DevicesProxy::new(handle),
             api_availability,
@@ -511,6 +511,36 @@ impl DeviceService {
         .await
         .map_err(Error::RestError)
     }
+
+    pub async fn list_devices_with_backoff(
+        &self,
+        token: AccountToken,
+    ) -> Result<Vec<Device>, Error> {
+        let proxy = self.proxy.clone();
+        let api_handle = self.api_availability.clone();
+
+        let retry_strategy = Jittered::jitter(
+            ExponentialBackoff::new(
+                RETRY_BACKOFF_INTERVAL_INITIAL,
+                RETRY_BACKOFF_INTERVAL_FACTOR,
+            )
+            .max_delay(RETRY_BACKOFF_INTERVAL_MAX),
+        );
+        retry_future(
+            move || {
+                let wait_online = api_handle.wait_online();
+                let fut = proxy.list(token.clone());
+                async move {
+                    let _ = wait_online.await;
+                    fut.await
+                }
+            },
+            should_retry_backoff,
+            retry_strategy,
+        )
+        .await
+        .map_err(Error::RestError)
+    }
 }
 
 pub struct DeviceCacher {
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 8d38eb8f54..c51cb8510e 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -345,6 +345,8 @@ pub(crate) enum InternalDaemonEvent {
     GenerateApiConnectionMode(api::ApiConnectionModeRequest),
     /// Sent when a device key is rotated.
     DeviceKeyEvent(device::DeviceKeyEvent),
+    /// Handles updates from versions without devices.
+    DeviceMigrationEvent(DeviceData),
     /// The split tunnel paths or state were updated.
     #[cfg(target_os = "windows")]
     ExcludedPathsEvent(ExcludedPathsUpdate, oneshot::Sender<Result<(), Error>>),
@@ -603,18 +605,6 @@ where
 
         let (internal_event_tx, internal_event_rx) = command_channel.destructure();
 
-        if let Err(error) = migrations::migrate_all(&cache_dir, &settings_dir).await {
-            log::error!(
-                "{}",
-                error.display_chain_with_msg("Failed to migrate settings or cache")
-            );
-        }
-        let settings = SettingsPersister::load(&settings_dir).await;
-
-        let tunnel_parameters_generator = MullvadTunnelParametersGenerator {
-            tx: internal_event_tx.clone(),
-        };
-
         let rpc_runtime = mullvad_rpc::MullvadRpcRuntime::with_cache(
             &cache_dir,
             true,
@@ -637,6 +627,25 @@ where
             .mullvad_rest_handle(proxy_provider, endpoint_updater.callback())
             .await;
 
+        if let Err(error) = migrations::migrate_all(
+            &cache_dir,
+            &settings_dir,
+            rpc_handle.clone(),
+            internal_event_tx.clone(),
+        )
+        .await
+        {
+            log::error!(
+                "{}",
+                error.display_chain_with_msg("Failed to migrate settings or cache")
+            );
+        }
+        let settings = SettingsPersister::load(&settings_dir).await;
+
+        let tunnel_parameters_generator = MullvadTunnelParametersGenerator {
+            tx: internal_event_tx.clone(),
+        };
+
         let mut account_manager = device::AccountManager::new(
             runtime.clone(),
             rpc_handle.clone(),
@@ -909,6 +918,7 @@ where
                 self.handle_generate_api_connection_mode(request).await
             }
             DeviceKeyEvent(event) => self.handle_device_key_event(event).await,
+            DeviceMigrationEvent(event) => self.handle_device_migration_event(event).await,
             #[cfg(windows)]
             ExcludedPathsEvent(update, tx) => self.handle_new_excluded_paths(update, tx).await,
         }
@@ -1374,6 +1384,18 @@ where
             .notify_device_event(DeviceEvent(Some(Device::from(event.0))));
     }
 
+    async fn handle_device_migration_event(&mut self, data: DeviceData) {
+        if self.account_manager.get().is_some() {
+            // Discard stale device
+            return;
+        }
+        let device = data.device.clone();
+        self.account_manager.set(data);
+        self.reconnect_tunnel();
+        self.event_listener
+            .notify_device_event(DeviceEvent(Some(device)));
+    }
+
     #[cfg(windows)]
     async fn handle_new_excluded_paths(
         &mut self,
diff --git a/mullvad-daemon/src/migrations/mod.rs b/mullvad-daemon/src/migrations/mod.rs
index 98ad71c23c..8347b3cd76 100644
--- a/mullvad-daemon/src/migrations/mod.rs
+++ b/mullvad-daemon/src/migrations/mod.rs
@@ -87,7 +87,12 @@ pub enum Error {
 
 pub type Result<T> = std::result::Result<T, Error>;
 
-pub async fn migrate_all(cache_dir: &Path, settings_dir: &Path) -> Result<()> {
+pub(crate) async fn migrate_all(
+    cache_dir: &Path,
+    settings_dir: &Path,
+    rest_handle: mullvad_rpc::rest::MullvadRestHandle,
+    daemon_tx: crate::DaemonEventSender,
+) -> Result<()> {
     #[cfg(windows)]
     windows::migrate_after_windows_update(settings_dir)
         .await
@@ -114,11 +119,12 @@ pub async fn migrate_all(cache_dir: &Path, settings_dir: &Path) -> Result<()> {
     v2::migrate(&mut settings)?;
     v3::migrate(&mut settings)?;
     v4::migrate(&mut settings)?;
-    v5::migrate(&mut settings)?;
 
     account_history::migrate_location(cache_dir, settings_dir).await;
     account_history::migrate_formats(settings_dir, &mut settings).await?;
 
+    v5::migrate(&mut settings, rest_handle, daemon_tx).await?;
+
     if settings == old_settings {
         // Nothing changed
         return Ok(());
diff --git a/mullvad-daemon/src/migrations/v5.rs b/mullvad-daemon/src/migrations/v5.rs
index 0fcaca4e08..0695ee8c7e 100644
--- a/mullvad-daemon/src/migrations/v5.rs
+++ b/mullvad-daemon/src/migrations/v5.rs
@@ -1,5 +1,10 @@
 use super::{Error, Result};
-use mullvad_types::settings::SettingsVersion;
+use crate::{device::DeviceService, DaemonEventSender, InternalDaemonEvent};
+use mullvad_types::{
+    account::AccountToken, device::DeviceData, settings::SettingsVersion, wireguard::WireguardData,
+};
+use talpid_core::mpsc::Sender;
+use talpid_types::ErrorExt;
 
 // ======================================================
 // Section for vendoring types and values that
@@ -21,16 +26,48 @@ use mullvad_types::settings::SettingsVersion;
 ///  * `use_mulithop` was not present in the settings
 ///  * A multihop entry location had been previously specified.
 ///
-/// This change is backwards compatible since older daemons will just ignore `use_multihop` if
-/// present.
-///
 /// It is also no longer valid to have `entry_location` set to null. So remove the field if it
 /// is null in order to make it default back to the default location.
-pub fn migrate(settings: &mut serde_json::Value) -> Result<()> {
-    if !version_matches(settings) {
-        return Ok(());
+///
+/// This also removes the account token and WireGuard key from the settings, looks up the
+/// corresponding device, and eventually stores them in `device.json` instead. This is done by
+/// sending the `DeviceMigrationEvent` event to the daemon. Because this is fallible, it can
+/// result in the account token and private key being lost. This should not be not critical since
+/// the account token is also stored in the account history.
+pub(crate) async fn migrate(
+    settings: &mut serde_json::Value,
+    rest_handle: mullvad_rpc::rest::MullvadRestHandle,
+    daemon_tx: DaemonEventSender,
+) -> Result<()> {
+    let migration_data = migrate_inner(settings).await?;
+
+    if let Some(migration_data) = migration_data {
+        let api_handle = rest_handle.availability.clone();
+        let service = DeviceService::new(rest_handle, api_handle);
+        match (migration_data.token, migration_data.wg_data) {
+            (token, Some(wg_data)) => {
+                log::info!("Creating a new device cache from previous settings");
+                tokio::spawn(cache_from_wireguard_key(daemon_tx, service, token, wg_data));
+            }
+            (token, None) => {
+                log::info!("Generating a new device for the account");
+                tokio::spawn(cache_from_account(daemon_tx, service, token));
+            }
+        }
     }
 
+    Ok(())
+}
+
+struct MigrationData {
+    token: AccountToken,
+    wg_data: Option<WireguardData>,
+}
+
+async fn migrate_inner(settings: &mut serde_json::Value) -> Result<Option<MigrationData>> {
+    if !version_matches(settings) {
+        return Ok(None);
+    }
     let wireguard_constraints = || -> Option<&serde_json::Value> {
         settings
             .get("relay_settings")?
@@ -54,11 +91,35 @@ pub fn migrate(settings: &mut serde_json::Value) -> Result<()> {
         }
     }
 
+    if let Some(token) = settings.get("account_token").filter(|t| !t.is_null()) {
+        let token: AccountToken =
+            serde_json::from_value(token.clone()).map_err(Error::ParseError)?;
+        let mig_data = if let Some(wg_data) = settings.get("wireguard").filter(|wg| !wg.is_null()) {
+            let wg_data: WireguardData =
+                serde_json::from_value(wg_data.clone()).map_err(Error::ParseError)?;
+            Ok(Some(MigrationData {
+                token,
+                wg_data: Some(wg_data),
+            }))
+        } else {
+            Ok(Some(MigrationData {
+                token,
+                wg_data: None,
+            }))
+        };
+
+        let settings_map = settings.as_object_mut().ok_or(Error::NoMatchingVersion)?;
+        settings_map.remove("account_token");
+        settings_map.remove("wireguard");
+
+        return mig_data;
+    }
+
     // Note: Not incrementing the version number yet, since this migration is still open
     // for future modification.
     // settings["settings_version"] = serde_json::json!(SettingsVersion::V6);
 
-    Ok(())
+    Ok(None)
 }
 
 fn version_matches(settings: &mut serde_json::Value) -> bool {
@@ -68,9 +129,56 @@ fn version_matches(settings: &mut serde_json::Value) -> bool {
         .unwrap_or(false)
 }
 
+async fn cache_from_wireguard_key(
+    daemon_tx: DaemonEventSender,
+    service: DeviceService,
+    token: AccountToken,
+    wg_data: WireguardData,
+) {
+    let devices = match service.list_devices_with_backoff(token.clone()).await {
+        Ok(devices) => devices,
+        Err(error) => {
+            log::error!(
+                "{}",
+                error.display_chain_with_msg("Failed to enumerate devices for account")
+            );
+            return;
+        }
+    };
+
+    for device in devices.into_iter() {
+        if device.pubkey == wg_data.private_key.public_key() {
+            let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(DeviceData {
+                token,
+                device,
+                wg_data,
+            }));
+            return;
+        }
+    }
+    log::info!("The existing WireGuard key is not valid; generating a new device");
+    cache_from_account(daemon_tx, service, token).await;
+}
+
+async fn cache_from_account(
+    daemon_tx: DaemonEventSender,
+    service: DeviceService,
+    token: AccountToken,
+) {
+    match service.generate_for_account_with_backoff(token).await {
+        Ok(device_data) => {
+            let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(device_data));
+        }
+        Err(error) => log::error!(
+            "{}",
+            error.display_chain_with_msg("Failed to generate new device for account")
+        ),
+    }
+}
+
 #[cfg(test)]
 mod test {
-    use super::{migrate, version_matches};
+    use super::{migrate_inner, version_matches};
     use serde_json;
 
     pub const V5_SETTINGS_V1: &str = r#"
@@ -144,7 +252,6 @@ mod test {
 
     pub const V5_SETTINGS_V2: &str = r#"
 {
-  "account_token": "1234",
   "relay_settings": {
     "normal": {
       "location": {
@@ -212,13 +319,12 @@ mod test {
 }
 "#;
 
-    #[test]
-    fn test_v5_v1_migration() {
+    #[tokio::test]
+    async fn test_v5_v1_migration() {
         let mut old_settings = serde_json::from_str(V5_SETTINGS_V1).unwrap();
 
         assert!(version_matches(&mut old_settings));
-
-        migrate(&mut old_settings).unwrap();
+        migrate_inner(&mut old_settings).await.unwrap();
         let new_settings: serde_json::Value = serde_json::from_str(V5_SETTINGS_V2).unwrap();
 
         assert_eq!(&old_settings, &new_settings);
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index a49e392320..43c5f362e1 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -420,6 +420,7 @@ pub struct DevicesProxy {
 struct DeviceResponse {
     id: DeviceId,
     name: DeviceName,
+    pubkey: wireguard::PublicKey,
     ipv4_address: ipnetwork::Ipv4Network,
     ipv6_address: ipnetwork::Ipv6Network,
 }
@@ -467,13 +468,14 @@ impl DevicesProxy {
             let DeviceResponse {
                 id,
                 name,
+                pubkey,
                 ipv4_address,
                 ipv6_address,
                 ..
             } = response;
 
             Ok((
-                Device { id, name },
+                Device { id, name, pubkey },
                 mullvad_types::wireguard::AssociatedAddresses {
                     ipv4_address,
                     ipv6_address,
diff --git a/mullvad-rpc/src/rest.rs b/mullvad-rpc/src/rest.rs
index c7e5d02fb1..21599fc5a0 100644
--- a/mullvad-rpc/src/rest.rs
+++ b/mullvad-rpc/src/rest.rs
@@ -583,7 +583,7 @@ pub async fn handle_error_response<T>(response: Response) -> Result<T> {
 pub struct MullvadRestHandle {
     pub(crate) service: RequestServiceHandle,
     pub factory: RequestFactory,
-    availability: ApiAvailabilityHandle,
+    pub availability: ApiAvailabilityHandle,
     pub token_store: AccessTokenProxy,
 }