Update the device before passing it to frontends

2 files changed, 52 insertions, 20 deletions

diff --git a/mullvad-daemon/src/device.rs b/mullvad-daemon/src/device.rs
index bc5069929a..de7bfd9a7f 100644
--- a/mullvad-daemon/src/device.rs
+++ b/mullvad-daemon/src/device.rs
@@ -80,8 +80,8 @@ impl Error {
 pub enum ValidationResult {
     /// The device and key were valid.
     Valid,
-    /// The device was valid but the key had to be replaced.
-    RotatedKey(WireguardData),
+    /// The device was valid but one or more fields, such as the key, were replaced
+    Updated,
     /// The device was not found remotely and was removed from the cache.
     Removed,
 }
@@ -317,20 +317,38 @@ impl AccountManager {
     }
 
     /// Check if the device is valid for the account, and yank it if it no longer exists.
+    /// This also updates any associated data and returns whether it changed.
     pub async fn validate_device(&mut self) -> Result<ValidationResult, Error> {
-        let data = {
+        let mut data = {
             let inner = self.inner.lock().unwrap();
             inner.data.as_ref().ok_or(Error::NoDevice)?.clone()
         };
 
-        match self.device_service.get(data.token, data.device.id).await {
+        match self
+            .device_service
+            .get(data.token.clone(), data.device.id.clone())
+            .await
+        {
             Ok(device) => {
                 if device.pubkey == data.device.pubkey {
-                    log::debug!("The current device is still valid");
-                    Ok(ValidationResult::Valid)
+                    if device == data.device {
+                        log::debug!("The current device is still valid");
+                        Ok(ValidationResult::Valid)
+                    } else {
+                        log::debug!("Updating data for the current device");
+                        data.device = device;
+                        {
+                            let mut inner = self.inner.lock().unwrap();
+                            inner.data.replace(data.clone());
+                            let (result_tx, _result_rx) = oneshot::channel();
+                            let _ = self.cache_update_tx.unbounded_send((Some(data), result_tx));
+                        }
+                        Ok(ValidationResult::Updated)
+                    }
                 } else {
                     log::debug!("Rotating invalid WireGuard key");
-                    Ok(ValidationResult::RotatedKey(self.rotate_key().await?))
+                    self.rotate_key().await?;
+                    Ok(ValidationResult::Updated)
                 }
             }
             Err(Error::InvalidAccount) | Err(Error::InvalidDevice) => {
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 404e9f9834..00cce7bfc0 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -1162,19 +1162,7 @@ where
                 {
                     match self.account_manager.validate_device().await {
                         Ok(status) => {
-                            match status {
-                                device::ValidationResult::Valid => (),
-                                device::ValidationResult::Removed => {
-                                    self.event_listener
-                                        .notify_device_event(DeviceEvent::revoke(true));
-                                }
-                                device::ValidationResult::RotatedKey(_) => {
-                                    self.event_listener.notify_device_event(DeviceEvent::new(
-                                        self.account_manager.get(),
-                                        true,
-                                    ));
-                                }
-                            }
+                            self.handle_validation_result(status);
                             self.wg_check_validity = false;
                         }
                         Err(error) => {
@@ -1217,6 +1205,21 @@ where
         }
     }
 
+    // Emit the appropriate events for an updated device.
+    fn handle_validation_result(&mut self, result: device::ValidationResult) {
+        match result {
+            device::ValidationResult::Valid => (),
+            device::ValidationResult::Removed => {
+                self.event_listener
+                    .notify_device_event(DeviceEvent::revoke(true));
+            }
+            device::ValidationResult::Updated => {
+                self.event_listener
+                    .notify_device_event(DeviceEvent::new(self.account_manager.get(), true));
+            }
+        }
+    }
+
     async fn schedule_reconnect(&mut self, delay: Duration) {
         self.unschedule_reconnect();
 
@@ -1737,6 +1740,17 @@ where
     }
 
     async fn on_get_device(&mut self, tx: ResponseTx<Option<DeviceConfig>, Error>) {
+        // Make sure the device is updated
+        match self.account_manager.validate_device().await {
+            Ok(status) => self.handle_validation_result(status),
+            Err(error) => {
+                log::error!(
+                    "{}",
+                    error.display_chain_with_msg("Failed to update device data")
+                );
+            }
+        }
+
         Self::oneshot_send(
             tx,
             Ok(self.account_manager.get().map(DeviceConfig::from)),