Retry key removals when switching between accounts or doing factory

reset

4 files changed, 63 insertions, 16 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 5e4b9453a3..15733d2f29 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -1591,7 +1591,7 @@ where
                 {
                     let remove_key = self
                         .wireguard_key_manager
-                        .remove_key(previous_token, previous_key);
+                        .remove_key_with_backoff(previous_token, previous_key);
                     tokio::spawn(async move {
                         if let Err(error) = remove_key.await {
                             log::error!(
diff --git a/mullvad-daemon/src/wireguard.rs b/mullvad-daemon/src/wireguard.rs
index 117bcf33c2..84aecc911b 100644
--- a/mullvad-daemon/src/wireguard.rs
+++ b/mullvad-daemon/src/wireguard.rs
@@ -9,8 +9,10 @@ pub use mullvad_types::wireguard::*;
 use std::{future::Future, pin::Pin, time::Duration};
 
 use futures::future::{abortable, AbortHandle};
+#[cfg(not(target_os = "android"))]
+use talpid_core::future_retry::constant_interval;
 use talpid_core::{
-    future_retry::{retry_future, ExponentialBackoff, Jittered},
+    future_retry::{retry_future, retry_future_n, ExponentialBackoff, Jittered},
     mpsc::Sender,
 };
 
@@ -30,6 +32,11 @@ const RETRY_INTERVAL_INITIAL: Duration = Duration::from_secs(4);
 const RETRY_INTERVAL_FACTOR: u32 = 5;
 const RETRY_INTERVAL_MAX: Duration = Duration::from_secs(24 * 60 * 60);
 
+#[cfg(not(target_os = "android"))]
+const SHORT_RETRY_INTERVAL: Duration = Duration::from_millis(500);
+
+const MAX_KEY_REMOVAL_RETRIES: usize = 2;
+
 #[derive(err_derive::Error, Debug)]
 pub enum Error {
     #[error(display = "Unexpected HTTP request error")]
@@ -138,17 +145,56 @@ impl KeyManager {
     }
 
     /// Removes a key from an account
+    #[cfg(not(target_os = "android"))]
     pub fn remove_key(
         &self,
         account: AccountToken,
         key: talpid_types::net::wireguard::PublicKey,
     ) -> impl Future<Output = Result<()>> {
+        self.remove_key_inner(account, key, constant_interval(SHORT_RETRY_INTERVAL), false)
+    }
+
+    /// Removes a key from an account
+    pub fn remove_key_with_backoff(
+        &self,
+        account: AccountToken,
+        key: talpid_types::net::wireguard::PublicKey,
+    ) -> impl Future<Output = Result<()>> {
+        let retry_strategy = Jittered::jitter(
+            ExponentialBackoff::new(RETRY_INTERVAL_INITIAL, RETRY_INTERVAL_FACTOR)
+                .max_delay(RETRY_INTERVAL_MAX),
+        );
+        self.remove_key_inner(account, key, retry_strategy, true)
+    }
+
+    fn remove_key_inner<D: Iterator<Item = Duration> + 'static>(
+        &self,
+        account: AccountToken,
+        key: talpid_types::net::wireguard::PublicKey,
+        retry_strategy: D,
+        offline_check: bool,
+    ) -> impl Future<Output = Result<()>> {
         let mut rpc = mullvad_rpc::WireguardKeyProxy::new(self.http_handle.clone());
-        async move {
-            rpc.remove_wireguard_key(account, &key)
-                .await
-                .map_err(Self::map_rpc_error)
-        }
+        let api_handle = self.availability_handle.clone();
+        let future = retry_future_n(
+            move || {
+                let remove_key = rpc.remove_wireguard_key(account.clone(), key.clone());
+                let wait_future = api_handle.wait_online();
+                async move {
+                    if offline_check {
+                        let _ = wait_future.await;
+                    }
+                    remove_key.await
+                }
+            },
+            move |result| match result {
+                Ok(_) => false,
+                Err(error) => Self::should_retry(error),
+            },
+            retry_strategy,
+            MAX_KEY_REMOVAL_RETRIES,
+        );
+        async move { future.await.map_err(Self::map_rpc_error) }
     }
 
     fn should_retry(error: &RestError) -> bool {
diff --git a/mullvad-rpc/src/lib.rs b/mullvad-rpc/src/lib.rs
index 098ed2f0b4..883128f8aa 100644
--- a/mullvad-rpc/src/lib.rs
+++ b/mullvad-rpc/src/lib.rs
@@ -492,14 +492,13 @@ impl WireguardKeyProxy {
         rest::deserialize_body(response).await
     }
 
-    pub async fn remove_wireguard_key(
+    pub fn remove_wireguard_key(
         &mut self,
         account_token: AccountToken,
-        key: &wireguard::PublicKey,
-    ) -> Result<(), rest::Error> {
+        key: wireguard::PublicKey,
+    ) -> impl Future<Output = Result<(), rest::Error>> {
         let service = self.handle.service.clone();
-
-        let _ = rest::send_request(
+        let future = rest::send_request(
             &self.handle.factory,
             service,
             &format!(
@@ -509,9 +508,11 @@ impl WireguardKeyProxy {
             Method::DELETE,
             Some(account_token),
             StatusCode::NO_CONTENT,
-        )
-        .await?;
-        Ok(())
+        );
+        async move {
+            let _ = future.await?;
+            Ok(())
+        }
     }
 }
 
diff --git a/mullvad-setup/src/main.rs b/mullvad-setup/src/main.rs
index 5f88e3456f..e809052e05 100644
--- a/mullvad-setup/src/main.rs
+++ b/mullvad-setup/src/main.rs
@@ -179,7 +179,7 @@ async fn remove_wireguard_key() -> Result<(), Error> {
             let mut key_proxy =
                 mullvad_rpc::WireguardKeyProxy::new(rpc_runtime.mullvad_rest_handle());
             key_proxy
-                .remove_wireguard_key(token, &wg_data.private_key.public_key())
+                .remove_wireguard_key(token, wg_data.private_key.public_key())
                 .await
                 .map_err(Error::RemoveKeyError)?;
             settings