diff --git a/crates/security-engine/tests/windows_tests.rs b/crates/security-engine/tests/windows_tests.rs index 74c70c7..343f9d2 100644 --- a/crates/security-engine/tests/windows_tests.rs +++ b/crates/security-engine/tests/windows_tests.rs @@ -1,4 +1,3 @@ - #[cfg(windows)] #[test] fn test_get_vessel_credentials_consistency() { diff --git a/crates/vessel/src/main.rs b/crates/vessel/src/main.rs index f19ef8e..f3b64cf 100644 --- a/crates/vessel/src/main.rs +++ b/crates/vessel/src/main.rs @@ -114,6 +114,22 @@ async fn main() -> Result<()> { } }; + tokio::spawn(async move { + if let Ok(_) = tokio::signal::ctrl_c().await { + log::info!("Received Ctrl-C, initiating shutdown sequence..."); + sanitation::restore_system_defaults(); + std::process::exit(0); + } + }); + + struct ShutdownGuard; + impl Drop for ShutdownGuard { + fn drop(&mut self) { + sanitation::restore_system_defaults(); + } + } + let _shutdown_guard = ShutdownGuard; + sanitation::clean_high_risk_directories(); // Create secondary desktop before security initialization so that @@ -132,6 +148,7 @@ async fn main() -> Result<()> { Commands::Stop => { log::info!("Stopping Vessel agent..."); // TODO: Clean up isolated desktop sessions and guest applications + sanitation::restore_system_defaults(); } Commands::Credentials(args) => { match args.command { diff --git a/crates/vessel/src/sanitation.rs b/crates/vessel/src/sanitation.rs index 2f87e07..e9b6ff9 100644 --- a/crates/vessel/src/sanitation.rs +++ b/crates/vessel/src/sanitation.rs @@ -9,16 +9,21 @@ mod windows_impl { type DWORD = u32; type WCHAR = u16; + type PVOID = *mut std::ffi::c_void; + const MAXIMUM_ALLOWED: DWORD = 0x02000000; + const SPI_GETFOREGROUNDLOCKTIMEOUT: DWORD = 0x2000; + const SPI_SETFOREGROUNDLOCKTIMEOUT: DWORD = 0x2001; + const SPIF_UPDATEINIFILE: DWORD = 0x01; + const SPIF_SENDCHANGE: DWORD = 0x02; type DESKTOPENUMPROCW = - Option BOOL>; + Option BOOL>; type WINSTAENUMPROCW = - Option BOOL>; + Option BOOL>; #[link(name = "user32")] extern "system" { - fn GetProcessWindowStation() -> HWINSTA; fn EnumWindowStationsW(lpEnumFunc: WINSTAENUMPROCW, lParam: LPARAM) -> BOOL; fn EnumDesktopsW(hwinsta: HWINSTA, lpEnumFunc: DESKTOPENUMPROCW, lParam: LPARAM) -> BOOL; fn OpenWindowStationW( @@ -34,6 +39,19 @@ mod windows_impl { dwDesiredAccess: DWORD, ) -> HDESK; fn CloseDesktop(hDesktop: HDESK) -> BOOL; + fn SystemParametersInfoW( + uiAction: DWORD, + uiParam: DWORD, + pvParam: PVOID, + fWinIni: DWORD, + ) -> BOOL; + } + + static ORIGINAL_FOREGROUND_LOCK_TIMEOUT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + static TIMEOUT_CAPTURED: std::sync::atomic::AtomicBool = std::sync::atomic::AtomicBool::new(false); + + fn get_timeout_backup_path() -> std::path::PathBuf { + std::env::temp_dir().join("vessel_foreground_timeout.txt") } unsafe extern "system" fn enum_desktops_callback( @@ -90,6 +108,30 @@ mod windows_impl { let count = AtomicUsize::new(0); unsafe { + let mut timeout: DWORD = 0; + let res = SystemParametersInfoW( + SPI_GETFOREGROUNDLOCKTIMEOUT, + 0, + &mut timeout as *mut _ as PVOID, + 0, + ); + if res != 0 { + ORIGINAL_FOREGROUND_LOCK_TIMEOUT.store(timeout as usize, Ordering::SeqCst); + TIMEOUT_CAPTURED.store(true, Ordering::SeqCst); + let _ = std::fs::write(get_timeout_backup_path(), timeout.to_string()); + log::info!("Sanitation: Captured original ForegroundLockTimeout: {} ms", timeout); + + let new_timeout: DWORD = 0; // 0 allows focus stealing, preventing taskbar flash + SystemParametersInfoW( + SPI_SETFOREGROUNDLOCKTIMEOUT, + 0, + new_timeout as usize as PVOID, + SPIF_UPDATEINIFILE | SPIF_SENDCHANGE, + ); + } else { + log::warn!("Sanitation: Failed to capture original ForegroundLockTimeout"); + } + EnumWindowStationsW( Some(enum_window_stations_callback), &count as *const _ as LPARAM, @@ -103,6 +145,60 @@ mod windows_impl { Ok(()) } + + pub fn restore_timeout() { + let mut timeout_to_restore = None; + + if TIMEOUT_CAPTURED.load(Ordering::SeqCst) { + timeout_to_restore = Some(ORIGINAL_FOREGROUND_LOCK_TIMEOUT.load(Ordering::SeqCst) as DWORD); + } else { + // Try to read from backup file in case this is a separate process (e.g. vessel stop) + if let Ok(content) = std::fs::read_to_string(get_timeout_backup_path()) { + if let Ok(parsed) = content.trim().parse::() { + timeout_to_restore = Some(parsed); + } + } + } + + if let Some(timeout) = timeout_to_restore { + unsafe { + SystemParametersInfoW( + SPI_SETFOREGROUNDLOCKTIMEOUT, + 0, + timeout as usize as PVOID, + SPIF_UPDATEINIFILE | SPIF_SENDCHANGE, + ); + } + let _ = std::fs::remove_file(get_timeout_backup_path()); + log::info!("Sanitation: Restored original ForegroundLockTimeout: {} ms", timeout); + } + } + + #[cfg(test)] + mod tests { + use super::*; + use std::sync::atomic::Ordering; + + #[test] + fn test_get_timeout_backup_path() { + let path = get_timeout_backup_path(); + assert!(path.to_string_lossy().contains("vessel_foreground_timeout.txt")); + } + + #[test] + fn test_timeout_captured_initial_state() { + // Because this is global state, we can't assume it's strictly false if other tests run, + // but we can test we can read it. + let _ = TIMEOUT_CAPTURED.load(Ordering::SeqCst); + } + } +} + +pub fn restore_system_defaults() { + #[cfg(windows)] + { + windows_impl::restore_timeout(); + } } pub fn sanitize_vessel_desktops() -> anyhow::Result<()> { @@ -207,3 +303,21 @@ pub fn clean_high_risk_directories() { wipe_dir(&target); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_restore_system_defaults_does_not_panic() { + // This should run smoothly and not panic on any platform. + restore_system_defaults(); + } + + #[test] + fn test_sanitize_vessel_desktops_does_not_panic() { + // This should not panic + let res = sanitize_vessel_desktops(); + assert!(res.is_ok()); + } +}