Reset daemon environment when needed

6 files changed, 256 insertions, 35 deletions

diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs
index 005d0c89d4..73c72ebc2e 100644
--- a/test/test-manager/src/run_tests.rs
+++ b/test/test-manager/src/run_tests.rs
@@ -105,8 +105,8 @@ pub async fn run(
             // Try to reset the daemon state if the test failed OR if the test doesn't explicitly
             // disabled cleanup.
             if test.cleanup || matches!(test_result.result, Err(_) | Ok(Err(_))) {
-                let mut client = test_context.rpc_provider.new_client().await;
-                crate::tests::cleanup_after_test(&mut client).await?;
+                crate::tests::cleanup_after_test(client.clone(), &test_context.rpc_provider)
+                    .await?;
             }
         }
 
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index b89b4f0fd0..ce38308b0d 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -13,14 +13,16 @@ mod tunnel;
 mod tunnel_state;
 mod ui;
 
-use crate::mullvad_daemon::{MullvadClientArgument, RpcClientProvider};
+use crate::{
+    mullvad_daemon::{MullvadClientArgument, RpcClientProvider},
+    tests::helpers::get_app_env,
+};
 use anyhow::Context;
 pub use test_metadata::TestMetadata;
 use test_rpc::ServiceClient;
 
 use futures::future::BoxFuture;
 
-use mullvad_management_interface::MullvadProxyClient;
 use std::time::Duration;
 
 const WAIT_FOR_TUNNEL_STATE_TIMEOUT: Duration = Duration::from_secs(40);
@@ -69,10 +71,16 @@ pub fn get_tests() -> Vec<&'static TestMetadata> {
 }
 
 /// Restore settings to the defaults.
