Refactor key rotation

4 files changed, 50 insertions, 77 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 739a5bf51a..e1641f7852 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1279,7 +1279,6 @@ dependencies = [
  "talpid-core 0.1.0",
  "talpid-ipc 0.1.0",
  "talpid-types 0.1.0",
- "tokio 0.1.22 (registry+https://github.com/rust-lang/crates.io-index)",
  "tokio-core 0.1.17 (registry+https://github.com/rust-lang/crates.io-index)",
  "tokio-retry 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "tokio-timer 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)",
diff --git a/mullvad-daemon/Cargo.toml b/mullvad-daemon/Cargo.toml
index 37365a8c90..fec93c0a79 100644
--- a/mullvad-daemon/Cargo.toml
+++ b/mullvad-daemon/Cargo.toml
@@ -33,7 +33,6 @@ rand = "0.7"
 regex = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
-tokio = "0.1"
 tokio-core = "0.1"
 tokio-retry = "0.2"
 tokio-timer = "0.1"
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index e4cb5d750d..c05e668a2a 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -452,7 +452,6 @@ where
             internal_event_tx.clone(),
             rpc_handle.clone(),
             tokio_remote.clone(),
-            settings.get_account_token(),
         );
 
         // Attempt to download a fresh relay list
@@ -483,15 +482,18 @@ where
 
         daemon.ensure_wireguard_keys_for_current_account();
 
-        daemon.wireguard_key_manager.set_rotation_interval(
-            &mut daemon.account_history,
-            daemon
-                .settings
-                .get_tunnel_options()
-                .wireguard
-                .automatic_rotation
-                .map(|hours| 60 * hours),
-        );
+        if let Some(token) = daemon.settings.get_account_token() {
+            daemon.wireguard_key_manager.set_rotation_interval(
+                &mut daemon.account_history,
+                token,
+                daemon
+                    .settings
+                    .get_tunnel_options()
+                    .wireguard
+                    .automatic_rotation
+                    .map(|hours| 60 * hours),
+            );
+        }
 
         Ok(daemon)
     }
@@ -1125,8 +1127,11 @@ where
 
             self.ensure_wireguard_keys_for_current_account();
 
-            self.wireguard_key_manager
-                .set_account_token(&mut self.account_history, account_token);
+            if let Some(token) = account_token {
+                // update automatic rotation
+                self.wireguard_key_manager
+                    .reset_rotation(&mut self.account_history, token);
+            }
         }
         Ok(account_changed)
     }
@@ -1374,12 +1379,17 @@ where
             Ok(settings_changed) => {
                 Self::oneshot_send(tx, (), "set_wireguard_automatic_rotation response");
                 if settings_changed {
-                    self.event_listener.notify_settings(self.settings.clone());
+                    let account_token = self.settings.get_account_token();
 
-                    self.wireguard_key_manager.set_rotation_interval(
-                        &mut self.account_history,
-                        interval.map(|mins| 60 * mins),
-                    );
+                    if let Some(token) = account_token {
+                        self.wireguard_key_manager.set_rotation_interval(
+                            &mut self.account_history,
+                            token,
+                            interval.map(|hours| 60 * hours),
+                        );
+                    }
+
+                    self.event_listener.notify_settings(self.settings.clone());
                 }
             }
             Err(e) => error!("{}", e.display_chain_with_msg("Unable to save settings")),
@@ -1450,9 +1460,10 @@ where
                     let keygen_event = KeygenEvent::NewKey(public_key);
                     self.event_listener.notify_key_event(keygen_event.clone());
 
