Start key rotation after first key created

3 files changed, 31 insertions, 40 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7a72f205df..f7379d65c8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -40,6 +40,7 @@ Line wrap the file at 100 chars.                                              Th
 
 ### Fixed
 - Stop resetting the firewall after an upgrade to not leak after an upgrade.
+- Start key rotation when WireGuard key is first created.
 
 #### Windows
 - Remove firewall filters (unblock internet access) when "Always require VPN" is enabled and the app
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 737ffab185..31f6e656f8 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -623,22 +623,6 @@ where
 
         daemon.ensure_wireguard_keys_for_current_account().await;
 
-        if let Some(token) = daemon.settings.get_account_token() {
-            daemon
-                .wireguard_key_manager
-                .set_rotation_interval(
-                    &mut daemon.account_history,
-                    token,
-                    daemon
-                        .settings
-                        .tunnel_options
-                        .wireguard
-                        .automatic_rotation
-                        .map(|hours| Duration::from_secs(60u64 * 60u64 * hours as u64)),
-                )
-                .await;
-        }
-
         Ok(daemon)
     }
 
@@ -1084,6 +1068,8 @@ where
                         account: account.clone(),
                         wireguard: None,
                     });
+                // if no key existed before
+                let first_key_for_account_on_host = account_entry.wireguard.is_none();
                 account_entry.wireguard = Some(data);
                 match self.account_history.insert(account_entry).await {
                     Ok(_) => {
@@ -1091,7 +1077,10 @@ where
                             self.schedule_reconnect(WG_RECONNECT_DELAY).await;
                         }
                         self.event_listener
-                            .notify_key_event(KeygenEvent::NewKey(public_key))
+                            .notify_key_event(KeygenEvent::NewKey(public_key));
+                        if first_key_for_account_on_host {
+                            self.ensure_key_rotation().await;
+                        }
                     }
                     Err(e) => {
                         log::error!(
@@ -1120,6 +1109,21 @@ where
         }
     }
 
+    async fn ensure_key_rotation(&mut self) {
+        if let Some(token) = self.settings.get_account_token() {
+            let rotation_interval = self
+                .settings
+                .tunnel_options
+                .wireguard
+                .automatic_rotation
+                .map(|hours| Duration::from_secs(60u64 * 60u64 * hours as u64));
+
+            self.wireguard_key_manager
+                .set_rotation_interval(&mut self.account_history, token, rotation_interval)
+                .await;
+        }
+    }
+
     async fn handle_new_account_event(
         &mut self,
         new_token: AccountToken,
@@ -1342,13 +1346,6 @@ where
             }
 
             self.ensure_wireguard_keys_for_current_account().await;
-
-            if let Some(token) = account_token {
-                // update automatic rotation
-                self.wireguard_key_manager
-                    .reset_rotation(&mut self.account_history, token)
-                    .await;
-            }
         }
         Ok(account_changed)
     }
@@ -1679,19 +1676,7 @@ where
             Ok(settings_changed) => {
                 Self::oneshot_send(tx, (), "set_wireguard_rotation_interval response");
                 if settings_changed {
-                    let account_token = self.settings.get_account_token();
-
-                    if let Some(token) = account_token {
-                        self.wireguard_key_manager
-                            .set_rotation_interval(
-                                &mut self.account_history,
-                                token,
-                                interval
-                                    .map(|hours| Duration::from_secs(60u64 * 60u64 * hours as u64)),
-                            )
-                            .await;
-                    }
-
+                    self.ensure_key_rotation().await;
                     self.event_listener
                         .notify_settings(self.settings.to_settings());
                 }
@@ -1711,10 +1696,11 @@ where
             {
                 log::info!("Automatically generating new wireguard key for account");
                 self.wireguard_key_manager
-                    .generate_key_async(account, Some(FIRST_KEY_PUSH_TIMEOUT))
+                    .spawn_key_generation_task(account, Some(FIRST_KEY_PUSH_TIMEOUT))
                     .await;
             } else {
-                log::info!("Account already has wireguard key");
+                log::info!("Account already has wireguard key, starting key rotation.");
+                self.ensure_key_rotation().await;
             }
         }
     }
diff --git a/mullvad-daemon/src/wireguard.rs b/mullvad-daemon/src/wireguard.rs
index 0e13496b0a..28595fe1da 100644
--- a/mullvad-daemon/src/wireguard.rs
+++ b/mullvad-daemon/src/wireguard.rs
@@ -150,7 +150,11 @@ impl KeyManager {
 
 
     /// Generate a new private key asynchronously. The new keys will be sent to the daemon channel.
-    pub async fn generate_key_async(&mut self, account: AccountToken, timeout: Option<Duration>) {
+    pub async fn spawn_key_generation_task(
+        &mut self,
+        account: AccountToken,
+        timeout: Option<Duration>,
+    ) {
         self.reset();
         let private_key = PrivateKey::new_from_random();