Validate device in state transition handler

3 files changed, 68 insertions, 37 deletions

diff --git a/mullvad-daemon/src/device.rs b/mullvad-daemon/src/device.rs
index 4796ad31d4..07686e1fcc 100644
--- a/mullvad-daemon/src/device.rs
+++ b/mullvad-daemon/src/device.rs
@@ -80,7 +80,9 @@ impl Error {
 pub enum ValidationResult {
     /// The device and key were valid.
     Valid,
-    /// The device was valid but one or more fields, such as the key, were replaced
+    /// The device was valid but the key was replaced
+    RotatedKey,
+    /// The device was valid but one or more fields, such as ports, were replaced
     Updated,
     /// The device was not found remotely and was removed from the cache.
     Removed,
@@ -282,8 +284,11 @@ impl AccountManager {
             let mut inner = self.inner.lock().unwrap();
             inner.data.replace(data.clone());
             let (result_tx, _result_rx) = oneshot::channel();
-            let _ = self.cache_update_tx.unbounded_send((Some(data), result_tx));
+            let _ = self
+                .cache_update_tx
+                .unbounded_send((Some(data.clone()), result_tx));
             // NOTE: No need to wait on cache update
+            let _ = self.key_update_tx.send(DeviceKeyEvent(data));
         }
         self.start_key_rotation();
         result
@@ -317,6 +322,8 @@ impl AccountManager {
             inner.data.as_ref().ok_or(Error::NoDevice)?.clone()
         };
 
+        log::debug!("Checking whether the device is still valid");
+
         match self
             .device_service
             .get(data.token.clone(), data.device.id.clone())
@@ -341,7 +348,7 @@ impl AccountManager {
                 } else {
                     log::debug!("Rotating invalid WireGuard key");
                     self.rotate_key().await?;
-                    Ok(ValidationResult::Updated)
+                    Ok(ValidationResult::RotatedKey)
                 }
             }
             Err(Error::InvalidAccount) | Err(Error::InvalidDevice) => {
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index e8f475aaa5..f8102f8c3f 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -944,6 +944,7 @@ where
     ) {
         self.reset_rpc_sockets_on_tunnel_state_transition(&tunnel_state_transition)
             .await;
+
         let tunnel_state = match tunnel_state_transition {
             TunnelStateTransition::Disconnected => TunnelState::Disconnected,
             TunnelStateTransition::Connecting(endpoint) => TunnelState::Connecting {
@@ -960,7 +961,12 @@ where
             TunnelStateTransition::Error(error_state) => TunnelState::Error(error_state),
         };
 
-        self.unschedule_reconnect();
+        if !tunnel_state.is_connected() {
+            // Cancel reconnects except when entering the connected state.
+            // Exempt the latter because a reconnect scheduled while connecting should not be
+            // aborted.
+            self.unschedule_reconnect();
+        }
 
         log::debug!("New tunnel state: {:?}", tunnel_state);
         match tunnel_state {
@@ -979,14 +985,17 @@ where
                 }
 
                 if let ErrorStateCause::AuthFailed(_) = error_state.cause() {
-                    self.schedule_reconnect(Duration::from_secs(60)).await
+                    self.schedule_reconnect(Duration::from_secs(60))
                 }
             }
             _ => {}
         }
 
         self.tunnel_state = tunnel_state.clone();
-        self.event_listener.notify_new_state(tunnel_state);
+        self.event_listener.notify_new_state(tunnel_state.clone());
+
+        // Check device validity last so that the broadcast is not delayed.
+        self.maybe_validate_device(&tunnel_state).await;
     }
 
     async fn reset_rpc_sockets_on_tunnel_state_transition(
@@ -1002,6 +1011,41 @@ where
         };
     }
 
