From e86ac2e10ad4c1e4e69f6fdab6e0d230ad0c9b31 Mon Sep 17 00:00:00 2001 From: Markus Pettersson Date: Tue, 9 Apr 2024 13:48:07 +0200 Subject: [PATCH] Reset daemon environment when needed --- test/test-manager/src/run_tests.rs | 4 +- test/test-manager/src/tests/mod.rs | 34 +++- test/test-rpc/src/client.rs | 14 ++ test/test-rpc/src/lib.rs | 3 + test/test-runner/src/main.rs | 7 + test/test-runner/src/sys.rs | 241 ++++++++++++++++++++++++++--- 6 files changed, 277 insertions(+), 26 deletions(-) diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs index 0846f4bf32c0..697406e40112 100644 --- a/test/test-manager/src/run_tests.rs +++ b/test/test-manager/src/run_tests.rs @@ -107,8 +107,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 bb7b8aa749a2..87aa139aaa7d 100644 --- a/test/test-manager/src/tests/mod.rs +++ b/test/test-manager/src/tests/mod.rs @@ -13,14 +13,17 @@ 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; +use mullvad_management_interface::MullvadProxyClient; 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); @@ -62,10 +65,15 @@ pub enum Error { } /// 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 + let mut mullvad_client = restart_daemon(rpc, rpc_provider).await?; - helpers::disconnect_and_wait(mullvad_client).await?; + helpers::disconnect_and_wait(&mut mullvad_client).await?; let default_settings = mullvad_types::settings::Settings::default(); @@ -122,3 +130,21 @@ pub async fn cleanup_after_test(mullvad_client: &mut MullvadProxyClient) -> anyh Ok(()) } + +/// Conditonally 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, + rpc_provider: &RpcClientProvider, +) -> anyhow::Result { + let current_env = rpc.get_daemon_environment().await?; + let default_env = get_app_env(); + if current_env != default_env { + rpc.set_daemon_environment(default_env).await?; + } + let mullvad_client = rpc_provider.new_client().await; + Ok(mullvad_client) +} diff --git a/test/test-rpc/src/client.rs b/test/test-rpc/src/client.rs index 79c6a038ca96..433a2e78f83e 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, Error> { + let mut ctx = tarpc::context::current(); + ctx.deadline = SystemTime::now().checked_add(LOG_LEVEL_TIMEOUT).unwrap(); + let env = self.client.get_daemon_environment(ctx).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 fc9ea67c8449..3673a215d954 100644 --- a/test/test-rpc/src/lib.rs +++ b/test/test-rpc/src/lib.rs @@ -213,6 +213,9 @@ mod service { /// service. async fn set_daemon_environment(env: HashMap) -> Result<(), Error>; + /// Get the environment variables for the running daemon service. + async fn get_daemon_environment() -> Result, 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 29f0f936d363..79e11e0506f5 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, 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 5015f26ca126..e2c647ed7ccf 100644 --- a/test/test-runner/src/sys.rs +++ b/test/test-runner/src/sys.rs @@ -1,6 +1,6 @@ -use std::collections::HashMap; #[cfg(target_os = "windows")] use std::io; +use std::{collections::HashMap, path::Path, str::FromStr}; use test_rpc::{meta::OsVersion, mullvad_daemon::Verbosity}; #[cfg(target_os = "windows")] @@ -154,6 +154,7 @@ 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; + // TODO(markus): Only define once const SYSTEMD_OVERRIDE_FILE: &str = "/etc/systemd/system/mullvad-daemon.service.d/override.conf"; @@ -388,6 +389,65 @@ pub async fn set_daemon_log_level(_verbosity_level: Verbosity) -> Result<(), tes Ok(()) } +// TODO(markus): Move somewhere sane! +#[derive(Debug)] +struct EnvVar { + var: String, + value: String, +} + +impl, V: Into> From<(K, V)> for EnvVar { + fn from(value: (K, V)) -> Self { + Self { + var: value.0.into(), + value: value.1.into(), + } + } +} + +// TODO(markus): This is only valid on Linux +impl std::str::FromStr for EnvVar { + type Err = &'static str; + + fn from_str(s: &str) -> Result { + // 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" => { + // Pre-proccess the input just a bit more - remove the leading and trailing quote ("). + let var = { + let var = input.next().ok_or(error)?; + var[1..].to_string() + }; + let value = { + let value = input.next().ok_or(error)?; + value[..(value.len() - 1)].to_string() + }; + Ok(EnvVar { var, value }) + } + _ => Err(error), + } + } +} + +impl EnvVar { + /// Convert `self` into a tuple of its parts. + pub fn tuple(self) -> (String, String) { + (self.var, self.value) + } + + #[cfg(target_os = "linux")] + 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) -> Result<(), test_rpc::Error> { use std::fmt::Write; @@ -398,8 +458,9 @@ pub async fn set_daemon_environment(env: HashMap) -> Result<(), 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 var in env.iter().map(EnvVar::from) { + writeln!(&mut override_content, "{}", var.to_systemd_string()) + .map_err(|err| test_rpc::Error::Service(err.to_string()))?; } let override_path = std::path::Path::new(SYSTEMD_OVERRIDE_FILE); @@ -409,15 +470,10 @@ pub async fn set_daemon_environment(env: HashMap) -> 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) + tokio::fs::File::create(override_path) .await - .map_err(|e| test_rpc::Error::Service(e.to_string()))?; - - file.write_all(override_content.as_bytes()) + .map_err(|e| test_rpc::Error::Service(e.to_string()))? + .write_all(override_content.as_bytes()) .await .map_err(|e| test_rpc::Error::Service(e.to_string()))?; @@ -440,7 +496,7 @@ pub async fn set_daemon_environment(env: HashMap) -> Result<(), #[cfg(target_os = "windows")] pub async fn set_daemon_environment(env: HashMap) -> 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 +504,19 @@ pub async fn set_daemon_environment(env: HashMap) -> 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); + // TODO(markus): Only define this path once! + let path = Path::new(r"SYSTEM\CurrentControlSet\Services\Mullvad VPN").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 +533,7 @@ pub fn get_system_path_var() -> Result { 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 @@ -479,11 +548,10 @@ pub async fn set_daemon_environment(env: HashMap) -> Result<(), 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_PATH) .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 +571,7 @@ pub async fn set_daemon_environment(env: HashMap) -> 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_PATH) .map_err(|e| test_rpc::Error::Service(format!("failed to open plist: {e}")))?; parsed_plist @@ -528,11 +593,13 @@ pub async fn set_daemon_environment(env: HashMap) -> Result<(), #[cfg(target_os = "macos")] async fn set_launch_daemon_state(on: bool) -> Result<(), test_rpc::Error> { + // TODO: Only define this constant once + const PLIST_OVERRIDE_FILE: &str = "/Library/LaunchDaemons/net.mullvad.daemon.plist"; tokio::process::Command::new("launchctl") .args([ if on { "load" } else { "unload" }, "-w", - "/Library/LaunchDaemons/net.mullvad.daemon.plist", + PLIST_OVERRIDE_FILE, ]) .status() .await @@ -540,6 +607,113 @@ async fn set_launch_daemon_state(on: bool) -> Result<(), test_rpc::Error> { Ok(()) } +#[cfg(target_os = "linux")] +pub async fn get_daemon_environment() -> Result, test_rpc::Error> { + use tokio::fs::File; + use tokio::io::AsyncReadExt; + const SYSTEMD_OVERRIDE_FILE: &str = "/etc/systemd/system/mullvad-daemon.service.d/env.conf"; + + let mut file = File::open(SYSTEMD_OVERRIDE_FILE) + .await + .map_err(|err| test_rpc::Error::FileSystem(err.to_string()))?; + let text = { + let mut buf = String::new(); + file.read_to_string(&mut buf) + .await + .map_err(|err| test_rpc::Error::FileSystem(err.to_string()))?; + buf + }; + + let env: HashMap = parse_systemd_env_file(&text) + .into_iter() + .map(EnvVar::tuple) + .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" +fn parse_systemd_env_file(input: &str) -> impl IntoIterator + '_ { + input + .lines() + // Skip the [Service] line + .skip(1) + .map(EnvVar::from_str) + .filter_map(|env_var| env_var.ok()) + .inspect(|env_var| println!("Parsed {env_var:?}")) +} + +#[cfg(target_os = "windows")] +pub async fn get_daemon_environment() -> Result, test_rpc::Error> { + use winreg::{enums::*, RegKey}; + + let hklm = RegKey::predef(HKEY_LOCAL_MACHINE); + let key = hklm + .open_subkey(r"SYSTEM\CurrentControlSet\Services\Mullvad VPN") + .map_err(|error| { + test_rpc::Error::Registry(format!("Failed to open Mullvad VPN subkey: {}", error)) + })?; + + let trim = |string: String| { + string + .trim_start_matches('"') + .trim_end_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()) + // The Strings will be quoted (surrounded by "") when read from the registry - we should + // trim that! + .map(|(k, v)| (trim(k), trim(v.to_string()))) + .collect(); + + Ok(env) +} + +#[cfg(target_os = "macos")] +pub async fn get_daemon_environment() -> Result, test_rpc::Error> { + // TODO: Only define this constant once + const PLIST_PATH: &str = "/Library/LaunchDaemons/net.mullvad.daemon.plist"; + + let plist = tokio::task::spawn_blocking(|| { + let parsed_plist = plist::Value::from_file(PLIST_PATH) + .map_err(|error| test_rpc::Error::Service(format!("failed to parse plist: {error}")))?; + + Ok::(parsed_plist) + }) + .await + .unwrap()?; + + 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() + // TODO(markus): Do not unwrap! + .map(|(key, value)| (key, value.into_string().unwrap())) + .collect(); + + Ok(env) +} + #[cfg(target_os = "linux")] enum ServiceState { Running, @@ -618,3 +792,30 @@ pub fn get_os_version() -> Result { pub fn get_os_version() -> Result { Ok(OsVersion::Linux) } + +#[cfg(test)] +mod test { + + #[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).into_iter().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"); + } +}