From 6f7e87799b71335e0518199677d11f6f6c6c41c9 Mon Sep 17 00:00:00 2001
From: LGUG2Z <jadeiqbal@fastmail.com>
Date: Thu, 19 Aug 2021 14:51:35 -0700
Subject: [PATCH] fix(wm): handle winvd errors gracefully

Of course, the crate built to interact with an undocumented COM API is
not the best candidate for unwrap and expect calls...

fix #15
---
 komorebi/src/process_command.rs | 19 ++++++++++---------
 komorebi/src/process_event.rs   | 21 +++++++++++----------
 komorebi/src/window_manager.rs  |  5 ++---
 3 files changed, 23 insertions(+), 22 deletions(-)

diff --git a/komorebi/src/process_command.rs b/komorebi/src/process_command.rs
index c6eafa02..55e6350e 100644
--- a/komorebi/src/process_command.rs
+++ b/komorebi/src/process_command.rs
@@ -52,16 +52,17 @@ pub fn listen_for_commands(wm: Arc<Mutex<WindowManager>>) {
 impl WindowManager {
     #[tracing::instrument(skip(self))]
     pub fn process_command(&mut self, message: SocketMessage) -> Result<()> {
-        let virtual_desktop_id = winvd::helpers::get_current_desktop_number()
-            .expect("could not determine the current virtual desktop number");
-
-        if virtual_desktop_id != self.virtual_desktop_id {
-            tracing::warn!(
-                "ignoring commands while not on virtual desktop {}",
-                self.virtual_desktop_id
-            );
+        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 {
+                tracing::warn!(
+                    "ignoring events while not on virtual desktop {:?}",
+                    virtual_desktop_id
+                );
 
-            return Ok(());
+                return Ok(());
+            }
         }
 
         match message {
diff --git a/komorebi/src/process_event.rs b/komorebi/src/process_event.rs
index db97c080..de54ac66 100644
--- a/komorebi/src/process_event.rs
+++ b/komorebi/src/process_event.rs
@@ -48,16 +48,17 @@ impl WindowManager {
             return Ok(());
         }
 
-        let virtual_desktop_id = winvd::helpers::get_current_desktop_number()
-            .expect("could not determine the current virtual desktop number");
-
-        if virtual_desktop_id != self.virtual_desktop_id {
-            tracing::warn!(
-                "ignoring events while not on virtual desktop {}",
-                self.virtual_desktop_id
-            );
-
-            return Ok(());
+        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 {
+                tracing::warn!(
+                    "ignoring events while not on virtual desktop {:?}",
+                    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 8530aafc..4d9e1f7e 100644
--- a/komorebi/src/window_manager.rs
+++ b/komorebi/src/window_manager.rs
@@ -45,7 +45,7 @@ pub struct WindowManager {
     pub command_listener: UnixListener,
     pub is_paused: bool,
     pub hotwatch: Hotwatch,
-    pub virtual_desktop_id: usize,
+    pub virtual_desktop_id: Option<usize>,
 }
 
 #[derive(Debug, Serialize)]
@@ -123,8 +123,7 @@ impl WindowManager {
 
         let listener = UnixListener::bind(&socket)?;
 
-        let virtual_desktop_id = winvd::helpers::get_current_desktop_number()
-            .expect("could not determine the current virtual desktop number");
+        let virtual_desktop_id = winvd::helpers::get_current_desktop_number().ok();
 
         Ok(Self {
             monitors: Ring::default(),