+    /// Check whether the device is valid after a number of failed connection attempts.
+    async fn maybe_validate_device(&mut self, tunnel_state: &TunnelState) {
+        match tunnel_state {
+            TunnelState::Connecting { endpoint, .. } => {
+                if endpoint.tunnel_type != TunnelType::Wireguard {
+                    return;
+                }
+                self.wg_retry_attempt += 1;
+                if self.wg_check_validity && self.wg_retry_attempt % WG_DEVICE_CHECK_THRESHOLD == 0
+                {
+                    match self.account_manager.validate_device().await {
+                        Ok(status) => {
+                            self.handle_validation_result(status);
+                            self.wg_check_validity = false;
+                        }
+                        Err(error) => {
+                            log::error!(
+                                "{}",
+                                error.display_chain_with_msg("Failed to check device validity")
+                            );
+                            if !error.is_network_error() {
+                                self.wg_check_validity = false;
+                            }
+                        }
+                    }
+                }
+            }
+            TunnelState::Connected { .. } | TunnelState::Disconnected => {
+                self.wg_check_validity = true;
+                self.wg_retry_attempt = 0;
+            }
+            _ => (),
+        }
+    }
+
     async fn handle_generate_tunnel_parameters(
         &mut self,
         tunnel_parameters_tx: &sync_mpsc::Sender<
@@ -1088,10 +1132,6 @@ where
         let tunnel_options = self.settings.tunnel_options.clone();
         let location = relay.location.as_ref().expect("Relay has no location set");
         self.last_generated_bridge_relay = None;
-        if retry_attempt == 0 {
-            self.wg_retry_attempt = 0;
-            self.wg_check_validity = true;
-        }
         match endpoint {
             MullvadEndpoint::OpenVpn(endpoint) => {
                 let proxy_settings = match &self.settings.bridge_settings {
@@ -1156,26 +1196,6 @@ where
                 .into())
             }
             MullvadEndpoint::Wireguard(endpoint) => {
-                self.wg_retry_attempt += 1;
-                if self.wg_check_validity && self.wg_retry_attempt % WG_DEVICE_CHECK_THRESHOLD == 0
-                {
-                    match self.account_manager.validate_device().await {
-                        Ok(status) => {
-                            self.handle_validation_result(status);
-                            self.wg_check_validity = false;
-                        }
-                        Err(error) => {
-                            log::error!(
-                                "{}",
-                                error.display_chain_with_msg("Failed to check device validity")
-                            );
-                            if !error.is_network_error() {
-                                self.wg_check_validity = false;
-                            }
-                        }
-                    }
-                }
-
                 let wg_data = self
                     .account_manager
                     .data()
@@ -1207,7 +1227,7 @@ where
     // Emit the appropriate events for an updated device.
     fn handle_validation_result(&mut self, result: device::ValidationResult) {
         match result {
-            device::ValidationResult::Valid => (),
+            device::ValidationResult::RotatedKey | device::ValidationResult::Valid => (),
             device::ValidationResult::Removed => {
                 self.event_listener
                     .notify_device_event(DeviceEvent::revoke(true));
@@ -1219,7 +1239,7 @@ where
         }
     }
 
-    async fn schedule_reconnect(&mut self, delay: Duration) {
+    fn schedule_reconnect(&mut self, delay: Duration) {
         self.unschedule_reconnect();
 
         let tunnel_command_tx = self.tx.to_specialized_sender();
@@ -1431,7 +1451,7 @@ where
             return;
         }
         if let Some(TunnelType::Wireguard) = self.get_target_tunnel_type() {
-            self.schedule_reconnect(WG_RECONNECT_DELAY).await;
+            self.schedule_reconnect(WG_RECONNECT_DELAY);
         }
         self.event_listener
             .notify_device_event(DeviceEvent::from_device(event.0, false));
@@ -2410,10 +2430,6 @@ where
 
     async fn on_rotate_wireguard_key(&mut self, tx: ResponseTx<(), Error>) {
         let result = self.account_manager.rotate_key().await;
-        if let Ok(ref _wg_data) = result {
-            let device = DeviceEvent::new(self.account_manager.data(), false);
-            self.event_listener.notify_device_event(device);
-        }
         let _ = tx.send(result.map(|_| ()).map_err(Error::KeyRotationError));
     }
 
diff --git a/mullvad-types/src/states.rs b/mullvad-types/src/states.rs
index 9d3b188db4..86d9e816fe 100644
--- a/mullvad-types/src/states.rs
+++ b/mullvad-types/src/states.rs
@@ -55,4 +55,12 @@ impl TunnelState {
             _ => false,
         }
     }
+
+    /// Returns true if the tunnel state is in the connected state.
+    pub fn is_connected(&self) -> bool {
+        match self {
+            TunnelState::Connected { .. } => true,
+            _ => false,
+        }
+    }
 }