Check if a winrt functionality is available before using it

GuillaumeGomez · Apr 6, 2020 · 299fba6 · 299fba6
1 parent 202000e
commit 299fba6
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 9 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,7 +19,7 @@ doc-comment = "0.3"
 once_cell = "1.0"
 
 [target.'cfg(windows)'.dependencies]
-winapi = { version = "0.3", features = ["fileapi", "handleapi", "ifdef", "ioapiset", "minwindef", "pdh", "psapi", "synchapi", "sysinfoapi", "winbase", "winerror", "winioctl", "winnt", "oleauto", "wbemcli", "rpcdce", "combaseapi", "objidl", "powerbase", "netioapi", "lmcons", "lmaccess", "lmapibuf", "memoryapi", "shellapi"] }
+winapi = { version = "0.3", features = ["fileapi", "handleapi", "ifdef", "ioapiset", "minwindef", "pdh", "psapi", "synchapi", "sysinfoapi", "winbase", "winerror", "winioctl", "winnt", "oleauto", "wbemcli", "rpcdce", "combaseapi", "objidl", "powerbase", "netioapi", "lmcons", "lmaccess", "lmapibuf", "memoryapi", "shellapi", "roapi"] }
 ntapi = "0.3"
 winrt = { version = "0.6.0", features = ["windows-system"]}
 

diff --git a/src/windows/system.rs b/src/windows/system.rs
@@ -74,19 +74,59 @@ unsafe fn boot_time() -> u64 {
         }
     }
 }
+
+// To be removed once winrt PR is merged!
+#[derive(Clone, Copy, Debug, PartialEq)]
+enum RegistrationStatus {
+    Ok,
+    NotRegistered,
+    Error(u32),
+}
+
+// To be removed once winrt PR is merged!
+fn can_be_registered() -> RegistrationStatus {
+    use std::ptr;
+    use winapi::shared::ntdef::VOID;
+    use winapi::shared::winerror::{S_OK, CO_E_NOTINITIALIZED, REGDB_E_CLASSNOTREG};
+    use winapi::um::combaseapi::CoIncrementMTAUsage;
+    use winapi::winrt::roapi::RoGetActivationFactory;
+    use winrt::{ComIid, HStringReference, RtType};
+
+    let mut res: RtType::Abi = unsafe { RtType::uninitialized() };
+    let class_id = unsafe { HStringReference::from_utf16_unchecked(&[87,105,110,100,111,119,115,46,83,121,115,116,101,109,46,65,112,112,68,105,97,103,110,111,115,116,105,99,73,110,102,111,0]) };
+    let mut hr = unsafe { RoGetActivationFactory(class_id.get(), ComIid::iid().as_ref(), &mut res as *mut *mut _ as *mut *mut VOID) };
+    if hr == CO_E_NOTINITIALIZED {
+        let mut cookie = ptr::null_mut();
+        unsafe { CoIncrementMTAUsage(&mut cookie); }
+        hr = unsafe { RoGetActivationFactory(class_id.get(), ComIid::iid().as_ref(), &mut res as *mut *mut _ as *mut *mut VOID) };
+    }
+    if hr == S_OK {
+        RegistrationStatus::Ok
+    } else if hr == REGDB_E_CLASSNOTREG {
+        RegistrationStatus::NotRegistered
+    } else {
+        RegistrationStatus::Error(hr as u32)
+    }
+}
+
 impl SystemExt for System {
     #[allow(non_snake_case)]
     fn new_with_specifics(refreshes: RefreshKind) -> System {
-        match AppDiagnosticInfo::request_access_async() {
-            Ok(x) => {
-                match x.blocking_get() {
-                    Ok(DiagnosticAccessStatus::Unspecified) => sysinfo_debug!("app diagnostic access refused: Unspecified"),
-                    Ok(DiagnosticAccessStatus::Denied) => sysinfo_debug!("app diagnostic access refused: Denied"),
-                    Ok(_) => {}
-                    Err(_e) => sysinfo_debug!("failed to get access request status: {:?}", _e),
+        match can_be_registered() {
+            RegistrationStatus::Ok => {
+                match AppDiagnosticInfo::request_access_async() {
+                    Ok(x) => {
+                        match x.blocking_get() {
+                            Ok(DiagnosticAccessStatus::Unspecified) => sysinfo_debug!("app diagnostic access refused: Unspecified"),
+                            Ok(DiagnosticAccessStatus::Denied) => sysinfo_debug!("app diagnostic access refused: Denied"),
+                            Ok(_) => {}
+                            Err(_e) => sysinfo_debug!("failed to get access request status: {:?}", _e),
+                        }
+                    }
+                    Err(_e) => sysinfo_debug!("failed to request app diagnostic access: {:?}", _e),
                 }
             }
-            Err(_e) => sysinfo_debug!("failed to request app diagnostic access: {:?}", _e),
+            x => sysinfo_debug!("Cannot register AppDiagnosticInfo: {:?}", x),
         }
         let (processors, vendor_id, brand) = init_processors();
         let mut s = System {