From 85fe20ebba4fd5832f233efb75372ce4f5644ebf Mon Sep 17 00:00:00 2001 From: LGUG2Z Date: Thu, 2 Dec 2021 18:44:41 -0800 Subject: [PATCH] refactor(wm): validate virtual desktops via reg This commit refactors the validations that ensure that only commands and events originating on the same virtual desktop that komorebi was started on are managed. This was previously handled by the winvd crate which relied on undocumented APIs that broke as of Windows 11. This method, while not very elegant, seems like the best solution for now. In short, komorebi checks the registry (which has different paths on Win10 and Win11...) to keep track of the current virtual desktop id. This is problematic because we just end up comparing byte arrays, and there is no meaningful representation of the ids that are being compared, not even a GUID. Nevertheless, it works and it ensures that komorebi is limited to operating on a single virtual desktop. resolve #77 --- Cargo.lock | 44 +++++---------------------------- komorebi/Cargo.toml | 3 ++- komorebi/src/main.rs | 32 ++++++++++++++++++++++++ komorebi/src/process_command.rs | 11 ++++++--- komorebi/src/process_event.rs | 11 ++++++--- komorebi/src/window_manager.rs | 27 +++----------------- komorebi/src/windows_api.rs | 14 +++++++++++ 7 files changed, 73 insertions(+), 69 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 53f4b8a9..addfa728 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -141,37 +141,6 @@ dependencies = [ "tracing-error", ] -[[package]] -name = "com" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a30a2b2a013da986dc5cc3eda3d19c0d59d53f835be1b2356eb8d00f000c793" -dependencies = [ - "com_macros", -] - -[[package]] -name = "com_macros" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7606b05842fea68ddcc89e8053b8860ebcb2a0ba8d6abfe3a148e5d5a8d3f0c1" -dependencies = [ - "com_macros_support", - "proc-macro2", - "syn", -] - -[[package]] -name = "com_macros_support" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "97e9a6d20f4ac8830e309a455d7e9416e65c6af5a97c88c55fbb4c2012e107da" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - [[package]] name = "core-foundation-sys" version = "0.8.3" @@ -514,7 +483,7 @@ dependencies = [ "which", "windows", "winput", - "winvd", + "winreg", ] [[package]] @@ -1524,13 +1493,12 @@ dependencies = [ ] [[package]] -name = "winvd" -version = "0.0.20" -source = "git+https://github.com/Ciantic/VirtualDesktopAccessor?branch=rust#11502c442f0cc0e6b5304e6a45022406ac502842" +name = "winreg" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80d0f4e272c85def139476380b12f9ac60926689dd2e01d4923222f40580869d" dependencies = [ - "com", - "crossbeam-channel", - "once_cell", + "winapi 0.3.9", ] [[package]] diff --git a/komorebi/Cargo.toml b/komorebi/Cargo.toml index 0dbcea57..4c2f3bb6 100644 --- a/komorebi/Cargo.toml +++ b/komorebi/Cargo.toml @@ -36,8 +36,8 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] } uds_windows = "1" which = "4" winput = "0.2" -winvd = { git = "https://github.com/Ciantic/VirtualDesktopAccessor", branch = "rust" } miow = "0.4" +winreg = "0.10" [dependencies.windows] version = "0.28" @@ -47,6 +47,7 @@ features = [ "Win32_Graphics_Dwm", "Win32_Graphics_Gdi", "Win32_System_Threading", + "Win32_System_RemoteDesktop", "Win32_UI_Input_KeyboardAndMouse", "Win32_UI_Accessibility", "Win32_UI_WindowsAndMessaging" diff --git a/komorebi/src/main.rs b/komorebi/src/main.rs index 0da33069..fcf11458 100644 --- a/komorebi/src/main.rs +++ b/komorebi/src/main.rs @@ -6,6 +6,7 @@ use std::fs::File; use std::io::Write; use std::process::Command; use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU32; use std::sync::atomic::Ordering; use std::sync::Arc; #[cfg(feature = "deadlock_detection")] @@ -28,6 +29,8 @@ use tracing_appender::non_blocking::WorkerGuard; use tracing_subscriber::layer::SubscriberExt; use tracing_subscriber::EnvFilter; use which::which; +use winreg::enums::HKEY_CURRENT_USER; +use winreg::RegKey; use komorebi_core::HidingBehaviour; use komorebi_core::SocketMessage; @@ -98,6 +101,7 @@ lazy_static! { } pub static CUSTOM_FFM: AtomicBool = AtomicBool::new(false); +pub static SESSION_ID: AtomicU32 = AtomicU32::new(0); fn setup() -> Result<(WorkerGuard, WorkerGuard)> { if std::env::var("RUST_LIB_BACKTRACE").is_err() { @@ -198,6 +202,31 @@ pub fn load_configuration() -> Result<()> { Ok(()) } +pub fn current_virtual_desktop() -> Result> { + let hkcu = RegKey::predef(HKEY_CURRENT_USER); + + // This is the path on Windows 10 + let current = match hkcu.open_subkey(format!( + r#"SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\SessionInfo\{}\VirtualDesktops"#, + SESSION_ID.load(Ordering::SeqCst) + )) { + Ok(desktops) => { + if let Ok(current) = desktops.get_raw_value("CurrentVirtualDesktop") { + current.bytes + } else { + // This is the path on Windows 11 + let desktops = hkcu.open_subkey( + r#"SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VirtualDesktops"#, + )?; + desktops.get_raw_value("CurrentVirtualDesktop")?.bytes + } + } + Err(_) => unreachable!(), + }; + + Ok(current) +} + #[derive(Debug, Serialize)] #[serde(untagged)] pub enum NotificationEvent { @@ -285,6 +314,9 @@ fn main() -> Result<()> { let has_valid_args = arg_count == 1 || (arg_count == 2 && CUSTOM_FFM.load(Ordering::SeqCst)); if has_valid_args { + let session_id = WindowsApi::process_id_to_session_id()?; + SESSION_ID.store(session_id, Ordering::SeqCst); + let mut system = sysinfo::System::new_all(); system.refresh_processes(); diff --git a/komorebi/src/process_command.rs b/komorebi/src/process_command.rs index 66d52227..c85d0ff9 100644 --- a/komorebi/src/process_command.rs +++ b/komorebi/src/process_command.rs @@ -68,9 +68,14 @@ pub fn listen_for_commands(wm: Arc>) { impl WindowManager { #[tracing::instrument(skip(self))] pub fn process_command(&mut self, message: SocketMessage) -> Result<()> { - if let Err(error) = self.validate_virtual_desktop_id() { - tracing::info!("{}", error); - return Ok(()); + if let Ok(id) = crate::current_virtual_desktop() { + if id != self.virtual_desktop_id { + tracing::info!( + "ignoring events and commands while not on virtual desktop {:?}", + self.virtual_desktop_id + ); + return Ok(()); + } } match message { diff --git a/komorebi/src/process_event.rs b/komorebi/src/process_event.rs index da1c05f4..34bc9e7c 100644 --- a/komorebi/src/process_event.rs +++ b/komorebi/src/process_event.rs @@ -51,9 +51,14 @@ impl WindowManager { return Ok(()); } - if let Err(error) = self.validate_virtual_desktop_id() { - tracing::info!("{}", error); - return Ok(()); + if let Ok(id) = crate::current_virtual_desktop() { + if id != self.virtual_desktop_id { + tracing::info!( + "ignoring events and commands while not on virtual desktop {:?}", + self.virtual_desktop_id + ); + return Ok(()); + } } // Make sure we have the most recently focused monitor from any event diff --git a/komorebi/src/window_manager.rs b/komorebi/src/window_manager.rs index 25c8262f..c23022e6 100644 --- a/komorebi/src/window_manager.rs +++ b/komorebi/src/window_manager.rs @@ -27,6 +27,7 @@ use komorebi_core::Sizing; use komorebi_core::WindowContainerBehaviour; use crate::container::Container; +use crate::current_virtual_desktop; use crate::load_configuration; use crate::monitor::Monitor; use crate::ring::Ring; @@ -55,7 +56,7 @@ pub struct WindowManager { pub focus_follows_mouse: Option, pub mouse_follows_focus: bool, pub hotwatch: Hotwatch, - pub virtual_desktop_id: Option, + pub virtual_desktop_id: Vec, pub has_pending_raise_op: bool, pub pending_move_op: Option<(usize, usize, usize)>, } @@ -66,7 +67,6 @@ pub struct State { pub is_paused: bool, pub invisible_borders: Rect, pub resize_delta: i32, - pub virtual_desktop_id: Option, pub new_window_behaviour: WindowContainerBehaviour, pub work_area_offset: Option, pub focus_follows_mouse: Option, @@ -87,7 +87,6 @@ impl From<&WindowManager> for State { invisible_borders: wm.invisible_borders, work_area_offset: wm.work_area_offset, resize_delta: wm.resize_delta, - virtual_desktop_id: wm.virtual_desktop_id, new_window_behaviour: wm.window_container_behaviour, focus_follows_mouse: wm.focus_follows_mouse.clone(), mouse_follows_focus: wm.mouse_follows_focus, @@ -148,8 +147,6 @@ impl WindowManager { let listener = UnixListener::bind(&socket)?; - let virtual_desktop_id = winvd::helpers::get_current_desktop_number().ok(); - Ok(Self { monitors: Ring::default(), incoming_events: incoming, @@ -161,13 +158,13 @@ impl WindowManager { right: 14, bottom: 7, }, + virtual_desktop_id: current_virtual_desktop()?, work_area_offset: None, window_container_behaviour: WindowContainerBehaviour::Create, resize_delta: 50, focus_follows_mouse: None, mouse_follows_focus: true, hotwatch: Hotwatch::new()?, - virtual_desktop_id, has_pending_raise_op: false, pending_move_op: None, }) @@ -468,24 +465,6 @@ impl WindowManager { Ok(()) } - #[tracing::instrument(skip(self))] - pub fn validate_virtual_desktop_id(&self) -> Result<()> { - let virtual_desktop_id = winvd::helpers::get_current_desktop_number().ok(); - if let (Some(id), Some(virtual_desktop_id)) = (virtual_desktop_id, self.virtual_desktop_id) - { - if id != virtual_desktop_id { - return Err(anyhow!( - "ignoring events and commands while not on virtual desktop {}", - virtual_desktop_id - )); - } - } else { - tracing::warn!("unable to look up virtual desktop id, skipping validation"); - } - - Ok(()) - } - #[tracing::instrument(skip(self))] pub fn manage_focused_window(&mut self) -> Result<()> { let hwnd = WindowsApi::foreground_window()?; diff --git a/komorebi/src/windows_api.rs b/komorebi/src/windows_api.rs index 85898e12..5b081679 100644 --- a/komorebi/src/windows_api.rs +++ b/komorebi/src/windows_api.rs @@ -31,6 +31,7 @@ use windows::Win32::Graphics::Gdi::HMONITOR; use windows::Win32::Graphics::Gdi::MONITORENUMPROC; use windows::Win32::Graphics::Gdi::MONITORINFO; use windows::Win32::Graphics::Gdi::MONITOR_DEFAULTTONEAREST; +use windows::Win32::System::RemoteDesktop::ProcessIdToSessionId; use windows::Win32::System::Threading::AttachThreadInput; use windows::Win32::System::Threading::GetCurrentProcessId; use windows::Win32::System::Threading::GetCurrentThreadId; @@ -382,6 +383,19 @@ impl WindowsApi { unsafe { GetCurrentProcessId() } } + pub fn process_id_to_session_id() -> Result { + let process_id = Self::current_process_id(); + let mut session_id = 0; + + unsafe { + if ProcessIdToSessionId(process_id, &mut session_id).as_bool() { + Ok(session_id) + } else { + Err(anyhow!("could not determine current session id")) + } + } + } + pub fn attach_thread_input(thread_id: u32, target_thread_id: u32, attach: bool) -> Result<()> { unsafe { AttachThreadInput(thread_id, target_thread_id, attach) } .ok()