use std::collections::HashMap; use std::ffi::OsString; use std::io; use std::path::Path; use test_rpc::mullvad_daemon::ServiceStatus; use test_rpc::{meta::OsVersion, mullvad_daemon::Verbosity}; use windows_service::{ service::{Service, ServiceAccess, ServiceInfo, ServiceState}, service_manager::{ServiceManager, ServiceManagerAccess}, }; use windows_sys::Win32::{ System::Shutdown::{ EWX_REBOOT, ExitWindowsEx, SHTDN_REASON_FLAG_PLANNED, SHTDN_REASON_MAJOR_APPLICATION, SHTDN_REASON_MINOR_OTHER, }, UI::WindowsAndMessaging::EWX_FORCEIFHUNG, }; const MULLVAD_WIN_REGISTRY: &str = r"SYSTEM\CurrentControlSet\Services\Mullvad VPN"; pub fn reboot() -> Result<(), test_rpc::Error> { grant_shutdown_privilege()?; std::thread::spawn(|| { std::thread::sleep(std::time::Duration::from_secs(5)); let shutdown_result = unsafe { ExitWindowsEx( EWX_FORCEIFHUNG | EWX_REBOOT, SHTDN_REASON_MAJOR_APPLICATION | SHTDN_REASON_MINOR_OTHER | SHTDN_REASON_FLAG_PLANNED, ) }; if shutdown_result == 0 { log::error!( "Failed to restart test machine: {}", io::Error::last_os_error() ); std::process::exit(1); } std::process::exit(0); }); // NOTE: We do not bother to revoke the privilege. Ok(()) } fn grant_shutdown_privilege() -> Result<(), test_rpc::Error> { use windows_sys::Win32::{ Foundation::{CloseHandle, HANDLE, LUID}, Security::{ AdjustTokenPrivileges, LUID_AND_ATTRIBUTES, LookupPrivilegeValueW, SE_PRIVILEGE_ENABLED, TOKEN_ADJUST_PRIVILEGES, TOKEN_PRIVILEGES, }, System::{ SystemServices::SE_SHUTDOWN_NAME, Threading::{GetCurrentProcess, OpenProcessToken}, }, }; let mut privileges = TOKEN_PRIVILEGES { PrivilegeCount: 1, Privileges: [LUID_AND_ATTRIBUTES { Luid: LUID { HighPart: 0, LowPart: 0, }, Attributes: SE_PRIVILEGE_ENABLED, }], }; if unsafe { LookupPrivilegeValueW( std::ptr::null(), SE_SHUTDOWN_NAME, &mut privileges.Privileges[0].Luid, ) } == 0 { log::error!( "Failed to lookup shutdown privilege LUID: {}", io::Error::last_os_error() ); return Err(test_rpc::Error::Syscall); } let mut token_handle: HANDLE = 0; if unsafe { OpenProcessToken( GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &mut token_handle, ) } == 0 { log::error!("OpenProcessToken() failed: {}", io::Error::last_os_error()); return Err(test_rpc::Error::Syscall); } let result = unsafe { AdjustTokenPrivileges( token_handle, 0, &privileges, 0, std::ptr::null_mut(), std::ptr::null_mut(), ) }; unsafe { CloseHandle(token_handle) }; if result == 0 { log::error!( "Failed to enable SE_SHUTDOWN_NAME: {}", io::Error::last_os_error() ); return Err(test_rpc::Error::Syscall); } Ok(()) } /// Restart the Mullvad VPN application. /// /// This function waits for the app to successfully start again. pub async fn restart_app() -> Result<(), test_rpc::Error> { stop_app().await?; start_app().await?; Ok(()) } /// Stop the Mullvad VPN application. /// /// This function waits for the app to successfully shut down. pub async fn stop_app() -> Result<(), test_rpc::Error> { let _ = tokio::process::Command::new("net") .args(["stop", "mullvadvpn"]) .status() .await .map_err(|e| test_rpc::Error::ServiceStop(e.to_string()))?; Ok(()) } /// Start the Mullvad VPN application. /// /// This function waits for the app to successfully start again. pub async fn start_app() -> Result<(), test_rpc::Error> { let _ = tokio::process::Command::new("net") .args(["start", "mullvadvpn"]) .status() .await .map_err(|e| test_rpc::Error::ServiceStart(e.to_string()))?; Ok(()) } /// Disable the Mullvad VPN system service startup. This will not trigger the service to stop /// immediately, but it will prevent it from starting on the next system boot. pub async fn disable_system_service_startup() -> Result<(), test_rpc::Error> { let status = tokio::process::Command::new("powershell") .args([ "-NoProfile", "-Command", "Set-Service -Name MullvadVPN -StartupType Disabled", ]) .status() .await .map_err(|e| test_rpc::Error::ServiceChange(e.to_string()))?; if !status.success() { return Err(test_rpc::Error::ServiceChange( "Failed to disable MullvadVPN service".to_string(), )); } Ok(()) } /// Enable the Mullvad VPN system service startup. This will configure the service to start automatically on system boot. pub async fn enable_system_service_startup() -> Result<(), test_rpc::Error> { let status = tokio::process::Command::new("powershell") .args([ "-NoProfile", "-Command", "Set-Service -Name MullvadVPN -StartupType Automatic", ]) .status() .await .map_err(|e| test_rpc::Error::ServiceChange(e.to_string()))?; if !status.success() { return Err(test_rpc::Error::ServiceChange( "Failed to enable MullvadVPN service".to_string(), )); } Ok(()) } pub fn get_daemon_system_service_status() -> Result { let manager = ServiceManager::local_computer(None::<&str>, ServiceManagerAccess::CONNECT) .map_err(|e| test_rpc::Error::ServiceNotFound(e.to_string()))?; let service = manager .open_service("mullvadvpn", ServiceAccess::QUERY_STATUS) .map_err(|e| test_rpc::Error::ServiceNotFound(e.to_string()))?; get_daemon_system_service_status_inner(&service) } fn get_daemon_system_service_status_inner( service: &Service, ) -> Result { let status = service .query_status() .map_err(|e| test_rpc::Error::Other(e.to_string()))?; let status = match status.current_state { ServiceState::Running => ServiceStatus::Running, // NOTE: not counting pending start as running, since we cannot set log level then _ => ServiceStatus::NotRunning, }; Ok(status) } async fn wait_for_service_status( service: &Service, accept_fn: impl Fn(&windows_service::service::ServiceStatus) -> bool, ) -> Result<(), test_rpc::Error> { const MAX_ATTEMPTS: usize = 10; const POLL_INTERVAL: std::time::Duration = std::time::Duration::from_secs(3); for _ in 0..MAX_ATTEMPTS { let status = service .query_status() .map_err(|e| test_rpc::Error::Other(e.to_string()))?; if accept_fn(&status) { return Ok(()); } tokio::time::sleep(POLL_INTERVAL).await; } Err(test_rpc::Error::ServiceStart( "Awaiting new service state timed out".to_string(), )) } pub async fn set_daemon_log_level(verbosity_level: Verbosity) -> Result<(), test_rpc::Error> { use std::error::Error; fn error_with_source(e: &impl Error) -> String { if let Some(source) = e.source() { format!("{e}: {source}") } else { e.to_string() } } log::debug!("Setting log level"); let verbosity = match verbosity_level { Verbosity::Info => "", Verbosity::Debug => "-v", Verbosity::Trace => "-vv", }; let manager = ServiceManager::local_computer(None::<&str>, ServiceManagerAccess::CONNECT) .map_err(|e| test_rpc::Error::ServiceNotFound(e.to_string()))?; let service = manager .open_service( "mullvadvpn", ServiceAccess::QUERY_CONFIG | ServiceAccess::QUERY_STATUS | ServiceAccess::CHANGE_CONFIG | ServiceAccess::START | ServiceAccess::STOP, ) .map_err(|e| { test_rpc::Error::ServiceNotFound(format!( "Failed to open service: {}", error_with_source(&e) )) })?; log::info!("Stopping service"); // Stop the service service .stop() .map_err(|e| test_rpc::Error::ServiceStop(e.to_string()))?; // Wait until the service is fully stopped wait_for_service_status(&service, |status| { status.current_state == ServiceState::Stopped }) .await?; // Get the current service configuration let config = service.query_config().map_err(|e| { test_rpc::Error::ServiceNotFound(format!( "Failed to query service config: {}", error_with_source(&e) )) })?; let executable_path = "C:\\Program Files\\Mullvad VPN\\resources\\mullvad-daemon.exe"; let launch_arguments = vec![ OsString::from("--run-as-service"), OsString::from(verbosity), ]; // Update the service binary arguments let updated_config = ServiceInfo { name: config.display_name.clone(), display_name: config.display_name.clone(), service_type: config.service_type, start_type: config.start_type, error_control: config.error_control, executable_path: std::path::PathBuf::from(executable_path), launch_arguments, dependencies: config.dependencies.clone(), account_name: config.account_name.clone(), account_password: None, }; // Apply the updated configuration service.change_config(&updated_config).map_err(|e| { test_rpc::Error::ServiceChange(format!("Update service config: {}", error_with_source(&e))) })?; // Start the service service.start::(&[]).map_err(|e| { test_rpc::Error::ServiceNotFound(format!( "Failed to start service: {}", error_with_source(&e) )) })?; // Wait until the service is fully started wait_for_service_status(&service, |status| { status.current_state == ServiceState::Running }) .await?; Ok(()) } 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.clone() { tokio::process::Command::new("setx") .arg("/m") .args([k, v]) .status() .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::{RegKey, enums::*}; 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?; start_app().await?; Ok(()) } pub fn get_system_path_var() -> Result { use winreg::{enums::*, *}; let hklm = RegKey::predef(HKEY_LOCAL_MACHINE); 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}")) })?; let path: String = key .get_value("Path") .map_err(|error| test_rpc::Error::Registry(format!("Failed to get PATH: {error}")))?; Ok(path) } pub async fn get_daemon_environment() -> Result, test_rpc::Error> { use winreg::{RegKey, enums::*}; let env = tokio::task::spawn_blocking(|| -> Result, 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) } pub fn get_os_version() -> Result { use test_rpc::meta::WindowsVersion; let version = talpid_platform_metadata::WindowsVersion::new() .inspect_err(|error| { log::error!("Failed to obtain OS version: {error}"); }) .map_err(|_| test_rpc::Error::Syscall)?; Ok(OsVersion::Windows(WindowsVersion { major: version.release_version().0, })) }