Implement correctly cancellable voucher submissions in the account manager

5 files changed, 110 insertions, 20 deletions

diff --git a/mullvad-daemon/src/device/api.rs b/mullvad-daemon/src/device/api.rs
index 6cd46ccc48..93dae78261 100644
--- a/mullvad-daemon/src/device/api.rs
+++ b/mullvad-daemon/src/device/api.rs
@@ -2,7 +2,7 @@ use std::pin::Pin;
 
 use chrono::{DateTime, Utc};
 use futures::{future::FusedFuture, Future};
-use mullvad_types::{device::Device, wireguard::WireguardData};
+use mullvad_types::{account::VoucherSubmission, device::Device, wireguard::WireguardData};
 
 use super::{Error, PrivateAccountAndDevice, ResponseTx};
 
@@ -39,6 +39,14 @@ impl CurrentApiCall {
         self.current_call = Some(Call::ExpiryCheck(expiry_call));
     }
 
+    pub fn set_voucher_submission(
+        &mut self,
+        voucher_call: ApiCall<VoucherSubmission>,
+        tx: ResponseTx<VoucherSubmission>,
+    ) {
+        self.current_call = Some(Call::VoucherSubmission(voucher_call, Some(tx)));
+    }
+
     pub fn is_validating(&self) -> bool {
         matches!(
             &self.current_call,
@@ -97,6 +105,10 @@ enum Call {
     TimerKeyRotation(ApiCall<WireguardData>),
     OneshotKeyRotation(ApiCall<WireguardData>),
     Validation(ApiCall<Device>),
+    VoucherSubmission(
+        ApiCall<VoucherSubmission>,
+        Option<ResponseTx<VoucherSubmission>>,
+    ),
     ExpiryCheck(ApiCall<DateTime<Utc>>),
 }
 
@@ -120,6 +132,16 @@ impl futures::Future for Call {
                 Pin::new(call).poll(cx).map(ApiResult::Rotation)
             }
             Validation(call) => Pin::new(call).poll(cx).map(ApiResult::Validation),
+            VoucherSubmission(call, tx) => {
+                if let std::task::Poll::Ready(response) = Pin::new(call).poll(cx) {
+                    std::task::Poll::Ready(ApiResult::VoucherSubmission(
+                        response,
+                        tx.take().unwrap(),
+                    ))
+                } else {
+                    std::task::Poll::Pending
+                }
+            }
             ExpiryCheck(call) => Pin::new(call).poll(cx).map(ApiResult::ExpiryCheck),
         }
     }
@@ -129,5 +151,9 @@ pub(crate) enum ApiResult {
     Login(Result<PrivateAccountAndDevice, Error>, ResponseTx<()>),
     Rotation(Result<WireguardData, Error>),
     Validation(Result<Device, Error>),
+    VoucherSubmission(
+        Result<VoucherSubmission, Error>,
+        ResponseTx<VoucherSubmission>,
+    ),
     ExpiryCheck(Result<DateTime<Utc>, Error>),
 }
