Unify. NET wrappers of callbacks and more optimizations

Microsoft.O365.Security.Native.ETW/Callbacks.hpp

@@ -0,0 +1,165 @@
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+
+#pragma once
+
+#include <krabs.hpp>
+
+#include "EventRecord.hpp"
+#include "EventRecordError.hpp"
+
+using namespace System;
+using namespace System::Runtime::InteropServices;
+
+namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
+
+    /// <summary>
+    /// Delegate called when a new ETW <see cref="O365::Security::ETW::EventRecord"/> is received.
+    /// </summary>
+    public delegate void IEventRecordDelegate(IEventRecord^ record);
+
+    /// <summary>
+    /// Delegate called when a new ETW <see cref="O365::Security::ETW::EventRecordMetadata"/> is received.
+    /// </summary>
+    public delegate void IEventRecordMetadataDelegate(IEventRecordMetadata^ record);
+
+    /// <summary>
+    /// Delegate called on errors when processing an <see cref="O365::Security::ETW::EventRecord"/>.
+    /// </summary>
+    public delegate void EventRecordErrorDelegate(IEventRecordError^ error);
+
+    delegate void EventNativeDelegate(const EVENT_RECORD&, const krabs::trace_context&);
+
+    delegate void EventErrorNativeDelegate(const EVENT_RECORD&, const std::string&);
+
+    ref class CallbackBridge {
+    private:
+
+        EventNativeDelegate^ eventDelegateKeepAlive_;
+        EventErrorNativeDelegate^ errorDelegateKeepAlive_;
+        EventRecordMetadata^ metadata_;
+        EventRecord^ record_;
+        EventRecordError^ error_;
+
+        EventRecordMetadata^ WrapMetadata(const EVENT_RECORD& record);
+        EventRecord^ WrapRecord(const EVENT_RECORD& record, const krabs::schema& schema, krabs::parser& parser);
+        EventRecordError^ WrapError(const std::string& error_message, const EVENT_RECORD& record);
+
+        void EventNotification(const EVENT_RECORD& record, const krabs::trace_context& trace_context);
+        void ErrorNotification(const EVENT_RECORD& record, const std::string& error_message);
+
+    public:
+        IEventRecordMetadataDelegate^ OnMetadata;
+        IEventRecordDelegate^ OnEvent;
+        EventRecordErrorDelegate^ OnError;
+
+        krabs::c_provider_callback GetOnEventBridge()
+        {
+            if (!eventDelegateKeepAlive_)
+                eventDelegateKeepAlive_ = gcnew EventNativeDelegate(this, &CallbackBridge::EventNotification);
+            return (krabs::c_provider_callback)Marshal::GetFunctionPointerForDelegate(eventDelegateKeepAlive_).ToPointer();
+        }
+
+        krabs::c_provider_error_callback GetOnErrorBridge()
+        {
+            if (!errorDelegateKeepAlive_)
+                errorDelegateKeepAlive_ = gcnew EventErrorNativeDelegate(this, &CallbackBridge::ErrorNotification);
+            return (krabs::c_provider_error_callback)Marshal::GetFunctionPointerForDelegate(errorDelegateKeepAlive_).ToPointer();
+        }
+    };
+
+    EventRecordMetadata^ CallbackBridge::WrapMetadata(const EVENT_RECORD& record)
+    {
+        auto value = metadata_;
+        if (!value)
+            return metadata_ = gcnew EventRecordMetadata(record);
+
+        value->Update(record);
+        return value;
+    }
+
+    EventRecord^ CallbackBridge::WrapRecord(const EVENT_RECORD& record, const krabs::schema& schema, krabs::parser& parser)
+    {
+        auto value = record_;
+        if (!value) 
+            return record_ = gcnew EventRecord(record, schema, parser);
+
+        value->Update(record, schema, parser);
+        return value;
+    }
+
+    EventRecordError^ CallbackBridge::WrapError(const std::string& error_message, const EVENT_RECORD& record)
+    {
+        auto msg = gcnew String(error_message.c_str());
+        auto value = error_;
+        if (!value)
+            return error_ = gcnew EventRecordError(msg, gcnew EventRecordMetadata(record));
+
+        value->Update(msg, record);
+        return value;
+    }
+
+    void CallbackBridge::EventNotification(const EVENT_RECORD& record, const krabs::trace_context& trace_context)
+    {
+        auto onMetadata = OnMetadata;
+        if (onMetadata) {
+            auto metadata = WrapMetadata(record);
+
+            onMetadata(metadata);
+        }
+
+        auto onEvent = OnEvent;
+        if (onEvent) {
+            TDHSTATUS status = ERROR_SUCCESS;
+            trace_context.schema_locator.get_event_schema_no_throw(record, status);
+
+            if (status == ERROR_SUCCESS) {
+                krabs::schema schema(record, trace_context.schema_locator);
+                krabs::parser parser(schema);
+                auto evt = WrapRecord(record, schema, parser);
+
+                onEvent(evt);
+            }
+            else {
+                auto error_message = krabs::get_status_and_record_context(status, record);
+                ErrorNotification(record, error_message);
+            }
+        }
+    }
+
+    void CallbackBridge::ErrorNotification(const EVENT_RECORD& record, const std::string& error_message)
+    {
+        auto onError = OnError;
+        if (onError) {
+            auto error = WrapError(error_message, record);
+
+            onError(error);
+        }
+    }
+
+    template<typename T>
+    inline void CombineOrRemoveDelegate(T% target, T value, bool remove)
+    {
+        T original = target;
+        T comparand;
+        do
+        {
+            comparand = original;
+            T newValue = remove ? (T)System::Delegate::Remove(original, value) : (T)System::Delegate::Combine(original, value);
+            original = (T)System::Threading::Interlocked::CompareExchange(target, newValue, comparand);
+        } while (original != comparand);
+    }
+
+    template<typename T>
+    inline void CombineDelegate(T% target, T value)
+    {
+        CombineOrRemoveDelegate(target, value, false);
+    }
+
+    template<typename T>
+    inline void RemoveDelegate(T% target, T value)
+    {
+        CombineOrRemoveDelegate(target, value, true);
+    }
+
+} } } }

