Merge branch 'test_automatic_wireguard_rotation-fails-due-to-timeout-des-1094'

8 files changed, 132 insertions, 128 deletions

diff --git a/mullvad-daemon/src/tunnel.rs b/mullvad-daemon/src/tunnel.rs
index 31f4686df7..839f07e299 100644
--- a/mullvad-daemon/src/tunnel.rs
+++ b/mullvad-daemon/src/tunnel.rs
@@ -25,7 +25,7 @@ use talpid_types::net::{obfuscation::ObfuscatorConfig, wireguard, TunnelParamete
 
 use talpid_types::{tunnel::ParameterGenerationError, ErrorExt};
 
-use crate::device::{AccountManagerHandle, PrivateAccountAndDevice};
+use crate::device::{AccountManagerHandle, Error as DeviceError, PrivateAccountAndDevice};
 
 /// The IP-addresses that the client uses when it connects to a server that supports the
 /// "Same IP" functionality. This means all clients have the same in-tunnel IP on these
@@ -49,6 +49,9 @@ pub enum Error {
 
     #[error("Failed to resolve hostname for custom relay")]
     ResolveCustomHostname,
+
+    #[error("Failed to get device data")]
+    Device(#[from] DeviceError),
 }
 
 #[derive(Clone)]
@@ -257,13 +260,8 @@ impl InnerParametersGenerator {
     }
 
     async fn device(&self) -> Result<PrivateAccountAndDevice, Error> {
-        self.account_manager
-            .data()
-            .await
-            .map(|s| s.into_device())
-            .ok()
-            .flatten()
-            .ok_or(Error::NoAuthDetails)
+        let device_state = self.account_manager.data().await?;
+        device_state.into_device().ok_or(Error::NoAuthDetails)
     }
 }
 
@@ -279,25 +277,33 @@ impl TunnelParametersGenerator for ParametersGenerator {
             inner
                 .generate(retry_attempt, ipv6)
                 .await
-                .map_err(|error| match error {
-                    Error::SelectRelay(mullvad_relay_selector::Error::NoBridge) => {
-                        ParameterGenerationError::NoMatchingBridgeRelay
-                    }
-                    Error::ResolveCustomHostname => {
-                        ParameterGenerationError::CustomTunnelHostResultionError
-                    }
-                    error => {
-                        log::error!(
-                            "{}",
-                            error.display_chain_with_msg("Failed to generate tunnel parameters")
-                        );
-                        ParameterGenerationError::NoMatchingRelay
-                    }
+                .inspect_err(|error| {
+                    log::error!(
+                        "{}",
+                        error.display_chain_with_msg("Failed to generate tunnel parameters")
+                    );
                 })
+                .map_err(ParameterGenerationError::from)
         })
     }
 }
 
