Dropping CallbackData at the right time

Wrapping it into an Arc ensures we're not dropping it when the trace is stopped, but we're waiting for the potential callbacks to terminate first

This fixes "Race 2" in n4r1b#45 (n4r1b#45)

src/native/etw_types.rs

@@ -12,6 +12,7 @@ use crate::trace::{CallbackData, TraceProperties, TraceTrait};
 use std::ffi::c_void;
 use std::fmt::Formatter;
 use std::marker::PhantomData;
+use std::sync::Arc;
 use windows::core::GUID;
 use windows::core::PWSTR;
 use windows::Win32::System::Diagnostics::Etw;
@@ -242,7 +243,7 @@ pub struct EventTraceLogfile<'callbackdata> {
 
 impl<'callbackdata> EventTraceLogfile<'callbackdata> {
     /// Create a new instance
-    pub fn create(callback_data: &'callbackdata Box<CallbackData>, trace_name: &str, callback: unsafe extern "system" fn(*mut Etw::EVENT_RECORD)) -> Self {
+    pub fn create(callback_data: &'callbackdata Box<Arc<CallbackData>>, trace_name: &str, callback: unsafe extern "system" fn(*mut Etw::EVENT_RECORD)) -> Self {
         let mut native = Etw::EVENT_TRACE_LOGFILEW::default();
 
         let mut wide_logger_name = U16CString::from_str_truncate(trace_name);
@@ -252,7 +253,7 @@ impl<'callbackdata> EventTraceLogfile<'callbackdata> {
 
         native.Anonymous2.EventRecordCallback = Some(callback);
 
-        let not_really_mut_ptr = callback_data.as_ref() as *const CallbackData as *const c_void as *mut c_void; // That's kind-of fine because the user context is _not supposed_ to be changed by Windows APIs
+        let not_really_mut_ptr = callback_data.as_ref() as *const Arc<CallbackData> as *const c_void as *mut c_void; // That's kind-of fine because the user context is _not supposed_ to be changed by Windows APIs
         native.Context = not_really_mut_ptr;
 
         Self {

src/native/evntrace.rs

@@ -59,7 +59,7 @@ extern "system" fn trace_callback_thunk(p_record: *mut Etw::EVENT_RECORD) {
         };
 
         if let Some(event_record) = record_from_ptr {
-            let p_user_context = event_record.user_context().cast::<CallbackData>();
+            let p_user_context = event_record.user_context().cast::<Arc<CallbackData>>();
             let user_context = unsafe {
                 // Safety:
                 //  * the API of this create guarantees this points to a `TraceData` already created
@@ -68,7 +68,10 @@ extern "system" fn trace_callback_thunk(p_record: *mut Etw::EVENT_RECORD) {
                 p_user_context.as_ref()
             };
             if let Some(user_context) = user_context {
-                user_context.on_event(event_record);
+                // The UserContext is owned by the `NativeEtw` object. When it is dropped, so will the UserContext.
+                // We clone it now, so that the original Arc can be safely dropped at all times, but the callback data (including the closure captured context) will still be alive until the callback ends.
+                let cloned_arc = Arc::clone(user_context);
+                cloned_arc.on_event(event_record);
             }
         }
     })) {

src/trace.rs

@@ -2,6 +2,7 @@
 //!
 //! Provides both a Kernel and User trace that allows to start an ETW session
 use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
 use super::traits::*;
 use crate::native::etw_types::{EventRecord, EventTraceProperties};