Move imports out of v5 module

3 files changed, 128 insertions, 106 deletions

diff --git a/mullvad-daemon/src/migrations/device.rs b/mullvad-daemon/src/migrations/device.rs
new file mode 100644
index 0000000000..cc7b0cd5c9
--- /dev/null
+++ b/mullvad-daemon/src/migrations/device.rs
@@ -0,0 +1,97 @@
+//! Generates a `device.json` from a WireGuard key and account token by matching them against
+//! devices returned by the API and sending the `DeviceMigrationEvent` event to the daemon.
+//! The account token and private key may be lost if it fails, but this should not be not
+//! critical since the account history also contains the token.
+//!
+//! 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 crate::{
+    device::{self, DeviceService},
+    DaemonEventSender, InternalDaemonEvent,
+};
+use mullvad_types::{account::AccountToken, device::DeviceData, wireguard::WireguardData};
+use talpid_core::mpsc::Sender;
+use talpid_types::ErrorExt;
+
+pub(crate) fn generate_device(
+    migration_data: MigrationData,
+    rest_handle: mullvad_api::rest::MullvadRestHandle,
+    daemon_tx: DaemonEventSender,
+) {
+    tokio::spawn(async move {
+        let wg_data: Option<WireguardData> = migration_data.wg_data.and_then(|data| {
+            serde_json::from_value(data)
+                .map(Some)
+                .unwrap_or_else(|error| {
+                    log::error!(
+                        "{}",
+                        error.display_chain_with_msg("Failed to parse WireGuard data")
+                    );
+                    None
+                })
+        });
+
+        let api_handle = rest_handle.availability.clone();
+        let service = DeviceService::new(rest_handle, api_handle);
+        let result = match (migration_data.token, wg_data) {
+            (token, Some(wg_data)) => {
+                log::info!("Creating a new device cache from previous settings");
+                cache_from_wireguard_key(service, token, wg_data).await
+            }
+            (token, None) => {
+                log::info!("Generating a new device for the account");
+                cache_from_account(service, token).await
+            }
+        };
+        if let Ok(data) = result {
+            let _ = daemon_tx.send(InternalDaemonEvent::DeviceMigrationEvent(data));
+        }
+    });
+}
+
+async fn cache_from_wireguard_key(
+    service: DeviceService,
+    token: AccountToken,
+    wg_data: WireguardData,
+) -> Result<DeviceData, device::Error> {
+    let devices = service
+        .list_devices_with_backoff(token.clone())
+        .await
+        .map_err(|error| {
+            log::error!(
+                "{}",
+                error.display_chain_with_msg("Failed to enumerate devices for account")
+            );
+            error
+        })?;
+
+    for device in devices.into_iter() {
+        if device.pubkey == wg_data.private_key.public_key() {
+            return Ok(DeviceData {
+                token,
+                device,
+                wg_data,
+            });
+        }
+    }
+    log::info!("The existing WireGuard key is not valid; generating a new device");
+    cache_from_account(service, token).await
+}
+
+async fn cache_from_account(
+    service: DeviceService,
+    token: AccountToken,
+) -> Result<DeviceData, device::Error> {
+    service
+        .generate_for_account_with_backoff(token)
+        .await
+        .map_err(|error| {
+            log::error!(
+                "{}",
+                error.display_chain_with_msg("Failed to generate new device for account")
+            );
+            error
+        })
+}
diff --git a/mullvad-daemon/src/migrations/mod.rs b/mullvad-daemon/src/migrations/mod.rs
index 9c33072305..98f3781061 100644
--- a/mullvad-daemon/src/migrations/mod.rs
+++ b/mullvad-daemon/src/migrations/mod.rs
@@ -38,6 +38,7 @@ use tokio::{
 };
 
 mod account_history;
+mod device;
 mod v1;
 mod v2;
 mod v3;
@@ -123,7 +124,11 @@ pub(crate) async fn migrate_all(
     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?;
+    let migration_data = v5::migrate(&mut settings).await?;
+
+    if let Some(migration_data) = migration_data {
+        device::generate_device(migration_data, rest_handle, daemon_tx);
+    }
 
     if settings == old_settings {
         // Nothing changed
diff --git a/mullvad-daemon/src/migrations/v5.rs b/mullvad-daemon/src/migrations/v5.rs
index c74bdc0ca4..4231e52a23 100644
--- a/mullvad-daemon/src/migrations/v5.rs
+++ b/mullvad-daemon/src/migrations/v5.rs
@@ -1,17 +1,19 @@
 use super::{Error, Result};
-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;
+use mullvad_types::settings::SettingsVersion;
 
 // ======================================================
 // Section for vendoring types and values that
 // this settings version depend on. See `mod.rs`.
 
+pub type AccountToken = String;
+
 // ======================================================
 
+pub(crate) struct MigrationData {
+    pub token: AccountToken,
+    pub wg_data: Option<serde_json::Value>,
+}
+
 /// This is an open ended migration. There is no v6 yet!
 /// The migrations performed by this function are still backwards compatible.
 /// The JSON coming out of this migration can be read by any v5 compatible daemon.
@@ -29,42 +31,8 @@ use talpid_types::ErrorExt;
 /// 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.
 ///
-/// 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_api::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>> {
+/// This also removes the account token and WireGuard key from the settings.
+pub(crate) async fn migrate(settings: &mut serde_json::Value) -> Result<Option<MigrationData>> {
     if !version_matches(settings) {
         return Ok(None);
     }
@@ -94,25 +62,24 @@ async fn migrate_inner(settings: &mut serde_json::Value) -> Result<Option<Migrat
     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 migration_data =
+            if let Some(wg_data) = settings.get("wireguard").filter(|wg| !wg.is_null()) {
+                Ok(Some(MigrationData {
+                    token,
+                    wg_data: Some(wg_data.clone()),
+                }))
+            } 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;
+        return migration_data;
     }
 
     // Note: Not incrementing the version number yet, since this migration is still open
@@ -129,56 +96,9 @@ 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_inner, version_matches};
+    use super::{migrate, version_matches};
     use serde_json;
 
     pub const V5_SETTINGS_V1: &str = r#"
@@ -324,7 +244,7 @@ mod test {
         let mut old_settings = serde_json::from_str(V5_SETTINGS_V1).unwrap();
 
         assert!(version_matches(&mut old_settings));
-        migrate_inner(&mut old_settings).await.unwrap();
+        migrate(&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);