+impl From<Error> for ParameterGenerationError {
+    fn from(error: Error) -> Self {
+        match error {
+            Error::SelectRelay(mullvad_relay_selector::Error::NoBridge) => {
+                ParameterGenerationError::NoMatchingBridgeRelay
+            }
+            Error::ResolveCustomHostname => {
+                ParameterGenerationError::CustomTunnelHostResultionError
+            }
+            Error::NoAuthDetails | Error::SelectRelay(_) | Error::Device(_) => {
+                ParameterGenerationError::NoMatchingRelay
+            }
+        }
+    }
+}
+
 /// Contains all relays that were selected last time when tunnel parameters were generated.
 enum LastSelectedRelays {
     /// Represents all relays generated for a WireGuard tunnel.
diff --git a/mullvad-management-interface/src/types/conversions/device.rs b/mullvad-management-interface/src/types/conversions/device.rs
index b43c47ee99..fd9f05ccbc 100644
--- a/mullvad-management-interface/src/types/conversions/device.rs
+++ b/mullvad-management-interface/src/types/conversions/device.rs
@@ -78,9 +78,9 @@ impl From<mullvad_types::device::DeviceState> for proto::DeviceState {
     fn from(state: mullvad_types::device::DeviceState) -> Self {
         proto::DeviceState {
             state: proto::device_state::State::from(&state) as i32,
-            device: state.into_device().map(|device| proto::AccountAndDevice {
-                account_token: device.account_token,
-                device: Some(proto::Device::from(device.device)),
+            device: state.logged_in().map(|client| proto::AccountAndDevice {
+                account_token: client.account_token,
+                device: Some(proto::Device::from(client.device)),
             }),
         }
     }
diff --git a/mullvad-types/src/device.rs b/mullvad-types/src/device.rs
index 140815414d..5ee51d54ca 100644
--- a/mullvad-types/src/device.rs
+++ b/mullvad-types/src/device.rs
@@ -50,23 +50,18 @@ pub enum DeviceState {
 }
 
 impl DeviceState {
-    pub fn into_device(self) -> Option<AccountAndDevice> {
+    /// Returns the active account and device if the device is currently logged in to a valid
+    /// account.
+    pub fn logged_in(self) -> Option<AccountAndDevice> {
         match self {
-            DeviceState::LoggedIn(dev) => Some(dev),
+            DeviceState::LoggedIn(client) => Some(client),
             _ => None,
         }
     }
 
-    pub fn is_logged_in(&self) -> bool {
+    pub const fn is_logged_in(&self) -> bool {
         matches!(self, Self::LoggedIn(_))
     }
-
-    pub fn get_account(&self) -> Option<&AccountAndDevice> {
-        match self {
-            DeviceState::LoggedIn(ref account) => Some(account),
-            _ => None,
-        }
-    }
 }
 
 /// A [Device] and its associated account token.
diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs
index 4449a06a6f..04f72d575a 100644
--- a/test/test-manager/src/run_tests.rs
+++ b/test/test-manager/src/run_tests.rs
@@ -92,6 +92,7 @@ pub async fn run(
             logger.store_records(true);
         }
 
+        // TODO: Log how long each test took to run.
         let test_result = run_test(
             client.clone(),
             mullvad_client,
diff --git a/test/test-manager/src/tests/account.rs b/test/test-manager/src/tests/account.rs
index c67e1bed19..7b1f83cfac 100644
--- a/test/test-manager/src/tests/account.rs
+++ b/test/test-manager/src/tests/account.rs
@@ -10,6 +10,7 @@ use mullvad_types::{
 };
 use std::time::Duration;
 use talpid_types::net::wireguard;
+use talpid_types::net::wireguard::PublicKey;
 use test_macro::test_function;
 use test_rpc::ServiceClient;
 
@@ -145,8 +146,8 @@ pub async fn test_revoked_device(
         .get_device()
         .await
         .context("failed to get device data")?
-        .into_device()
-        .unwrap()
+        .logged_in()
+        .context("Client is not logged in to a valid account")?
         .device
         .id;
 
@@ -207,6 +208,85 @@ pub async fn test_revoked_device(
     Ok(())
 }
 
+/// Assert that an old Wireguard key is automatically rotated by the daemon.
+#[test_function]
+pub async fn test_automatic_wireguard_rotation(
+    ctx: TestContext,
+    rpc: ServiceClient,
+    mut mullvad_client: MullvadProxyClient,
+) -> anyhow::Result<()> {
+    const ROTATION_TIMEOUT: Duration = Duration::from_secs(120);
+
+    // Make note of current WG key
+    let old_key = get_current_wireguard_key(&mut mullvad_client).await?;
+
+    log::info!("Old wireguard key: {old_key}");
+
+    rpc.stop_mullvad_daemon().await?;
+
+    log::info!("Changing created field of `device.json` to more than 7 days ago");
+    rpc.make_device_json_old()
+        .await
+        .context("Could not change device.json to have an old created timestamp")?;
+
+    rpc.start_mullvad_daemon().await?;
+
+    // NOTE: Need to create a new `mullvad_client` here after the restart otherwise we can't
+    // communicate with the daemon
+    log::info!("Reconnecting to daemon");
+    drop(mullvad_client);
+    let mut mullvad_client = ctx.rpc_provider.new_client().await;
+
+    // Check if the key rotation has already occurred when connected to the daemon, otherwise
+    // listen for device daemon events until we observe the change. We have to register the event
+    // listener before polling the current key to be sure we don't miss the change.
+    log::info!("Verifying that wireguard key has changed");
+    let event_listener = mullvad_client
+        .events_listen()
+        .await
+        .context("Failed to begin listening for state changes")?;
+    let new_key = get_current_wireguard_key(&mut mullvad_client).await?;
+
+    // If key has not yet been updated, listen for changes to it
+    if new_key == old_key {
+        // Verify rotation has happened within `ROTATION_TIMEOUT` - if the key hasn't been rotated after
+        // that, the rotation probably won't happen anytime soon.
+        log::info!("Listening for device daemon event");
+        let device_event = |daemon_event| match daemon_event {
+            DaemonEvent::Device(device_event) => Some(device_event),
+            _ => None,
+        };
+        let device_event_listener = tokio::time::timeout(
+            ROTATION_TIMEOUT,
+            helpers::find_daemon_event(event_listener, device_event),
+        );
+        let _ = device_event_listener.await;
+
+        // Note: The key rotation could possible have happened without us noticing due to
+        // some raceiness in the timeframe between starting the daemon and us starting to
+        // listen for new daemon events. Thus, it is probably a good idea to check manually if the
+        // device key was rotated.
+        let new_key = get_current_wireguard_key(&mut mullvad_client).await?;
+
+        assert_ne!(old_key, new_key);
+    }
+
+    Ok(())
+}
+
+async fn get_current_wireguard_key(
+    mullvad_client: &mut MullvadProxyClient,
+) -> anyhow::Result<PublicKey> {
+    let pubkey = mullvad_client
+        .get_device()
+        .await?
+        .logged_in()
+        .context("Client is not logged in to a valid account")?
+        .device
+        .pubkey;
+    Ok(pubkey)
+}
+
 /// Remove all devices on the current account
 pub async fn clear_devices(device_client: &DevicesProxy) -> anyhow::Result<()> {
     log::info!("Removing all devices for account");
@@ -277,86 +357,3 @@ pub async fn retry_if_throttled<
         }
     }
 }
-
-#[test_function]
-pub async fn test_automatic_wireguard_rotation(
-    ctx: TestContext,
-    rpc: ServiceClient,
-    mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
-    // Make note of current WG key
-    let old_key = mullvad_client
-        .get_device()
-        .await
-        .unwrap()
-        .into_device()
-        .expect("Could not get device")
-        .device
-        .pubkey;
-
-    log::info!("Old wireguard key: {old_key}");
-
-    log::info!("Stopping daemon");
-    rpc.stop_mullvad_daemon()
-        .await
-        .expect("Could not stop system service");
-
-    log::info!("Changing created field of `device.json` to more than 7 days ago");
-    rpc.make_device_json_old()
-        .await
-        .expect("Could not change device.json to have an old created timestamp");
-
-    log::info!("Starting daemon");
-    rpc.start_mullvad_daemon()
-        .await
-        .expect("Could not start system service");
-
-    // NOTE: Need to create a new `mullvad_client` here after the restart otherwise we can't
-    // communicate with the daemon
-    log::info!("Reconnecting to daemon");
-    drop(mullvad_client);
-    let mut mullvad_client = ctx.rpc_provider.new_client().await;
-
-    log::info!("Verifying that wireguard key has change");
-
-    // Check if the key rotation has already occurred when connected to the daemon, otherwise
-    // listen for device daemon events until we observe the change. We have to register the event
-    // listener before polling the current key to be sure we don't miss the change.
-    let event_listener = mullvad_client.events_listen().await.unwrap();
-    let new_key = mullvad_client
-        .get_device()
-        .await
-        .unwrap()
-        .into_device()
-        .expect("Could not get device")
-        .device
-        .pubkey;
-
-    // If key has not yet been updated, listen for changes to it
-    if new_key == old_key {
-        log::info!("Listening for device daemon event");
-        // Verify rotation has happened within 100 seconds - if the key hasn't been rotated after
-        // that, the rotation probably won't happen anytime soon.
-        let device_event = tokio::task::spawn(tokio::time::timeout(
-            Duration::from_secs(100),
-            helpers::find_daemon_event(event_listener, |daemon_event| match daemon_event {
-                DaemonEvent::Device(device_event) => Some(device_event),
-                _ => None,
-            }),
-        ))
-        .await
-        .unwrap()
-        .map_err(|_error| Error::Daemon(String::from("Tunnel event listener timed out")))??;
-
-        let new_key = device_event
-            .new_state
-            .into_device()
-            .expect("Could not get device")
-            .device
-            .pubkey;
-
-        assert_ne!(old_key, new_key);
-    }
-
-    Ok(())
-}
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index b92f68b413..b1688d27ac 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -209,8 +209,8 @@ pub async fn test_uninstall_app(
         .get_device()
         .await
         .context("failed to get device data")?
-        .into_device()
-        .context("failed to get device")?
+        .logged_in()
+        .context("Client is not logged in to a valid account")?
         .device
         .id;
 
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index 312e8ae2b3..f68d25ffa3 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -80,7 +80,7 @@ pub async fn cleanup_after_test(
     rpc_provider: &RpcClientProvider,
 ) -> anyhow::Result<()> {
     log::debug!("Resetting daemon settings after test");
-    // Check if daemon should be restarted
+    // Check if daemon should be restarted.
     restart_daemon(rpc).await?;
     let mut mullvad_client = rpc_provider.new_client().await;
     mullvad_client
diff --git a/test/test-rpc/src/client.rs b/test/test-rpc/src/client.rs
index 0c621e8936..da64dc4658 100644
--- a/test/test-rpc/src/client.rs
+++ b/test/test-rpc/src/client.rs
@@ -10,7 +10,11 @@ use super::*;
 
 const INSTALL_TIMEOUT: Duration = Duration::from_secs(300);
 const REBOOT_TIMEOUT: Duration = Duration::from_secs(30);
+/// How long to wait before proceeding after a reboot and a connection to the test-runner has been
+/// re-established
+const POST_REBOOT_GRACE_PERIOD: Duration = Duration::from_secs(5);
 const LOG_LEVEL_TIMEOUT: Duration = Duration::from_secs(60);
+const DAEMON_RESTART_TIMEOUT: Duration = Duration::from_secs(30);
 
 #[derive(Debug, Clone)]
 pub struct ServiceClient {
@@ -263,10 +267,11 @@ impl ServiceClient {
     /// This function will return *after* the app has been stopped, thus
     /// blocking execution until then.
     pub async fn stop_mullvad_daemon(&self) -> Result<(), Error> {
-        let _ = self
-            .client
-            .stop_mullvad_daemon(tarpc::context::current())
-            .await?;
+        let mut ctx = tarpc::context::current();
+        ctx.deadline = SystemTime::now()
+            .checked_add(DAEMON_RESTART_TIMEOUT)
+            .unwrap();
+        let _ = self.client.stop_mullvad_daemon(ctx).await?;
         Ok(())
     }
 
@@ -367,7 +372,7 @@ impl ServiceClient {
         self.connection_handle.reset_connected_state().await;
         self.connection_handle.wait_for_server().await?;
 
-        tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+        tokio::time::sleep(POST_REBOOT_GRACE_PERIOD).await;
 
         Ok(())
     }