-                    // reset automatic rotation
+                    // update automatic rotation
                     self.wireguard_key_manager.set_rotation_interval(
                         &mut self.account_history,
+                        account_token.clone(),
                         self.settings
                             .get_tunnel_options()
                             .wireguard
diff --git a/mullvad-daemon/src/wireguard.rs b/mullvad-daemon/src/wireguard.rs
index 3da3b50e43..1e1b8f72a5 100644
--- a/mullvad-daemon/src/wireguard.rs
+++ b/mullvad-daemon/src/wireguard.rs
@@ -53,8 +53,6 @@ pub struct KeyManager {
     current_job: Option<CancelHandle>,
 
     abort_scheduler_tx: Option<CancelHandle>,
-    account_token: Option<AccountToken>,
-    public_key: Option<PublicKey>,
     // unit: minutes
     auto_rotation_interval: u32,
 }
@@ -64,7 +62,6 @@ impl KeyManager {
         daemon_tx: mpsc::Sender<InternalDaemonEvent>,
         http_handle: mullvad_rpc::HttpHandle,
         tokio_remote: Remote,
-        account_token: Option<AccountToken>,
     ) -> Self {
         Self {
             daemon_tx,
@@ -72,32 +69,30 @@ impl KeyManager {
             tokio_remote,
             current_job: None,
             abort_scheduler_tx: None,
-
-            account_token,
-            public_key: None,
             auto_rotation_interval: 0,
         }
     }
 
-    fn update_public_key(&mut self, account_history: &mut AccountHistory) {
-        log::debug!("update_public_key");
-        let _ = self.public_key.take();
-
-        let token = if let Some(token) = &self.account_token {
-            token
-        } else {
-            log::warn!("Cannot update public key; no account token is set");
-            return ();
-        };
+    /// Update automatic key rotation interval (given in minutes)
+    /// Passing `None` for the interval will use the default value.
+    /// A value of `0` disables automatic key rotation.
+    pub fn reset_rotation(
+        &mut self,
+        account_history: &mut AccountHistory,
+        account_token: AccountToken,
+    ) {
+        log::debug!("reset_rotation");
 
-        self.public_key = match account_history.get(&token) {
-            Ok(v) => v
-                .map(|entry| entry.wireguard.map(|wg| wg.get_public_key()))
-                .unwrap(),
-            Err(e) => {
-                log::error!("KeyManager failed to obtain public key. {}", e);
-                None
+        match account_history
+            .get(&account_token)
+            .map(|entry| entry.map(|entry| entry.wireguard.map(|wg| wg.get_public_key())))
+        {
+            Ok(Some(Some(public_key))) => self.run_automatic_rotation(account_token, public_key),
+            Ok(Some(None)) => {
+                log::error!("reset_rotation: failed to obtain public key for account entry.")
             }
+            Ok(None) => log::error!("reset_rotation: account entry not found."),
+            Err(e) => log::error!("reset_rotation: failed to obtain account entry. {}", e),
         };
     }
 
@@ -107,6 +102,7 @@ impl KeyManager {
     pub fn set_rotation_interval(
         &mut self,
         account_history: &mut AccountHistory,
+        account_token: AccountToken,
         auto_rotation_interval_mins: Option<u32>,
     ) {
         log::debug!("set_rotation_interval");
@@ -114,9 +110,7 @@ impl KeyManager {
         self.auto_rotation_interval =
             auto_rotation_interval_mins.unwrap_or(DEFAULT_AUTOMATIC_KEY_ROTATION);
 
-        self.stop_automatic_rotation();
-        self.update_public_key(account_history);
-        self.run_automatic_rotation();
+        self.reset_rotation(account_history, account_token);
     }
 
     /// Stop current key generation
@@ -288,18 +282,6 @@ impl KeyManager {
             _ => Error::RpcError(err),
         }
     }
-
-    pub fn set_account_token(
-        &mut self,
-        account_history: &mut AccountHistory,
-        account_token: Option<AccountToken>,
-    ) {
-        log::debug!("set_account_token");
-        self.account_token = account_token;
-
-        self.set_rotation_interval(account_history, Some(self.auto_rotation_interval));
-    }
-
     fn create_key_expiration_timer(
         public_key: PublicKey,
         rotation_interval_secs: u64,
@@ -436,7 +418,7 @@ impl KeyManager {
         Box::new(fut.then(create_repeat_future).map(|_| ()))
     }
 
-    fn run_automatic_rotation(&mut self) {
+    fn run_automatic_rotation(&mut self, account_token: AccountToken, public_key: PublicKey) {
         self.stop_automatic_rotation();
 
         if let 0 = self.auto_rotation_interval {
@@ -444,24 +426,6 @@ impl KeyManager {
             return;
         }
 
-        if let None = self.account_token {
-            log::warn!(
-                "Not running automatic rotation since no \
-                 account token is set"
-            );
-            return;
-        }
-        let account_token = self.account_token.as_ref().unwrap().to_string();
-
-        if let None = self.public_key {
-            log::warn!(
-                "Not running automatic rotation since no \
-                 public key is set"
-            );
-            return;
-        }
-        let public_key = self.public_key.as_ref().unwrap().clone();
-
         // Schedule cancellable series of repeating rotation tasks
         let fut = Self::create_automatic_rotation(
             self.daemon_tx.clone(),