-pub async fn cleanup_after_test(mullvad_client: &mut MullvadProxyClient) -> anyhow::Result<()> {
+pub async fn cleanup_after_test(
+    rpc: ServiceClient,
+    rpc_provider: &RpcClientProvider,
+) -> anyhow::Result<()> {
     log::debug!("Cleaning up daemon in test cleanup");
+    // Check if daemon should be restarted
+    restart_daemon(rpc).await?;
+    let mut mullvad_client = rpc_provider.new_client().await;
 
-    helpers::disconnect_and_wait(mullvad_client).await?;
+    helpers::disconnect_and_wait(&mut mullvad_client).await?;
 
     // Bring all the settings into scope so we remember to reset them.
     let mullvad_types::settings::Settings {
@@ -176,3 +184,18 @@ pub async fn cleanup_after_test(mullvad_client: &mut MullvadProxyClient) -> anyh
 
     Ok(())
 }
+
+/// Conditionally restart the running daemon
+///
+/// If the daemon was started with non-standard environment variables, subsequent tests may break
+/// due to assuming a default configuration. In that case, reset the environment variables and
+/// restart.
+async fn restart_daemon(rpc: ServiceClient) -> anyhow::Result<()> {
+    let current_env = rpc.get_daemon_environment().await?;
+    let default_env = get_app_env();
+    if current_env != default_env {
+        log::debug!("Restarting daemon due changed environment variables. Values since last test {current_env:?}");
+        rpc.set_daemon_environment(default_env).await?;
+    }
+    Ok(())
+}
diff --git a/test/test-rpc/src/client.rs b/test/test-rpc/src/client.rs
index 79c6a038ca..84d923a826 100644
--- a/test/test-rpc/src/client.rs
+++ b/test/test-rpc/src/client.rs
@@ -314,6 +314,20 @@ impl ServiceClient {
         Ok(())
     }
 
+    /// Get the current daemon's environment variables.
+    ///
+    /// # Returns
+    /// - `Result::Ok(env)` if the current environment variables could be read.
+    /// - `Result::Err(Error)` if communication with the daemon failed or the environment values
+    /// could not be parsed.
+    pub async fn get_daemon_environment(&self) -> Result<HashMap<String, String>, Error> {
+        let env = self
+            .client
+            .get_daemon_environment(tarpc::context::current())
+            .await??;
+        Ok(env)
+    }
+
     pub async fn copy_file(&self, src: String, dest: String) -> Result<(), Error> {
         log::debug!("Copying \"{src}\" to \"{dest}\"");
         self.client
diff --git a/test/test-rpc/src/lib.rs b/test/test-rpc/src/lib.rs
index fc9ea67c84..82fed91541 100644
--- a/test/test-rpc/src/lib.rs
+++ b/test/test-rpc/src/lib.rs
@@ -59,10 +59,19 @@ pub enum Error {
     TcpForward,
     #[error("Unknown process ID: {0}")]
     UnknownPid(u32),
+    #[error("Failed to join tokio task: {0}")]
+    TokioJoinError(String),
     #[error("{0}")]
     Other(String),
 }
 
+impl Error {
+    /// Convenient mapping from a Tokio error to the test_rpc Error type.
+    pub fn from_tokio_join_error(error: tokio::task::JoinError) -> Error {
+        Error::TokioJoinError(error.to_string())
+    }
+}
+
 /// Response from am.i.mullvad.net
 #[derive(Debug, Serialize, Deserialize)]
 pub struct AmIMullvad {
@@ -213,6 +222,9 @@ mod service {
         /// service.
         async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(), Error>;
 
+        /// Get the environment variables for the running daemon service.
+        async fn get_daemon_environment() -> Result<HashMap<String, String>, Error>;
+
         /// Copy a file from `src` to `dest` on the test runner.
         async fn copy_file(src: String, dest: String) -> Result<(), Error>;
 
diff --git a/test/test-runner/src/main.rs b/test/test-runner/src/main.rs
index 29f0f936d3..79e11e0506 100644
--- a/test/test-runner/src/main.rs
+++ b/test/test-runner/src/main.rs
@@ -309,6 +309,13 @@ impl Service for TestServer {
         sys::set_daemon_environment(env).await
     }
 
+    async fn get_daemon_environment(
+        self,
+        _: context::Context,
+    ) -> Result<HashMap<String, String>, test_rpc::Error> {
+        sys::get_daemon_environment().await
+    }
+
     async fn copy_file(
         self,
         _: context::Context,
diff --git a/test/test-runner/src/sys.rs b/test/test-runner/src/sys.rs
index 5015f26ca1..0a176b1c3d 100644
--- a/test/test-runner/src/sys.rs
+++ b/test/test-runner/src/sys.rs
@@ -4,13 +4,22 @@ use std::io;
 use test_rpc::{meta::OsVersion, mullvad_daemon::Verbosity};
 
 #[cfg(target_os = "windows")]
-use std::ffi::OsString;
+use std::{ffi::OsString, path::Path};
 #[cfg(target_os = "windows")]
 use windows_service::{
     service::{ServiceAccess, ServiceInfo},
     service_manager::{ServiceManager, ServiceManagerAccess},
 };
 
+#[cfg(target_os = "linux")]
+const SYSTEMD_OVERRIDE_FILE: &str = "/etc/systemd/system/mullvad-daemon.service.d/override.conf";
+
+#[cfg(target_os = "macos")]
+const PLIST_OVERRIDE_FILE: &str = "/Library/LaunchDaemons/net.mullvad.daemon.plist";
+
+#[cfg(target_os = "windows")]
+const MULLVAD_WIN_REGISTRY: &str = r"SYSTEM\CurrentControlSet\Services\Mullvad VPN";
+
 #[cfg(target_os = "windows")]
 pub fn reboot() -> Result<(), test_rpc::Error> {
     use windows_sys::Win32::{
@@ -154,9 +163,6 @@ pub fn reboot() -> Result<(), test_rpc::Error> {
 #[cfg(target_os = "linux")]
 pub async fn set_daemon_log_level(verbosity_level: Verbosity) -> Result<(), test_rpc::Error> {
     use tokio::io::AsyncWriteExt;
-    const SYSTEMD_OVERRIDE_FILE: &str =
-        "/etc/systemd/system/mullvad-daemon.service.d/override.conf";
-
     log::debug!("Setting log level");
 
     let verbosity = match verbosity_level {
@@ -389,17 +395,57 @@ pub async fn set_daemon_log_level(_verbosity_level: Verbosity) -> Result<(), tes
 }
 
 #[cfg(target_os = "linux")]
+#[derive(Debug)]
+struct EnvVar {
+    var: String,
+    value: String,
+}
+
+#[cfg(target_os = "linux")]
+impl EnvVar {
+    fn from_systemd_string(s: &str) -> Result<Self, &'static str> {
+        // Here, we are only concerned with parsing a line that starts with "Environment".
+        let error = "Failed to parse systemd env-config";
+        let mut input = s.trim().split('=');
+        let pre = input.next().ok_or(error)?;
+        match pre {
+            "Environment" => {
+                // Proccess the input just a bit more - remove the leading and trailing quote (").
+                let var = input
+                    .next()
+                    .ok_or(error)?
+                    .trim_start_matches('"')
+                    .to_string();
+                let value = input.next().ok_or(error)?.trim_end_matches('"').to_string();
+                Ok(EnvVar { var, value })
+            }
+            _ => Err(error),
+        }
+    }
+
+    fn to_systemd_string(&self) -> String {
+        format!(
+            "Environment=\"{key}={value}\"",
+            key = self.var,
+            value = self.value
+        )
+    }
+}
+
+#[cfg(target_os = "linux")]
 pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(), test_rpc::Error> {
     use std::fmt::Write;
-    use tokio::io::AsyncWriteExt;
-
-    const SYSTEMD_OVERRIDE_FILE: &str = "/etc/systemd/system/mullvad-daemon.service.d/env.conf";
 
     let mut override_content = String::new();
     override_content.push_str("[Service]\n");
 
-    for (k, v) in env {
-        writeln!(&mut override_content, "Environment=\"{k}={v}\"").unwrap();
+    for env_var in env
+        .into_iter()
+        .map(|(var, value)| EnvVar { var, value })
+        .map(|env_var| env_var.to_systemd_string())
+    {
+        writeln!(&mut override_content, "{env_var}")
+            .map_err(|err| test_rpc::Error::Service(err.to_string()))?;
     }
 
     let override_path = std::path::Path::new(SYSTEMD_OVERRIDE_FILE);
@@ -409,15 +455,7 @@ pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(),
             .map_err(|e| test_rpc::Error::Service(e.to_string()))?;
     }
 
-    let mut file = tokio::fs::OpenOptions::new()
-        .create(true)
-        .truncate(true)
-        .write(true)
-        .open(override_path)
-        .await
-        .map_err(|e| test_rpc::Error::Service(e.to_string()))?;
-
-    file.write_all(override_content.as_bytes())
+    tokio::fs::write(override_path, override_content)
         .await
         .map_err(|e| test_rpc::Error::Service(e.to_string()))?;
 
@@ -440,7 +478,7 @@ pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(),
 #[cfg(target_os = "windows")]
 pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(), test_rpc::Error> {
     // Set environment globally (not for service) to prevent it from being lost on upgrade
-    for (k, v) in env {
+    for (k, v) in env.clone() {
         tokio::process::Command::new("setx")
             .arg("/m")
             .args([k, v])
@@ -448,6 +486,18 @@ pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(),
             .await
             .map_err(|e| test_rpc::Error::Registry(e.to_string()))?;
     }
+    // Persist the changed environment variables, such that we can retrieve them at will.
+    use winreg::{enums::*, RegKey};
+    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
+    let path = Path::new(MULLVAD_WIN_REGISTRY).join("Environment");
+    let (registry, _) = hklm.create_subkey(&path).map_err(|error| {
+        test_rpc::Error::Registry(format!("Failed to open Mullvad VPN subkey: {}", error))
+    })?;
+    for (k, v) in env {
+        registry.set_value(k, &v).map_err(|error| {
+            test_rpc::Error::Registry(format!("Failed to set Environment var: {}", error))
+        })?;
+    }
 
     // Restart service
     stop_app().await?;
@@ -464,7 +514,7 @@ pub fn get_system_path_var() -> Result<String, test_rpc::Error> {
     let key = hklm
         .open_subkey("SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment")
         .map_err(|error| {
-            test_rpc::Error::Registry(format!("Failed to open environment subkey: {}", error))
+            test_rpc::Error::Registry(format!("Failed to open Environment subkey: {}", error))
         })?;
 
     let path: String = key
@@ -476,14 +526,11 @@ pub fn get_system_path_var() -> Result<String, test_rpc::Error> {
 
 #[cfg(target_os = "macos")]
 pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(), test_rpc::Error> {
-    const PLIST_PATH: &str = "/Library/LaunchDaemons/net.mullvad.daemon.plist";
-
     tokio::task::spawn_blocking(|| {
-        let mut parsed_plist: plist::Value = plist::from_file(PLIST_PATH)
+        let mut parsed_plist = plist::Value::from_file(PLIST_OVERRIDE_FILE)
             .map_err(|error| test_rpc::Error::Service(format!("failed to parse plist: {error}")))?;
 
         let mut vars = plist::Dictionary::new();
-
         for (k, v) in env {
             // Set environment globally (not for service) to prevent it from being lost on upgrade
             std::process::Command::new("launchctl")
@@ -503,10 +550,7 @@ pub async fn set_daemon_environment(env: HashMap<String, String>) -> Result<(),
                 plist::Value::Dictionary(vars),
             );
 
-        let daemon_plist = std::fs::OpenOptions::new()
-            .write(true)
-            .truncate(true)
-            .open(PLIST_PATH)
+        let daemon_plist = std::fs::File::create(PLIST_OVERRIDE_FILE)
             .map_err(|e| test_rpc::Error::Service(format!("failed to open plist: {e}")))?;
 
         parsed_plist
@@ -532,7 +576,7 @@ async fn set_launch_daemon_state(on: bool) -> Result<(), test_rpc::Error> {
         .args([
             if on { "load" } else { "unload" },
             "-w",
-            "/Library/LaunchDaemons/net.mullvad.daemon.plist",
+            PLIST_OVERRIDE_FILE,
         ])
         .status()
         .await
@@ -541,6 +585,99 @@ async fn set_launch_daemon_state(on: bool) -> Result<(), test_rpc::Error> {
 }
 
 #[cfg(target_os = "linux")]
+pub async fn get_daemon_environment() -> Result<HashMap<String, String>, test_rpc::Error> {
+    let text = tokio::fs::read_to_string(SYSTEMD_OVERRIDE_FILE)
+        .await
+        .map_err(|err| test_rpc::Error::FileSystem(err.to_string()))?;
+
+    let env: HashMap<String, String> = parse_systemd_env_file(&text)
+        .map(|EnvVar { var, value }| (var, value))
+        .collect();
+    Ok(env)
+}
+
+/// Parse a systemd env-file. `input` is assumed to be the entire text content of a systemd-env file.
+///
+/// Example systemd-env file:
+/// ```
+/// [Service]
+/// Environment="VAR1=pGNqduRFkB4K9C2vijOmUDa2kPtUhArN"
+/// Environment="VAR2=JP8YLOc2bsNlrGuD6LVTq7L36obpjzxd"
+/// ```
+#[cfg(target_os = "linux")]
+fn parse_systemd_env_file(input: &str) -> impl Iterator<Item = EnvVar> + '_ {
+    input
+        .lines()
+        .map(EnvVar::from_systemd_string)
+        .filter_map(|env_var| env_var.ok())
+        .inspect(|env_var| log::trace!("Parsed {env_var:?}"))
+}
+
+#[cfg(target_os = "windows")]
+pub async fn get_daemon_environment() -> Result<HashMap<String, String>, test_rpc::Error> {
+    use winreg::{enums::*, RegKey};
+
+    let env =
+        tokio::task::spawn_blocking(|| -> Result<HashMap<String, String>, test_rpc::Error> {
+            let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
+            let key = hklm.open_subkey(MULLVAD_WIN_REGISTRY).map_err(|error| {
+                test_rpc::Error::Registry(format!("Failed to open Mullvad VPN subkey: {}", error))
+            })?;
+
+            // The Strings will be quoted (surrounded by ") when read from the registry - we should trim that!
+            let trim = |string: String| string.trim_matches('"').to_owned();
+            let env = key
+                .open_subkey("Environment")
+                .map_err(|error| {
+                    test_rpc::Error::Registry(format!(
+                        "Failed to open Environment subkey: {}",
+                        error
+                    ))
+                })?
+                .enum_values()
+                .filter_map(|x| x.inspect_err(|err| log::trace!("{err}")).ok())
+                .map(|(k, v)| (trim(k), trim(v.to_string())))
+                .collect();
+            Ok(env)
+        })
+        .await
+        .map_err(test_rpc::Error::from_tokio_join_error)??;
+
+    Ok(env)
+}
+
+#[cfg(target_os = "macos")]
+pub async fn get_daemon_environment() -> Result<HashMap<String, String>, test_rpc::Error> {
+    let plist = tokio::task::spawn_blocking(|| {
+        let parsed_plist = plist::Value::from_file(PLIST_OVERRIDE_FILE)
+            .map_err(|error| test_rpc::Error::Service(format!("failed to parse plist: {error}")))?;
+
+        Ok::<plist::Value, test_rpc::Error>(parsed_plist)
+    })
+    .await
+    .map_err(test_rpc::Error::from_tokio_join_error)??;
+
+    let plist_tree = plist
+        .as_dictionary()
+        .ok_or_else(|| test_rpc::Error::Service("plist missing dict".to_owned()))?;
+    let Some(env_vars) = plist_tree.get("EnvironmentVariables") else {
+        // `EnvironmentVariables` does not exist in plist file, so there are no env variables to parse.
+        return Ok(HashMap::new());
+    };
+    let env_vars = env_vars.as_dictionary().ok_or_else(|| {
+        test_rpc::Error::Service("`EnvironmentVariables` is not a dict".to_owned())
+    })?;
+
+    let env = env_vars
+        .clone()
+        .into_iter()
+        .filter_map(|(key, value)| Some((key, value.into_string()?)))
+        .collect();
+
+    Ok(env)
+}
+
+#[cfg(target_os = "linux")]
 enum ServiceState {
     Running,
     Inactive,
@@ -618,3 +755,31 @@ pub fn get_os_version() -> Result<OsVersion, test_rpc::Error> {
 pub fn get_os_version() -> Result<OsVersion, test_rpc::Error> {
     Ok(OsVersion::Linux)
 }
+
+#[cfg(test)]
+mod test {
+
+    #[cfg(target_os = "linux")]
+    #[test]
+    fn parse_systemd_environment_variables() {
+        use super::parse_systemd_env_file;
+        // Define an example systemd environment file
+        let systemd_file = "
+        [Service]
+        Environment=\"var1=value1\"
+        Environment=\"var2=value2\"
+        ";
+
+        // Parse the "file"
+        let env_vars: Vec<_> = parse_systemd_env_file(systemd_file).collect();
+
+        // Assert that the environment variables it defines are parsed as expected.
+        assert_eq!(env_vars.len(), 2);
+        let first = env_vars.first().unwrap();
+        assert_eq!(first.var, "var1");
+        assert_eq!(first.value, "value1");
+        let second = env_vars.get(1).unwrap();
+        assert_eq!(second.var, "var2");
+        assert_eq!(second.value, "value2");
+    }
+}