diff --git a/mullvad-daemon/src/device/mod.rs b/mullvad-daemon/src/device/mod.rs
index 06d9a9b898..3121c91e7b 100644
--- a/mullvad-daemon/src/device/mod.rs
+++ b/mullvad-daemon/src/device/mod.rs
@@ -6,7 +6,7 @@ use futures::{
 
 use mullvad_api::rest;
 use mullvad_types::{
-    account::AccountToken,
+    account::{AccountToken, VoucherSubmission},
     device::{
         AccountAndDevice, Device, DeviceEvent, DeviceEventCause, DeviceId, DeviceName, DevicePort,
         DeviceState,
@@ -307,6 +307,7 @@ enum AccountManagerCommand {
     RotateKey(ResponseTx<()>),
     SetRotationInterval(RotationInterval, ResponseTx<()>),
     ValidateDevice(ResponseTx<()>),
+    SubmitVoucher(String, ResponseTx<VoucherSubmission>),
     CheckExpiry(ResponseTx<DateTime<Utc>>),
     Shutdown(oneshot::Sender<()>),
 }
@@ -356,6 +357,11 @@ impl AccountManagerHandle {
             .await
     }
 
+    pub async fn submit_voucher(&self, voucher: String) -> Result<VoucherSubmission, Error> {
+        self.send_command(move |tx| AccountManagerCommand::SubmitVoucher(voucher, tx))
+            .await
+    }
+
     pub async fn check_expiry(&self) -> Result<DateTime<Utc>, Error> {
         self.send_command(AccountManagerCommand::CheckExpiry).await
     }
@@ -502,6 +508,9 @@ impl AccountManager {
                         Some(AccountManagerCommand::ValidateDevice(tx)) => {
                             self.handle_validation_request(tx, &mut current_api_call);
                         },
+                        Some(AccountManagerCommand::SubmitVoucher(voucher, tx)) => {
+                            self.handle_voucher_submission(tx, voucher, &mut current_api_call);
+                        },
                         Some(AccountManagerCommand::CheckExpiry(tx)) => {
                             self.handle_expiry_request(tx, &mut current_api_call);
                         },
@@ -555,6 +564,34 @@ impl AccountManager {
         }
     }
 
+    fn handle_voucher_submission(
+        &mut self,
+        tx: ResponseTx<VoucherSubmission>,
+        voucher: String,
+        current_api_call: &mut api::CurrentApiCall,
+    ) {
+        if current_api_call.is_logging_in() {
+            let _ = tx.send(Err(Error::AccountChange));
+            return;
+        }
+
+        let create_submission = move || {
+            let old_config = self.data.device().ok_or(Error::NoDevice)?;
+            let account_token = old_config.account_token.clone();
+            let account_service = self.account_service.clone();
+            Ok(async move { account_service.submit_voucher(account_token, voucher).await })
+        };
+
+        match create_submission() {
+            Ok(call) => {
+                current_api_call.set_voucher_submission(Box::pin(call), tx);
+            }
+            Err(err) => {
+                let _ = tx.send(Err(err));
+            }
+        }
+    }
+
     fn handle_expiry_request(
         &mut self,
         tx: ResponseTx<DateTime<Utc>>,
@@ -590,6 +627,9 @@ impl AccountManager {
             Login(data, tx) => self.consume_login(data, tx).await,
             Rotation(rotation_response) => self.consume_rotation_result(rotation_response).await,
             Validation(data_response) => self.consume_validation(data_response, api_call).await,
+            VoucherSubmission(data_response, tx) => {
+                self.consume_voucher_result(data_response, tx).await
+            }
             ExpiryCheck(data_response) => self.consume_expiry_result(data_response).await,
         }
     }
@@ -605,6 +645,29 @@ impl AccountManager {
         Self::drain_requests(&mut self.data_requests, || Ok(data.clone()));
     }
 
+    async fn consume_voucher_result(
+        &mut self,
+        response: Result<VoucherSubmission, Error>,
+        tx: ResponseTx<VoucherSubmission>,
+    ) {
+        match &response {
+            Ok(submission) => {
+                // Send expiry update event
+                let event = PrivateDeviceEvent::AccountExpiry(submission.new_expiry);
+                self.listeners
+                    .retain(|listener| listener.send(event.clone()).is_ok());
+            }
+            Err(Error::InvalidAccount) => {
+                self.revoke_device(|| Error::InvalidAccount).await;
+            }
+            Err(Error::InvalidDevice) => {
+                self.revoke_device(|| Error::InvalidDevice).await;
+            }
+            Err(err) => log::error!("Failed to submit voucher: {}", err),
+        }
+        let _ = tx.send(response);
+    }
+
     async fn consume_expiry_result(&mut self, response: Result<DateTime<Utc>, Error>) {
         match response {
             Ok(expiry) => {
@@ -612,6 +675,7 @@ impl AccountManager {
                 let event = PrivateDeviceEvent::AccountExpiry(expiry);
                 self.listeners
                     .retain(|listener| listener.send(event.clone()).is_ok());
+
                 Self::drain_requests(&mut self.expiry_requests, || Ok(expiry));
             }
             Err(Error::InvalidAccount) => {
diff --git a/mullvad-daemon/src/device/service.rs b/mullvad-daemon/src/device/service.rs
index 61cbe0e4bd..294e41bab5 100644
--- a/mullvad-daemon/src/device/service.rs
+++ b/mullvad-daemon/src/device/service.rs
@@ -315,10 +315,10 @@ impl AccountService {
     }
 
     pub async fn submit_voucher(
-        &mut self,
+        &self,
         account_token: AccountToken,
         voucher: String,
-    ) -> Result<VoucherSubmission, rest::Error> {
+    ) -> Result<VoucherSubmission, Error> {
         let mut proxy = self.proxy.clone();
         let api_handle = self.api_availability.clone();
         let result = retry_future_n(
@@ -332,7 +332,7 @@ impl AccountService {
             self.initial_check_abort_handle.abort();
             self.api_availability.resume_background();
         }
-        result
+        result.map_err(map_rest_error)
     }
 }
 
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 4d65e3dc57..cfbcbb322c 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -126,6 +126,9 @@ pub enum Error {
     #[error(display = "Failed to update device")]
     UpdateDeviceError(#[error(source)] device::Error),
 
+    #[error(display = "Failed to submit voucher")]
+    VoucherSubmission(#[error(source)] device::Error),
+
     #[cfg(target_os = "linux")]
     #[error(display = "Unable to initialize split tunneling")]
     InitSplitTunneling(#[error(source)] split_tunnel::Error),
@@ -1315,21 +1318,17 @@ where
         tx: ResponseTx<VoucherSubmission, Error>,
         voucher: String,
     ) {
-        if let Ok(Some(device)) = self.account_manager.data().await.map(|s| s.into_device()) {
-            let mut account = self.account_manager.account_service.clone();
-            tokio::spawn(async move {
-                Self::oneshot_send(
-                    tx,
-                    account
-                        .submit_voucher(device.account_token, voucher)
-                        .await
-                        .map_err(Error::RestError),
-                    "submit_voucher response",
-                );
-            });
-        } else {
-            Self::oneshot_send(tx, Err(Error::NoAccountToken), "submit_voucher response");
-        }
+        let manager = self.account_manager.clone();
+        tokio::spawn(async move {
+            Self::oneshot_send(
+                tx,
+                manager
+                    .submit_voucher(voucher)
+                    .await
+                    .map_err(Error::VoucherSubmission),
+                "submit_voucher response",
+            );
+        });
     }
 
     fn on_get_relay_locations(&mut self, tx: oneshot::Sender<RelayList>) {
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index 68ceed4c16..a8c4484be8 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -954,6 +954,7 @@ fn map_daemon_error(error: crate::Error) -> Status {
         DaemonError::ListDevicesError(error) => map_device_error(&error),
         DaemonError::RemoveDeviceError(error) => map_device_error(&error),
         DaemonError::UpdateDeviceError(error) => map_device_error(&error),
+        DaemonError::VoucherSubmission(error) => map_device_error(&error),
         #[cfg(windows)]
         DaemonError::SplitTunnelError(error) => map_split_tunnel_error(error),
         DaemonError::AccountHistory(error) => map_account_history_error(error),