Microsoft.O365.Security.Native.ETW/EventRecord.hpp

@@ -40,6 +40,16 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
             , schema_(&schema)
             , parser_(&parser) { }
 
+        /// <summary>
+        /// Updates this instance to point to the specified event record.
+        /// </summary>
+        void Update(const EVENT_RECORD& record, const krabs::schema& schema, krabs::parser& parser)
+        {
+            Update(record);
+            schema_ = &schema;
+            parser_ = &parser;
+        }
+
     public:
 
 #pragma region Schema

Microsoft.O365.Security.Native.ETW/EventRecordError.hpp

@@ -15,8 +15,25 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
     public ref struct EventRecordError : public IEventRecordError
     {
     private:
-        initonly System::String^ msg_;
-        initonly IEventRecordMetadata^ record_;
+        System::String^ msg_;
+        EventRecordMetadata^ record_;
+
+    internal:
+        EventRecordError(
+            System::String^ message,
+            EventRecordMetadata^ record)
+            : msg_(message)
+            , record_(record)
+        { }
+
+        /// <summary>
+        /// Updates this instance to point to the specified event record.
+        /// </summary>
+        void Update(System::String^ msg, const EVENT_RECORD& record)
+        {
+            msg_ = msg;
+            record_->Update(record);
+        }
 
     public:
         /// <summary>
@@ -39,14 +56,6 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
                 return record_;
             }
         }
-
-    internal:
-        EventRecordError(
-            System::String^ message,
-            IEventRecordMetadata^ record)
-            : msg_(message)
-            , record_(record)
-        { }
     };
 
 } } } }

Microsoft.O365.Security.Native.ETW/EventRecordMetadata.hpp

@@ -26,12 +26,10 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
             : record_(&record)
             , header_(&record.EventHeader) { }
 
-        EventRecordMetadata() { }
-
         /// <summary>
         /// Updates this instance to point to the specified event record.
         /// </summary>
-        virtual void Update(const EVENT_RECORD& record)
+        void Update(const EVENT_RECORD& record)
         {
             record_ = &record;
             header_ = &record.EventHeader;

Microsoft.O365.Security.Native.ETW/Filtering/EventFilter.hpp

@@ -5,13 +5,9 @@
 
 #include <krabs.hpp>
 
+#include "../Callbacks.hpp"
 #include "../Conversions.hpp"
-#include "../EventRecordError.hpp"
-#include "../EventRecord.hpp"
-#include "../EventRecordMetadata.hpp"
 #include "../Guid.hpp"
-#include "../IEventRecord.hpp"
-#include "../IEventRecordError.hpp"
 #include "../NativePtr.hpp"
 #include "Predicate.hpp"
 
@@ -20,16 +16,6 @@ using namespace System::Runtime::InteropServices;
 
 namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
 
-    /// <summary>
-    /// Delegate called when a new ETW <see cref="O365::Security::ETW::EventRecord"/> is received.
-    /// </summary>
-    public delegate void IEventRecordDelegate(O365::Security::ETW::IEventRecord^ record);
-
-    /// <summary>
-    /// Delegate called on errors when processing an <see cref="O365::Security::ETW::EventRecord"/>.
-    /// </summary>
-    public delegate void EventRecordErrorDelegate(O365::Security::ETW::IEventRecordError^ error);
-
     /// <summary>
     /// Allows for filtering an event in the native layer before it bubbles
     /// up to callbacks.
@@ -69,23 +55,23 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
         /// <param name="predicate">the predicate to use to filter an event</param>
         EventFilter(List<unsigned short>^ eventIds, O365::Security::ETW::Predicate^ predicate);
 
-        /// <summary>
-        /// Destructs an EventFilter.
-        /// </summary>
-        ~EventFilter();
-
         /// <summary>
         /// An event that is invoked when an ETW event is fired on this
         /// filter and the event meets the given predicate.
         /// </summary>
-        event IEventRecordDelegate^ OnEvent;
+        event IEventRecordDelegate^ OnEvent {
+            void add(IEventRecordDelegate^ value) { CombineDelegate(bridge_->OnEvent, value); }
+            void remove(IEventRecordDelegate^ value) { RemoveDelegate(bridge_->OnEvent, value); }
+        }
 
         /// <summary>
         /// An event that is invoked when an ETW event is received
         /// but an error occurs handling the record.
         /// </summary>
-        event EventRecordErrorDelegate^ OnError;
-
+        event EventRecordErrorDelegate^ OnError {
+            void add(EventRecordErrorDelegate^ value) { CombineDelegate(bridge_->OnError, value); }
+            void remove(EventRecordErrorDelegate^ value) { RemoveDelegate(bridge_->OnError, value); }
+        }
 
     internal:
         /// <summary>
@@ -97,20 +83,10 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
             return *filter_;
         }
 
