Skip to content

Commit

Permalink
Check if a winrt functionality is available before using it
Browse files Browse the repository at this point in the history
  • Loading branch information
@GuillaumeGomez
GuillaumeGomez committed Apr 7, 2020
1 parent 202000e commit bf16c48
Show file tree
Hide file tree
Showing 3 changed files with 94 additions and 42 deletions.
2 changes: 1 addition & 1 deletion Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"]}

Expand Down
56 changes: 24 additions & 32 deletions src/windows/process.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,26 @@ unsafe fn get_exe(process_handler: HANDLE, h_mod: *mut c_void) -> PathBuf {
PathBuf::from(String::from_utf16_lossy(&exe_buf[..pos]))
}

fn get_disk_for_process_id(pid: u32) -> Option<DiskInfo> {
use super::system::{can_be_registered, RegistrationStatus};

match can_be_registered::<ProcessDiagnosticInfo>() {
RegistrationStatus::Ok => {
ProcessDiagnosticInfo::try_get_for_process_id(pid as u32)
.ok()
.and_then(|x| x)
.map(|x| DiskInfo {
diag_info: PtrWrapper(x),
old_disk_bytes: RefCell::new(DiskBytes::default()),
})
}
_x => {
sysinfo_debug!("Cannot register AppDiagnosticInfo: {:?}", _x);
None
}
}
}

impl Process {
#[allow(clippy::uninit_assumed_init)]
pub(crate) fn new_from_pid(pid: Pid) -> Option<Process> {
Expand Down Expand Up @@ -236,14 +256,7 @@ impl Process {
old_user_cpu: 0,
start_time: unsafe { get_start_time(handle) },
updated: true,
disk_info:
ProcessDiagnosticInfo::try_get_for_process_id(pid as u32)
.ok()
.and_then(|x| x)
.map(|x| DiskInfo {
diag_info: PtrWrapper(x),
old_disk_bytes: RefCell::new(DiskBytes::default()),
})
disk_info: get_disk_for_process_id(pid as u32),
}
} else {
Process {
Expand All @@ -265,14 +278,7 @@ impl Process {
old_user_cpu: 0,
start_time: 0,
updated: true,
disk_info:
ProcessDiagnosticInfo::try_get_for_process_id(pid as u32)
.ok()
.and_then(|x| x)
.map(|x| DiskInfo {
diag_info: PtrWrapper(x),
old_disk_bytes: RefCell::new(DiskBytes::default()),
})
disk_info: get_disk_for_process_id(pid as u32),
}
}
}
Expand Down Expand Up @@ -310,14 +316,7 @@ impl Process {
old_user_cpu: 0,
start_time: get_start_time(process_handler),
updated: true,
disk_info:
ProcessDiagnosticInfo::try_get_for_process_id(pid as u32)
.ok()
.and_then(|x| x)
.map(|x| DiskInfo {
diag_info: PtrWrapper(x),
old_disk_bytes: RefCell::new(DiskBytes::default()),
})
disk_info: get_disk_for_process_id(pid as u32),
}
}
}
Expand Down Expand Up @@ -352,14 +351,7 @@ impl ProcessExt for Process {
old_user_cpu: 0,
start_time: 0,
updated: true,
disk_info:
ProcessDiagnosticInfo::try_get_for_process_id(pid as u32)
.ok()
.and_then(|x| x)
.map(|x| DiskInfo {
diag_info: PtrWrapper(x),
old_disk_bytes: RefCell::new(DiskBytes::default()),
})
disk_info: get_disk_for_process_id(pid as u32),
}
}
}
Expand Down
78 changes: 69 additions & 9 deletions src/windows/system.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ use winapi::um::sysinfoapi::{GetTickCount64, GlobalMemoryStatusEx, MEMORYSTATUSE
use winapi::um::winnt::HANDLE;
use rayon::prelude::*;
use winrt::windows::system::{AppDiagnosticInfo, DiagnosticAccessStatus};
use winrt::{RtAsyncOperation};
use winrt::{RtAsyncOperation, RtActivatable};

/// Struct containing the system's information.
pub struct System {
Expand Down Expand Up @@ -74,19 +74,79 @@ unsafe fn boot_time() -> u64 {
}
}
}

// To be removed once winrt PR is merged!
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) enum RegistrationStatus {
Ok,
NotRegistered,
Error(u32),
}

// To be removed once winrt PR is merged!
pub(crate) fn can_be_registered<T: RtActivatable<T>>() -> 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: <T as RtType>::Abi = unsafe {
<T as RtType>::uninitialized()
};
let class_id = unsafe { HStringReference::from_utf16_unchecked(T::name()) };
let mut hr = unsafe {
RoGetActivationFactory(
class_id.get(),
<T as 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(),
<T as 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::<AppDiagnosticInfo>() {
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 {
Expand Down

0 comments on commit bf16c48

Please sign in to comment.