Merge branch 'refactor-device-removal'

3 files changed, 59 insertions, 68 deletions

diff --git a/mullvad-daemon/src/device/mod.rs b/mullvad-daemon/src/device/mod.rs
index dc43ce9ee1..ebf51c5210 100644
--- a/mullvad-daemon/src/device/mod.rs
+++ b/mullvad-daemon/src/device/mod.rs
@@ -4,7 +4,7 @@ use futures::{
     stream::StreamExt,
 };
 
-use mullvad_api::{availability::ApiAvailabilityHandle, rest};
+use mullvad_api::rest;
 use mullvad_types::{
     account::AccountToken,
     device::{AccountAndDevice, Device, DeviceEvent, DeviceId, DeviceName, DevicePort},
@@ -201,7 +201,6 @@ enum AccountManagerCommand {
     RotateKey(ResponseTx<()>),
     SetRotationInterval(RotationInterval, ResponseTx<()>),
     ValidateDevice(ResponseTx<()>),
-    ReceiveEvents(Box<dyn Sender<PrivateDeviceEvent> + Send>, ResponseTx<()>),
     Shutdown(oneshot::Sender<()>),
 }
 
@@ -253,16 +252,6 @@ impl AccountManagerHandle {
             .await
     }
 
-    pub async fn receive_events(
-        &self,
-        events_tx: impl Sender<PrivateDeviceEvent> + Send + 'static,
-    ) -> Result<(), Error> {
-        self.send_command(|tx| {
-            AccountManagerCommand::ReceiveEvents(Box::new(events_tx) as Box<_>, tx)
-        })
-        .await
-    }
-
     pub async fn shutdown(self) {
         let (tx, rx) = oneshot::channel();
         let _ = self
@@ -296,14 +285,17 @@ pub(crate) struct AccountManager {
 }
 
 impl AccountManager {
+    /// Starts the account manager actor and returns a handle to it as well as the
+    /// current device.
     pub async fn spawn(
         rest_handle: rest::MullvadRestHandle,
-        api_availability: ApiAvailabilityHandle,
         settings_dir: &Path,
         initial_rotation_interval: RotationInterval,
-    ) -> Result<AccountManagerHandle, Error> {
+        listener_tx: impl Sender<PrivateDeviceEvent> + Send + 'static,
+    ) -> Result<(AccountManagerHandle, Option<PrivateAccountAndDevice>), Error> {
         let (cacher, data) = DeviceCacher::new(settings_dir).await?;
         let token = data.as_ref().map(|state| state.account_token.clone());
+        let api_availability = rest_handle.availability.clone();
         let account_service =
             service::spawn_account_service(rest_handle.clone(), token, api_availability.clone());
 
@@ -313,9 +305,9 @@ impl AccountManager {
         let manager = AccountManager {
             cacher,
             device_service: device_service.clone(),
-            data,
+            data: data.clone(),
             rotation_interval: initial_rotation_interval,
-            listeners: vec![],
+            listeners: vec![Box::new(listener_tx)],
             last_validation: None,
             validation_requests: vec![],
             rotation_requests: vec![],
@@ -328,7 +320,7 @@ impl AccountManager {
             account_service,
             device_service,
         };
-        Ok(handle)
+        Ok((handle, data))
     }
 
     async fn run(mut self, mut cmd_rx: mpsc::UnboundedReceiver<AccountManagerCommand>) {
@@ -397,9 +389,6 @@ impl AccountManager {
                         }
                         Some(AccountManagerCommand::ValidateDevice(tx)) => {
                             self.handle_validation_request(tx, &mut current_api_call);
-                        }
-                        Some(AccountManagerCommand::ReceiveEvents(events_tx, tx)) => {
-                            let _ = tx.send(Ok(self.listeners.push(events_tx)));
                         },
 
                         None => {
@@ -862,6 +851,7 @@ impl DeviceCacher {
         let _ = tokio::task::spawn_blocking(move || drop(std_file)).await;
     }
 }
+
 /// Checks if the current device is valid if a WireGuard tunnel cannot be set up
 /// after multiple attempts.
 pub(crate) struct TunnelStateChangeHandler {
diff --git a/mullvad-daemon/src/device/service.rs b/mullvad-daemon/src/device/service.rs
index f511a35f67..c2b09d322a 100644
--- a/mullvad-daemon/src/device/service.rs
+++ b/mullvad-daemon/src/device/service.rs
@@ -105,7 +105,38 @@ impl DeviceService {
         })
     }
 
-    pub async fn remove_device(&self, token: AccountToken, device: DeviceId) -> Result<(), Error> {
+    pub async fn remove_device(
+        &self,
+        account_token: AccountToken,
+        device_id: DeviceId,
+    ) -> Result<(Device, Vec<Device>), Error> {
+        let mut devices = self.list_devices(account_token.clone()).await?;
+        self.remove_device_inner(account_token, device_id.clone())
+            .await?;
+        if let Some(index) = devices.iter().position(|device| device.id == device_id) {
+            Ok((devices.swap_remove(index), devices))
+        } else {
+            // You would only end up here if the API service successfully removed a device that
+            // was previously not included in the list returned by it, which should be impossible.
+            // Just return a bogus device in its place.
+            log::error!("List did not contain the revoked device");
+            Ok((
+                Device {
+                    id: device_id,
+                    name: "unknown device".to_string(),
+                    pubkey: talpid_types::net::wireguard::PublicKey::from([0u8; 32]),
+                    ports: vec![],
+                },
+                devices,
+            ))
+        }
+    }
+
+    async fn remove_device_inner(
+        &self,
+        token: AccountToken,
+        device: DeviceId,
+    ) -> Result<(), Error> {
         let proxy = self.proxy.clone();
         let api_handle = self.api_availability.clone();
         retry_future_n(
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 302a0f8f57..fee22f1d70 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -597,26 +597,18 @@ where
             migrations::MigrationComplete::new(true)
         };
 
-        let account_manager = device::AccountManager::spawn(
+        let (account_manager, data) = device::AccountManager::spawn(
             api_handle.clone(),
-            api_availability.clone(),
             &settings_dir,
             settings
                 .tunnel_options
                 .wireguard
                 .rotation_interval
                 .unwrap_or_default(),
+            internal_event_tx.to_specialized_sender(),
         )
         .await
         .map_err(Error::LoadAccountManager)?;
-        account_manager
-            .receive_events(internal_event_tx.to_specialized_sender())
-            .await
-            .map_err(Error::LoadAccountManager)?;
-        let data = account_manager
-            .data()
-            .await
-            .map_err(Error::LoadAccountManager)?;
 
         let account_history = account_history::AccountHistory::new(
             &settings_dir,
@@ -1376,50 +1368,28 @@ where
     async fn on_remove_device(
         &mut self,
         tx: ResponseTx<(), Error>,
-        token: AccountToken,
+        account_token: AccountToken,
         device_id: DeviceId,
     ) {
         let device_service = self.account_manager.device_service.clone();
         let event_listener = self.event_listener.clone();
 
         tokio::spawn(async move {
-            let mut devices = match device_service
-                .list_devices(token.clone())
+            let result = device_service
+                .remove_device(account_token.clone(), device_id)
                 .await
-                .map_err(Error::ListDevicesError)
-            {
-                Ok(devices) => devices,
-                Err(error) => {
-                    Self::oneshot_send(tx, Err(error), "remove_device response");
-                    return;
-                }
-            };
-            if let Err(error) = device_service
-                .remove_device(token.clone(), device_id.clone())
-                .await
-                .map_err(Error::RemoveDeviceError)
-            {
-                Self::oneshot_send(tx, Err(error), "remove_device response");
-                return;
-            };
-            let removed_device =
-                if let Some(index) = devices.iter().position(|device| device.id == device_id) {
-                    devices.swap_remove(index)
-                } else {
-                    log::error!("List did not contain the revoked device");
-                    Device {
-                        id: device_id,
-                        name: "unknown device".to_string(),
-                        pubkey: talpid_types::net::wireguard::PublicKey::from([0u8; 32]),
-                        ports: vec![],
-                    }
-                };
-            event_listener.notify_remove_device_event(RemoveDeviceEvent {
-                account_token: token,
-                removed_device,
-                new_devices: devices,
-            });
-            Self::oneshot_send(tx, Ok(()), "remove_device response");
+                .map(move |(removed_device, new_devices)| {
+                    event_listener.notify_remove_device_event(RemoveDeviceEvent {
+                        account_token,
+                        removed_device,
+                        new_devices,
+                    });
+                });
+            Self::oneshot_send(
+                tx,
+                result.map_err(Error::RemoveDeviceError),
+                "remove_device response",
+            );
         });
     }