Include token in device update event

4 files changed, 27 insertions, 11 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index a51c8bf2c6..181ff2f313 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -1383,7 +1383,7 @@ where
             self.schedule_reconnect(WG_RECONNECT_DELAY).await;
         }
         self.event_listener
-            .notify_device_event(DeviceEvent(Some(Device::from(event.0))));
+            .notify_device_event(DeviceEvent::from(event.0));
     }
 
     async fn handle_device_migration_event(&mut self, data: DeviceData) {
@@ -1391,11 +1391,10 @@ where
             // Discard stale device
             return;
         }
-        let device = data.device.clone();
+        let event = DeviceEvent::from(data.clone());
         self.account_manager.set(data);
         self.reconnect_tunnel();
-        self.event_listener
-            .notify_device_event(DeviceEvent(Some(device)));
+        self.event_listener.notify_device_event(event);
     }
 
     #[cfg(windows)]
@@ -1658,7 +1657,7 @@ where
                     .await
                     .map_err(Error::LoginError)?;
                 self.event_listener
-                    .notify_device_event(DeviceEvent(Some(Device::from(device_data))));
+                    .notify_device_event(DeviceEvent::from(device_data));
             }
             None => {
                 self.account_manager.logout();
@@ -2302,9 +2301,8 @@ where
     async fn on_rotate_wireguard_key(&mut self, tx: ResponseTx<(), Error>) {
         let result = self.account_manager.rotate_key().await;
         if let Ok(ref _wg_data) = result {
-            let device = self.account_manager.get().map(Device::from);
-            self.event_listener
-                .notify_device_event(DeviceEvent(device.clone()));
+            let device = DeviceEvent::from(self.account_manager.get());
+            self.event_listener.notify_device_event(device);
         }
         let _ = tx.send(result.map(|_| ()).map_err(Error::KeyRotationError));
     }
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index 4ad231774c..1196ab2381 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -543,5 +543,5 @@ message DeviceRemoval {
 }
 
 message DeviceEvent {
-	Device device = 1;
+	DeviceConfig device = 1;
 }
diff --git a/mullvad-management-interface/src/types.rs b/mullvad-management-interface/src/types.rs
index c8e8189c0a..0d78cbc9c5 100644
--- a/mullvad-management-interface/src/types.rs
+++ b/mullvad-management-interface/src/types.rs
@@ -209,7 +209,10 @@ impl From<mullvad_types::device::Device> for Device {
 impl From<mullvad_types::device::DeviceEvent> for DeviceEvent {
     fn from(event: mullvad_types::device::DeviceEvent) -> Self {
         DeviceEvent {
-            device: event.0.map(|device| Device::from(device)),
+            device: event.0.map(|(account_token, device)| DeviceConfig {
+                account_token,
+                device: Some(Device::from(device)),
+            }),
         }
     }
 }
diff --git a/mullvad-types/src/device.rs b/mullvad-types/src/device.rs
index e40a3d7080..b795e0e47b 100644
--- a/mullvad-types/src/device.rs
+++ b/mullvad-types/src/device.rs
@@ -34,4 +34,19 @@ impl From<DeviceData> for Device {
 
 /// Emitted when logging in or out of an account, or when the device changes.
 #[derive(Clone, Debug)]
-pub struct DeviceEvent(pub Option<Device>);
+pub struct DeviceEvent(pub Option<(AccountToken, Device)>);
+
+impl From<DeviceData> for DeviceEvent {
+    fn from(data: DeviceData) -> DeviceEvent {
+        DeviceEvent(Some((data.token, data.device)))
+    }
+}
+
+impl From<Option<DeviceData>> for DeviceEvent {
+    fn from(data: Option<DeviceData>) -> DeviceEvent {
+        match data {
+            Some(data) => DeviceEvent::from(data),
+            None => DeviceEvent(None),
+        }
+    }
+}