Issue a delayed reconnect when after key rotation

To issue a delayed reconnect, the current reconnection mechanism was
reworked to use futures and to issue a `Reconnect` command instead of a
plain `SetTargetState`.

2 files changed, 36 insertions, 23 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 02dcc2053a..f6ddf617a0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,8 @@ Line wrap the file at 100 chars.                                              Th
 ### Added
 - Show system notification when account has expired.
 - Add fish shell completions for the mullvad CLI.
+- Reconnect with a new key when WireGuard key is rotated automatically, previously the tunnel would
+  time out before reconnecting.
 
 ### Changed
 - Upgrade from Electron 7 to Electron 8.
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index f1cb483f9a..f7755fbf75 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -28,7 +28,10 @@ use futures01::{
     Future, Stream,
 };
 use log::{debug, error, info, warn};
-use mullvad_rpc::AccountsProxy;
+use mullvad_rpc::{
+    rest::{CancelHandle, Cancellable},
+    AccountsProxy,
+};
 use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
     endpoint::MullvadEndpoint,
@@ -53,7 +56,6 @@ use std::{
     mem,
     path::PathBuf,
     sync::{mpsc, Arc, Weak},
-    thread,
     time::Duration,
 };
 #[cfg(target_os = "linux")]
@@ -86,6 +88,9 @@ const TUNNEL_STATE_MACHINE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
 /// Timeout for first WireGuard key pushing
 const FIRST_KEY_PUSH_TIMEOUT: Duration = Duration::from_secs(5);
 
+/// Delay between generating a new WireGuard key and reconnecting
+const WG_RECONNECT_DELAY: Duration = Duration::from_secs(30);
+
 #[derive(err_derive::Error, Debug)]
 #[error(no_from)]
 pub enum Error {
@@ -459,7 +464,7 @@ pub struct Daemon<L: EventListener> {
     exclude_pids: split_tunnel::PidManager,
     rx: Wait<UnboundedReceiver<InternalDaemonEvent>>,
     tx: DaemonEventSender,
-    reconnection_loop_tx: Option<mpsc::Sender<()>>,
+    reconnection_job: Option<CancelHandle>,
     event_listener: L,
     settings: SettingsPersister,
     account_history: account_history::AccountHistory,
@@ -599,7 +604,7 @@ where
             exclude_pids: split_tunnel::PidManager::new().map_err(Error::InitSplitTunneling)?,
             rx: internal_event_rx.wait(),
             tx: internal_event_tx,
-            reconnection_loop_tx: None,
+            reconnection_job: None,
             event_listener,
             settings,
             account_history,
@@ -945,29 +950,31 @@ where
 
     fn schedule_reconnect(&mut self, delay: Duration) {
         let tunnel_command_tx = self.tx.to_specialized_sender();
-        let (tx, rx) = mpsc::channel();
-
-        self.reconnection_loop_tx = Some(tx);
-
-        thread::spawn(move || {
-            let (result_tx, _result_rx) = oneshot::channel();
+        let (future, cancel_handle) = Cancellable::new(Box::pin(async move {
+            tokio02::time::delay_for(delay).await;
+            log::debug!("Attempting to reconnect");
+            let _ = tunnel_command_tx.send(DaemonCommand::Reconnect);
+        }));
 
-            if let Err(mpsc::RecvTimeoutError::Timeout) = rx.recv_timeout(delay) {
-                debug!("Attempting to reconnect");
-                let _ = tunnel_command_tx.send(DaemonCommand::SetTargetState(
-                    result_tx,
-                    TargetState::Secured,
-                ));
-            }
-        });
+        self.spawn_future(future);
+        self.reconnection_job = Some(cancel_handle);
     }
 
     fn unschedule_reconnect(&mut self) {
-        if let Some(tx) = self.reconnection_loop_tx.take() {
-            let _ = tx.send(());
+        if let Some(job) = self.reconnection_job.take() {
+            job.cancel();
         }
     }
 
+    fn spawn_future<F>(&mut self, fut: F)
+    where
+        F: std::future::Future + Send + 'static,
+        F::Output: Send,
+    {
+        self.rpc_runtime.runtime().spawn(fut);
+    }
+
+
     fn handle_command(&mut self, command: DaemonCommand) {
         use self::DaemonCommand::*;
         if !self.state.is_running() {
@@ -1063,9 +1070,13 @@ where
                     });
                 account_entry.wireguard = Some(data);
                 match self.account_history.insert(account_entry) {
-                    Ok(_) => self
-                        .event_listener
-                        .notify_key_event(KeygenEvent::NewKey(public_key)),
+                    Ok(_) => {
+                        if let Some(TunnelType::Wireguard) = self.get_connected_tunnel_type() {
+                            self.schedule_reconnect(WG_RECONNECT_DELAY);
+                        }
+                        self.event_listener
+                            .notify_key_event(KeygenEvent::NewKey(public_key))
+                    }
                     Err(e) => {
                         log::error!(
                             "{}",