-        void EventNotification(const EVENT_RECORD &, const krabs::trace_context &);
-        void ErrorNotification(const EVENT_RECORD&, const std::string&);
-
     internal:
-        delegate void EventReceivedNativeHookDelegate(const EVENT_RECORD &, const krabs::trace_context &);
-        delegate void ErrorReceivedNativeHookDelegate(const EVENT_RECORD &, const std::string &);
-
         NativePtr<krabs::event_filter> filter_;
-        EventReceivedNativeHookDelegate^ eventReceivedDelegate_;
-        ErrorReceivedNativeHookDelegate^ errorReceivedDelegate_;
-        GCHandle eventReceivedDelegateHookHandle_;
-        GCHandle errorReceivedDelegateHookHandle_;
-        GCHandle eventReceivedDelegateHandle_;
-        GCHandle errorReceivedDelegateHandle_;
+        CallbackBridge^ bridge_ = gcnew CallbackBridge();
+
         void RegisterCallbacks();
     };
 
@@ -147,67 +123,10 @@ namespace Microsoft { namespace O365 { namespace Security { namespace ETW {
         RegisterCallbacks();
     }
 
-    inline EventFilter::~EventFilter()
-    {
-        if (eventReceivedDelegateHandle_.IsAllocated)
-        {
-            eventReceivedDelegateHandle_.Free();
-        }
-
-        if (eventReceivedDelegateHookHandle_.IsAllocated)
-        {
-            eventReceivedDelegateHookHandle_.Free();
-        }
-
-        if (errorReceivedDelegateHandle_.IsAllocated)
-        {
-            errorReceivedDelegateHandle_.Free();
-        }
-
-        if (errorReceivedDelegateHookHandle_.IsAllocated)
-        {
-            errorReceivedDelegateHookHandle_.Free();
-        }
-    }
-
     inline void EventFilter::RegisterCallbacks()
     {
-        eventReceivedDelegate_ = gcnew EventReceivedNativeHookDelegate(this, &EventFilter::EventNotification);
-        eventReceivedDelegateHandle_ = GCHandle::Alloc(eventReceivedDelegate_);
-        auto bridgedEventDelegate = Marshal::GetFunctionPointerForDelegate(eventReceivedDelegate_);
-        eventReceivedDelegateHookHandle_ = GCHandle::Alloc(bridgedEventDelegate);
-
-        filter_->add_on_event_callback((krabs::c_provider_callback)bridgedEventDelegate.ToPointer());
-
-        errorReceivedDelegate_ = gcnew ErrorReceivedNativeHookDelegate(this, &EventFilter::ErrorNotification);
-        errorReceivedDelegateHandle_ = GCHandle::Alloc(errorReceivedDelegate_);
-        auto bridgedErrorDelegate = Marshal::GetFunctionPointerForDelegate(errorReceivedDelegate_);
-        errorReceivedDelegateHookHandle_ = GCHandle::Alloc(bridgedErrorDelegate);
-
-        filter_->add_on_error_callback((krabs::c_provider_error_callback)bridgedErrorDelegate.ToPointer());
-    }
-
-    inline void EventFilter::EventNotification(const EVENT_RECORD &record, const krabs::trace_context &trace_context)
-    {
-        try
-        {
-            krabs::schema schema(record, trace_context.schema_locator);
-            krabs::parser parser(schema);
-
-            OnEvent(gcnew EventRecord(record, schema, parser));
-        }
-        catch (const krabs::could_not_find_schema& ex)
-        {
-            ErrorNotification(record, ex.what());
-        }
-    }
-
-    inline void EventFilter::ErrorNotification(const EVENT_RECORD& record, const std::string& error_message)
-    {
-        auto msg = gcnew String(error_message.c_str());
-        auto metadata = gcnew EventRecordMetadata(record);
-
-        OnError(gcnew EventRecordError(msg, metadata));
+        filter_->add_on_event_callback(bridge_->GetOnEventBridge());
+        filter_->add_on_error_callback(bridge_->GetOnErrorBridge());
     }
 
 } } } }