diff --git a/Cargo.toml b/Cargo.toml index 1f051e1..f89f153 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,7 +19,7 @@ default-target = "x86_64-pc-windows-msvc" [features] test = ["lazy_static", "async-std"] -async-query = ["com-impl", "wio", "async-channel", "futures"] +async-query = ["async-channel", "futures", "com"] [target.'cfg(target_os = "windows")'.dependencies] winapi = { version = "0.3.9", features = ["objbase", "wbemcli", "objidlbase", "oaidl", "oleauto", "errhandlingapi", "wtypes", "wtypesbase"] } @@ -29,12 +29,10 @@ serde = { version = "1.0", features = ["derive"] } chrono = { version = "0.4", features = ["serde"] } lazy_static = { version = "1.4.0", optional = true } thiserror = "^1" -# dependencies for async queries -com-impl = { git = "https://github.com/Connicpu/com-impl", optional = true} -wio = {version = "0.2.2", optional = true} async-channel = {version = "1.5.1", optional = true} async-std = { version = "1.8.0", features = ["attributes"], optional = true } -futures = {version = "0.3", optional = true} +futures = { version = "0.3", optional = true} +com = { version = "0.4", features = ["production"], optional = true } [dev-dependencies] lazy_static = "1.4.0" diff --git a/azure-pipelines.yml b/azure-pipelines.yml index 593453c..9fecc54 100644 --- a/azure-pipelines.yml +++ b/azure-pipelines.yml @@ -29,7 +29,7 @@ steps: displayName: Install rust (windows) condition: eq( variables['Agent.OS'], 'Windows_NT' ) - - script: cargo build --release + - script: cargo build --all-features displayName: Cargo build - script: cargo test --all-features diff --git a/src/async_query.rs b/src/async_query.rs index 91cdb44..4c5b5d3 100644 --- a/src/async_query.rs +++ b/src/async_query.rs @@ -16,17 +16,17 @@ use crate::{ utils::check_hres, WMIResult, }; -use crate::query_sink::QuerySink; +use crate::query_sink::{QuerySink, IWbemObjectSink}; use std::{collections::HashMap, ptr}; +use com::production::ClassAllocation; +use com::AbiTransferable; use winapi::{ um::{ - wbemcli::IWbemObjectSink, wbemcli::{ WBEM_FLAG_BIDIRECTIONAL, }, }, }; -use wio::com::ComPtr; use serde::de; use futures::stream::{Stream, TryStreamExt, StreamExt}; @@ -48,20 +48,22 @@ impl WMIConnection { let query = BStr::from_str(query.as_ref())?; let (tx, rx) = async_channel::unbounded(); - // ComPtr keeps an internal RefCount with initial value = 1 - let p_sink: ComPtr = QuerySink::new(tx); - + // The internal RefCount has initial value = 1. + let p_sink: ClassAllocation = QuerySink::allocate(Some(tx)); + let p_sink_handel = p_sink.query_interface::().unwrap(); + unsafe { - // As p_sink.RefCount = 1 before this call, + // As p_sink's RefCount = 1 before this call, // p_sink won't be dropped at the end of ExecQueryAsync check_hres((*self.svc()).ExecQueryAsync( query_language.as_bstr(), query.as_bstr(), WBEM_FLAG_BIDIRECTIONAL as i32, ptr::null_mut(), - p_sink.as_raw(), + p_sink_handel.get_abi().as_ptr() as * mut _, ))?; } + Ok(rx) } @@ -158,7 +160,7 @@ mod tests { use futures::stream::StreamExt; #[async_std::test] - async fn _async_it_works_async() { + async fn async_it_works_async() { let wmi_con = wmi_con(); let result = wmi_con @@ -171,7 +173,7 @@ mod tests { } #[async_std::test] - async fn _async_it_handles_invalid_query() { + async fn async_it_handles_invalid_query() { let wmi_con = wmi_con(); let result = wmi_con @@ -180,11 +182,11 @@ mod tests { .collect::>() .await; - assert_eq!(result.len(), 0); + assert_eq!(result.len(), 0); } #[async_std::test] - async fn _async_it_provides_raw_query_result() { + async fn async_it_provides_raw_query_result() { let wmi_con = wmi_con(); let results: Vec> = diff --git a/src/de/meta.rs b/src/de/meta.rs index 8e69912..b0330da 100644 --- a/src/de/meta.rs +++ b/src/de/meta.rs @@ -2,7 +2,7 @@ use serde::de::{self, Deserialize, Deserializer, Visitor}; use serde::forward_to_deserialize_any; /// Return the fields of a struct. -/// Taken directly from https://github.com/serde-rs/serde/issues/1110 +/// Taken directly from /// pub fn struct_name_and_fields<'de, T>( ) -> Result<(&'static str, &'static [&'static str]), serde::de::value::Error> diff --git a/src/lib.rs b/src/lib.rs index 3d29cde..04ad84a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -86,7 +86,7 @@ //! # Internals //! //! [`WMIConnection`](WMIConnection) is used to create and execute a WMI query, returning -//! [`IWbemClassWrapper`](query::IWbemClassWrapper) which is a wrapper for a WMI object pointer. +//! [`IWbemClassWrapper`](result_enumerator::IWbemClassWrapper) which is a wrapper for a WMI object pointer. //! //! Then, [`from_wbem_class_obj`](de::wbem_class_de::from_wbem_class_obj) is used to create a Rust struct with the equivalent data. //! diff --git a/src/query.rs b/src/query.rs index c462b05..e95237c 100644 --- a/src/query.rs +++ b/src/query.rs @@ -339,7 +339,9 @@ impl WMIConnection { ))?; } - let pcls_wrapper = unsafe { IWbemClassWrapper::new(NonNull::new(pcls_obj)) }; + let pcls_ptr = NonNull::new(pcls_obj).ok_or(WMIError::NullPointerResult)?; + + let pcls_wrapper = unsafe { IWbemClassWrapper::new(pcls_ptr) }; Ok(pcls_wrapper) } @@ -435,7 +437,7 @@ impl WMIConnection { /// Query all the associators of type T of the given object. /// The `object_path` argument can be provided by querying an object wih it's `__Path` property. /// `AssocClass` must be have the name as the conneting association class between the original object and the results. - /// See https://docs.microsoft.com/en-us/windows/desktop/cimwin32prov/win32-diskdrivetodiskpartition for example. + /// See for example. /// /// ```edition2018 /// # fn main() -> Result<(), wmi::WMIError> { diff --git a/src/query_sink.rs b/src/query_sink.rs index 631c93e..ec36dec 100644 --- a/src/query_sink.rs +++ b/src/query_sink.rs @@ -1,6 +1,6 @@ use winapi::{ um::wbemcli::{ - {IWbemClassObject,IWbemObjectSink, IWbemObjectSinkVtbl}, + {IWbemClassObject}, WBEM_S_NO_ERROR, WBEM_STATUS_COMPLETE, }, @@ -13,100 +13,105 @@ use winapi::{ c_long, }, }; -use com_impl::{ComImpl, VTable, Refcount}; use log::{trace, warn}; use std::ptr::NonNull; -use wio::com::ComPtr; use async_channel::Sender; use crate::result_enumerator::IWbemClassWrapper; use crate::WMIError; -/// Implementation for [IWbemObjectSink](https://docs.microsoft.com/en-us/windows/win32/api/wbemcli/nn-wbemcli-iwbemobjectsink). -/// This [Sink](https://en.wikipedia.org/wiki/Sink_(computing)) -/// receives asynchronously the result of the query, -/// through Indicate calls. When finished,the SetStatus method -/// is called. -/// # -#[repr(C)] -#[derive(ComImpl)] -#[interfaces(IWbemObjectSink)] -pub struct QuerySink { - vtbl: VTable, - refcount: Refcount, - sender: Sender>, -} - -impl QuerySink { - pub fn new(sender: Sender>) -> ComPtr { - let ptr = QuerySink::create_raw(sender); - let ptr = ptr as *mut IWbemObjectSink; - unsafe { ComPtr::from_raw(ptr) } +com::interfaces! { + #[uuid("7C857801-7381-11CF-884D-00AA004B2E24")] + pub unsafe interface IWbemObjectSink: com::interfaces::IUnknown { + unsafe fn indicate( + &self, + lObjectCount: c_long, + apObjArray: *mut *mut IWbemClassObject + ) -> HRESULT; + + unsafe fn set_status( + &self, + lFlags: c_long, + _hResult: HRESULT, + _strParam: BSTR, + _pObjParam: *mut IWbemClassObject + ) -> HRESULT; } } -// QueryInterface, AddRef and Release methods are provided by com_impl -#[com_impl::com_impl] -unsafe impl IWbemObjectSink for QuerySink { - unsafe fn indicate( - &self, - lObjectCount: c_long, - apObjArray: *mut *mut IWbemClassObject - ) -> HRESULT { - trace!("Indicate call with {} objects", lObjectCount); - // Case of an incorrect or too restrictive query - if lObjectCount <= 0 { - return WBEM_S_NO_ERROR as i32; - } - - let lObjectCount = lObjectCount as usize; - let tx = self.sender.clone(); - unsafe { - // The array memory of apObjArray is read-only - // and is owned by the caller of the Indicate method. - // IWbemClassWrapper::clone calls AddRef on each element - // of apObjArray to make sure that they are not released, - // according to COM rules. - // https://docs.microsoft.com/en-us/windows/win32/api/wbemcli/nf-wbemcli-iwbemobjectsink-indicate - // For error codes, see https://docs.microsoft.com/en-us/windows/win32/learnwin32/error-handling-in-com - for i in 0..lObjectCount { - if let Some(p_el) = NonNull::new(*apObjArray.offset(i as isize)) { - let wbemClassObject = IWbemClassWrapper::clone(p_el); - - if let Err(e) = tx.try_send(Ok(wbemClassObject)) { - warn!("Error while sending object through async_channel: {}", e); - return E_FAIL; - } - } else { - if let Err(e) = tx.try_send(Err(WMIError::NullPointerResult)) { - warn!("Error while sending error code through async_channel: {}", e); +com::class! { + // Option is required because `Default` is required by the `class!` macro. + pub class QuerySink: IWbemObjectSink { + sender: Option>>, + } + + /// Implementation for [IWbemObjectSink](https://docs.microsoft.com/en-us/windows/win32/api/wbemcli/nn-wbemcli-iwbemobjectsink). + /// This [Sink](https://en.wikipedia.org/wiki/Sink_(computing)) + /// receives asynchronously the result of the query, through Indicate calls. + /// When finished,the SetStatus method is called. + /// # + impl IWbemObjectSink for QuerySink { + unsafe fn indicate( + &self, + lObjectCount: c_long, + apObjArray: *mut *mut IWbemClassObject + ) -> HRESULT { + trace!("Indicate call with {} objects", lObjectCount); + // Case of an incorrect or too restrictive query + if lObjectCount <= 0 { + return WBEM_S_NO_ERROR as i32; + } + + let lObjectCount = lObjectCount as usize; + let tx = self.sender.as_ref().unwrap().clone(); + + unsafe { + // The array memory of apObjArray is read-only + // and is owned by the caller of the Indicate method. + // IWbemClassWrapper::clone calls AddRef on each element + // of apObjArray to make sure that they are not released, + // according to COM rules. + // https://docs.microsoft.com/en-us/windows/win32/api/wbemcli/nf-wbemcli-iwbemobjectsink-indicate + // For error codes, see https://docs.microsoft.com/en-us/windows/win32/learnwin32/error-handling-in-com + for i in 0..lObjectCount { + if let Some(p_el) = NonNull::new(*apObjArray.offset(i as isize)) { + let wbemClassObject = IWbemClassWrapper::clone(p_el); + + if let Err(e) = tx.try_send(Ok(wbemClassObject)) { + warn!("Error while sending object through async_channel: {}", e); + return E_FAIL; + } + } else { + if let Err(e) = tx.try_send(Err(WMIError::NullPointerResult)) { + warn!("Error while sending error code through async_channel: {}", e); + } + return E_POINTER; } - return E_POINTER; + } - } + + WBEM_S_NO_ERROR as i32 } - - WBEM_S_NO_ERROR as i32 - } - - unsafe fn set_status( - &self, - lFlags: c_long, - _hResult: HRESULT, - _strParam: BSTR, - _pObjParam: *mut IWbemClassObject - ) -> HRESULT { - // SetStatus is called only once as flag=WBEM_FLAG_BIDIRECTIONAL in ExecQueryAsync - // https://docs.microsoft.com/en-us/windows/win32/api/wbemcli/nf-wbemcli-iwbemobjectsink-setstatus - // If you do not specify WBEM_FLAG_SEND_STATUS when calling your provider or service method, - // you are guaranteed to receive one and only one call to SetStatus - - if lFlags == WBEM_STATUS_COMPLETE as i32 { - trace!("End of async result, closing transmitter"); - self.sender.close(); + + unsafe fn set_status( + &self, + lFlags: c_long, + _hResult: HRESULT, + _strParam: BSTR, + _pObjParam: *mut IWbemClassObject + ) -> HRESULT { + // SetStatus is called only once as flag=WBEM_FLAG_BIDIRECTIONAL in ExecQueryAsync + // https://docs.microsoft.com/en-us/windows/win32/api/wbemcli/nf-wbemcli-iwbemobjectsink-setstatus + // If you do not specify WBEM_FLAG_SEND_STATUS when calling your provider or service method, + // you are guaranteed to receive one and only one call to SetStatus + + if lFlags == WBEM_STATUS_COMPLETE as i32 { + trace!("End of async result, closing transmitter"); + self.sender.as_ref().unwrap().close(); + } + WBEM_S_NO_ERROR as i32 } - WBEM_S_NO_ERROR as i32 } } @@ -123,12 +128,13 @@ mod tests { async fn _async_it_should_use_async_channel_to_send_result() { let con = wmi_con(); let (tx, rx) = async_channel::unbounded(); - let p_sink: ComPtr = QuerySink::new(tx); + let sink = QuerySink::allocate(Some(tx)); + let p_sink = sink.query_interface::().unwrap(); let raw_os = con.get_raw_by_path(r#"\\.\root\cimv2:Win32_OperatingSystem=@"#).unwrap(); let raw_os2 = con.get_raw_by_path(r#"\\.\root\cimv2:Win32_OperatingSystem=@"#).unwrap(); - let ptr: *mut IWbemClassObject = raw_os.inner.unwrap().as_ptr(); - let ptr2: *mut IWbemClassObject = raw_os2.inner.unwrap().as_ptr(); + let ptr: *mut IWbemClassObject = raw_os.inner.as_ptr(); + let ptr2: *mut IWbemClassObject = raw_os2.inner.as_ptr(); let mut arr = vec![ptr, ptr2]; @@ -143,7 +149,7 @@ mod tests { assert_eq!(refcount, 1); } - unsafe {p_sink.Indicate(arr.len() as i32, arr.as_mut_ptr());} + unsafe {p_sink.indicate(arr.len() as i32, arr.as_mut_ptr());} // tests on ref count after Indicate call unsafe { let test_ptr = &ptr; @@ -178,11 +184,12 @@ mod tests { #[test] fn _async_it_should_close_async_channel_after_set_status_call() { let (tx, rx) = async_channel::unbounded(); - let p_sink: ComPtr = QuerySink::new(tx); + let sink = QuerySink::allocate(Some(tx)); + let p_sink = sink.query_interface::().unwrap(); assert!(!rx.is_closed()); - unsafe {p_sink.SetStatus(WBEM_STATUS_COMPLETE as i32, 0, NULL as BSTR, NULL as *mut IWbemClassObject);} + unsafe {p_sink.set_status(WBEM_STATUS_COMPLETE as i32, 0, NULL as BSTR, NULL as *mut IWbemClassObject);} assert!(rx.is_closed()); } @@ -190,12 +197,13 @@ mod tests { #[async_std::test] async fn _async_it_should_return_e_pointer_after_indicate_call_with_null_pointer() { let (tx, rx) = async_channel::unbounded(); - let p_sink: ComPtr = QuerySink::new(tx); + let sink = QuerySink::allocate(Some(tx)); + let p_sink = sink.query_interface::().unwrap(); let mut arr = vec![NULL as *mut IWbemClassObject]; let result; - unsafe { result = p_sink.Indicate(1, arr.as_mut_ptr()) } + unsafe { result = p_sink.indicate(1, arr.as_mut_ptr()) } match rx.recv().await.unwrap() { Err(WMIError::NullPointerResult) => assert!(true), diff --git a/src/result_enumerator.rs b/src/result_enumerator.rs index 90b55dc..fd04324 100644 --- a/src/result_enumerator.rs +++ b/src/result_enumerator.rs @@ -27,11 +27,11 @@ use winapi::{ /// #[derive(Debug)] pub struct IWbemClassWrapper { - pub inner: Option>, + pub inner: NonNull, } impl IWbemClassWrapper { - pub unsafe fn new(ptr: Option>) -> Self { + pub unsafe fn new(ptr: NonNull) -> Self { Self { inner: ptr } } @@ -45,7 +45,7 @@ impl IWbemClassWrapper { pub unsafe fn clone(ptr: NonNull) -> Self { let refcount = ptr.as_ref().AddRef(); trace!("Reference count: {}", refcount); - Self::new(Some(ptr)) + Self::new(ptr) } /// Return the names of all the properties of the given object. @@ -54,7 +54,7 @@ impl IWbemClassWrapper { // This will store the properties names from the GetNames call. let mut p_names = NULL as *mut SAFEARRAY; - let ptr = self.inner.unwrap().as_ptr(); + let ptr = self.inner.as_ptr(); unsafe { check_hres((*ptr).GetNames( @@ -79,18 +79,18 @@ impl IWbemClassWrapper { let mut vt_prop: VARIANT = unsafe { mem::zeroed() }; unsafe { - (*self.inner.unwrap().as_ptr()).Get( + check_hres((*self.inner.as_ptr()).Get( name_prop.as_lpcwstr(), 0, &mut vt_prop, ptr::null_mut(), ptr::null_mut(), - ); + ))?; let property_value = Variant::from_variant(vt_prop)?; - VariantClear(&mut vt_prop); + check_hres(VariantClear(&mut vt_prop))?; Ok(property_value) } @@ -114,12 +114,10 @@ impl IWbemClassWrapper { impl Drop for IWbemClassWrapper { fn drop(&mut self) { - if let Some(pcls_obj) = self.inner { - let ptr = pcls_obj.as_ptr(); + let ptr = self.inner.as_ptr(); - unsafe { - (*ptr).Release(); - } + unsafe { + (*ptr).Release(); } } } @@ -180,8 +178,11 @@ impl<'a> Iterator for QueryResultEnumerator<'a> { pcls_obj ); - let pcls_wrapper = unsafe { IWbemClassWrapper::new(NonNull::new(pcls_obj)) }; + let pcls_ptr = NonNull::new(pcls_obj).ok_or(WMIError::NullPointerResult); - Some(Ok(pcls_wrapper)) + match pcls_ptr { + Err(e) => Some(Err(e)), + Ok(pcls_ptr) => Some(Ok( unsafe { IWbemClassWrapper::new(pcls_ptr) })), + } } }