Reset rotation interval when the active account changes

2 files changed, 32 insertions, 16 deletions

diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 7827bd1edc..4410c05679 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -1088,13 +1088,26 @@ where
             Ok(account_changed) => {
                 if account_changed {
                     match account_token {
-                        Some(_) => {
+                        Some(account_token) => {
                             info!("Initiating tunnel restart because the account token changed");
+
+                            let public_key = self.account_history
+                                .get(&account_token)
+                                .unwrap_or(None)
+                                .and_then(|entry| entry.wireguard.map(|wg| wg.get_public_key()));
+                            self.wireguard_key_manager.update_rotation_interval(Some(
+                                wireguard::KeyRotationParameters {
+                                    public_key,
+                                    interval: self.settings.get_tunnel_options().wireguard.automatic_rotation,
+                                },
+                            ));
+
                             self.reconnect_tunnel();
                         }
                         None => {
                             info!("Disconnecting because account token was cleared");
                             self.set_target_state(TargetState::Unsecured);
+                            self.wireguard_key_manager.update_rotation_interval(None);
                         }
                     };
                 }
@@ -1378,12 +1391,12 @@ where
                             .and_then(|entry| entry.wireguard.map(|wg| wg.get_public_key()))
                     });
 
-                    self.wireguard_key_manager.update_rotation_interval(
+                    self.wireguard_key_manager.update_rotation_interval(Some(
                         wireguard::KeyRotationParameters {
                             public_key,
                             interval,
                         },
-                    );
+                    ));
                 }
             }
             Err(e) => error!("{}", e.display_chain_with_msg("Unable to save settings")),
diff --git a/mullvad-daemon/src/wireguard.rs b/mullvad-daemon/src/wireguard.rs
index 7654bed061..476cce2b32 100644
--- a/mullvad-daemon/src/wireguard.rs
+++ b/mullvad-daemon/src/wireguard.rs
@@ -197,7 +197,7 @@ impl KeyManager {
             current_job: None,
             abort_scheduler_tx: None,
         };
-        manager.update_rotation_interval(automatic_key_rotation);
+        manager.update_rotation_interval(Some(automatic_key_rotation));
 
         manager
     }
@@ -205,24 +205,27 @@ impl KeyManager {
     /// Update automatic key rotation interval (given in hours)
     /// Passing `None` for the interval will use the default value.
     /// A value of `0` disables automatic key rotation.
-    pub fn update_rotation_interval(&mut self, automatic_key_rotation: KeyRotationParameters) {
+    pub fn update_rotation_interval(&mut self, automatic_key_rotation: Option<KeyRotationParameters>) {
         log::debug!("update_rotation_interval");
         if self.abort_scheduler_tx.is_some() {
             // Stop existing scheduler, if one exists
             let tx = self.abort_scheduler_tx.take().unwrap();
             let _ = tx.send(());
         }
-        self.abort_scheduler_tx = match automatic_key_rotation.interval {
-            // Interval=0 disables automatic key rotation
-            Some(0) => None,
-            _ => KeyRotationScheduler::new(
-                self.tokio_remote.clone(),
-                self.daemon_tx.clone(),
-                automatic_key_rotation.public_key,
-                automatic_key_rotation.interval,
-            )
-            .ok(),
-        };
+
+        if let Some(automatic_key_rotation) = automatic_key_rotation {
+            self.abort_scheduler_tx = match automatic_key_rotation.interval {
+                // Interval=0 disables automatic key rotation
+                Some(0) => None,
+                _ => KeyRotationScheduler::new(
+                    self.tokio_remote.clone(),
+                    self.daemon_tx.clone(),
+                    automatic_key_rotation.public_key,
+                    automatic_key_rotation.interval,
+                )
+                .ok(),
+            };
+        }
     }
 
     /// Stop current key generation