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

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"

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<Vec<u8>> {
+    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();
 

komorebi/src/process_command.rs

@@ -68,9 +68,14 @@ pub fn listen_for_commands(wm: Arc<Mutex<WindowManager>>) {
 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 {

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

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<FocusFollowsMouseImplementation>,
     pub mouse_follows_focus: bool,
     pub hotwatch: Hotwatch,
-    pub virtual_desktop_id: Option<usize>,
+    pub virtual_desktop_id: Vec<u8>,
     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<usize>,
     pub new_window_behaviour: WindowContainerBehaviour,
     pub work_area_offset: Option<Rect>,
     pub focus_follows_mouse: Option<FocusFollowsMouseImplementation>,
@@ -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()?;

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<u32> {
+        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()