From 977611482a15138333171565ff9bddb6abd1bcc3 Mon Sep 17 00:00:00 2001 From: daladim Date: Tue, 9 Aug 2022 16:11:56 +0200 Subject: [PATCH] Using a bitfield for TraceFlags --- src/native/etw_types.rs | 6 +++--- src/provider.rs | 23 ++++++++++++++--------- src/provider/trace_flags.rs | 16 ++++++++++++++++ 3 files changed, 33 insertions(+), 12 deletions(-) create mode 100644 src/provider/trace_flags.rs diff --git a/src/native/etw_types.rs b/src/native/etw_types.rs index 5a0a50a..3361388 100644 --- a/src/native/etw_types.rs +++ b/src/native/etw_types.rs @@ -7,7 +7,7 @@ //! In most cases a user of the crate won't have to deal with this and can directly obtain the data //! needed by using the functions exposed by the modules at the crate level use crate::native::tdh_types::Property; -use crate::provider::Provider; +use crate::provider::{Provider, TraceFlags}; use crate::trace::{TraceData, TraceProperties, TraceTrait}; use crate::utils; use std::fmt::Formatter; @@ -239,12 +239,12 @@ impl std::ops::DerefMut for EventTraceLogfile { pub struct EnableTraceParameters(Etw::ENABLE_TRACE_PARAMETERS); impl EnableTraceParameters { - pub fn create(guid: GUID, trace_flags: u32) -> Self { + pub fn create(guid: GUID, trace_flags: TraceFlags) -> Self { let mut params = EnableTraceParameters::default(); params.0.ControlFlags = 0; params.0.Version = Etw::ENABLE_TRACE_PARAMETERS_VERSION_2; params.0.SourceId = guid; - params.0.EnableProperty = trace_flags; + params.0.EnableProperty = trace_flags.bits(); // TODO: Add Filters option params.0.EnableFilterDesc = std::ptr::null_mut(); diff --git a/src/provider.rs b/src/provider.rs index 34f6792..61a6ab9 100644 --- a/src/provider.rs +++ b/src/provider.rs @@ -8,6 +8,9 @@ use crate::schema; use std::sync::{Arc, RwLock}; use windows::core::GUID; +mod trace_flags; +pub use trace_flags::TraceFlags; + /// Provider module errors #[derive(Debug)] pub enum ProviderError { @@ -278,7 +281,9 @@ pub struct Provider { /// Provider level flag pub level: u8, /// Provider trace flags - pub trace_flags: u32, + /// + /// Used as `EnableParameters.EnableProperty` when starting the trace (using [EnableTraceEx2](https://docs.microsoft.com/en-us/windows/win32/api/evntrace/nf-evntrace-enabletraceex2)) + pub trace_flags: TraceFlags, /// Provider kernel flags, only apply to KernelProvider pub flags: u32, // Only applies to KernelProviders // perfinfo @@ -302,7 +307,7 @@ impl Provider { any: 0, all: 0, level: 5, - trace_flags: 0, + trace_flags: TraceFlags::empty(), flags: 0, callbacks: Arc::new(RwLock::new(Vec::new())), } @@ -318,7 +323,7 @@ impl Provider { any: 0, all: 0, level: 5, - trace_flags: 0, + trace_flags: TraceFlags::empty(), flags: kernel_provider.flags, callbacks: Arc::new(RwLock::new(Vec::new())), } @@ -417,11 +422,11 @@ impl Provider { /// /// # Example /// ``` - /// # use ferrisetw::provider::Provider; - /// let my_provider = Provider::new().trace_flags(0x1); + /// # use ferrisetw::provider::{Provider, TraceFlags}; + /// let my_provider = Provider::new().trace_flags(TraceFlags::EVENT_ENABLE_PROPERTY_SID); /// ``` - pub fn trace_flags(mut self, trace_flag: u32) -> Self { - self.trace_flags = trace_flag; + pub fn trace_flags(mut self, trace_flags: TraceFlags) -> Self { + self.trace_flags = trace_flags; self } @@ -538,8 +543,8 @@ mod test { #[test] fn test_set_trace_flags() { - let prov = Provider::new().trace_flags(100); - assert_eq!(100, prov.trace_flags); + let prov = Provider::new().trace_flags(TraceFlags::all()); + assert_eq!(prov.trace_flags, TraceFlags::all()); } #[test] diff --git a/src/provider/trace_flags.rs b/src/provider/trace_flags.rs new file mode 100644 index 0000000..3cd1d21 --- /dev/null +++ b/src/provider/trace_flags.rs @@ -0,0 +1,16 @@ +use bitflags::bitflags; + +use windows::Win32::System::Diagnostics::Etw; + +bitflags! { + pub struct TraceFlags: u32 { + const EVENT_ENABLE_PROPERTY_IGNORE_KEYWORD_0 = Etw::EVENT_ENABLE_PROPERTY_IGNORE_KEYWORD_0; + const EVENT_ENABLE_PROPERTY_PROVIDER_GROUP = Etw::EVENT_ENABLE_PROPERTY_PROVIDER_GROUP; + const EVENT_ENABLE_PROPERTY_PROCESS_START_KEY = Etw::EVENT_ENABLE_PROPERTY_PROCESS_START_KEY; + const EVENT_ENABLE_PROPERTY_EVENT_KEY = Etw::EVENT_ENABLE_PROPERTY_EVENT_KEY; + const EVENT_ENABLE_PROPERTY_EXCLUDE_INPRIVATE = Etw::EVENT_ENABLE_PROPERTY_EXCLUDE_INPRIVATE; + const EVENT_ENABLE_PROPERTY_SID = Etw::EVENT_ENABLE_PROPERTY_SID; + const EVENT_ENABLE_PROPERTY_TS_ID = Etw::EVENT_ENABLE_PROPERTY_TS_ID; + const EVENT_ENABLE_PROPERTY_STACK_TRACE = Etw::EVENT_ENABLE_PROPERTY_STACK_TRACE; + } +}