diff --git a/krabs/krabs/etw.hpp b/krabs/krabs/etw.hpp index 0bc2265..aff2b4a 100644 --- a/krabs/krabs/etw.hpp +++ b/krabs/krabs/etw.hpp @@ -202,6 +202,9 @@ namespace krabs { namespace details { PVOID trace_information, ULONG information_length) { + if (trace_.registrationHandle_ == INVALID_PROCESSTRACE_HANDLE) + return; + ULONG status = TraceSetInformation( trace_.registrationHandle_, information_class, @@ -283,12 +286,16 @@ namespace krabs { namespace details { template void trace_manager::stop_trace() { + if (trace_.registrationHandle_ == INVALID_PROCESSTRACE_HANDLE) + return; + trace_info info = fill_trace_info(); ULONG status = ControlTrace( NULL, trace_.name_.c_str(), &info.properties, EVENT_TRACE_CONTROL_STOP); + trace_.registrationHandle_ = INVALID_PROCESSTRACE_HANDLE; if (status != ERROR_WMI_INSTANCE_NOT_FOUND) { error_check_common_conditions(status); diff --git a/krabs/krabs/trace.hpp b/krabs/krabs/trace.hpp index 81b076e..5f768ff 100644 --- a/krabs/krabs/trace.hpp +++ b/krabs/krabs/trace.hpp @@ -282,6 +282,12 @@ namespace krabs { */ void set_default_event_callback(c_provider_callback callback); + template + void set_default_event_callback(U& callback); + + template + void set_default_event_callback(const U& callback); + private: /** @@ -348,7 +354,8 @@ namespace krabs { template trace::~trace() { - stop(); + // Never throw in destructor + try { stop(); } catch (...) {} } template @@ -440,7 +447,30 @@ namespace krabs { template void trace::set_default_event_callback(c_provider_callback callback) { + // C function pointers don't interact well with std::ref, so we + // overload to take care of this scenario. default_callback_ = callback; } + template + template + void trace::set_default_event_callback(U& callback) + { + // std::function copies its argument -- because our callbacks list + // is a list of std::function, this causes problems when a user + // intended for their particular instance to be called. + // std::ref lets us get around this and point to a specific instance + // that they handed us. + default_callback_ = std::ref(callback); + } + + template + template + void trace::set_default_event_callback(const U& callback) + { + // This is where temporaries bind to. Temporaries can't be wrapped in + // a std::ref because they'll go away very quickly. We are forced to + // actually copy these. + default_callback_ = callback; + } } diff --git a/krabs/krabs/ut.hpp b/krabs/krabs/ut.hpp index b331195..4cabde0 100644 --- a/krabs/krabs/ut.hpp +++ b/krabs/krabs/ut.hpp @@ -164,7 +164,7 @@ namespace krabs { namespace details { std::vector filterEventIdBuffer; auto filterEventIdCount = settings.provider_filter_event_ids_.size(); - if (filterEventIdCount > 0) { + if (filterEventIdCount > 0 && filterEventIdCount <= MAX_EVENT_FILTER_EVENT_ID_COUNT) { //event filters existing, set native filters using API parameters.FilterDescCount = 1; filterDesc.Type = EVENT_FILTER_TYPE_EVENT_ID; @@ -237,15 +237,18 @@ namespace krabs { namespace details { // for MOF providers, EventHeader.Provider is the *Message* GUID // we need to ask TDH for event information in order to determine the - // correct provider to pass this event to - auto schema = get_event_schema_from_tdh(record); - auto eventInfo = reinterpret_cast(schema.get()); - for (auto& provider : trace.providers_) { - if (eventInfo->ProviderGuid == provider.get().guid_) { - provider.get().on_event(record, trace.context_); - return; + // correct provider to pass this event to. Don't throw in callback. + try + { + auto eventInfo = trace.context_.schema_locator.get_event_schema(record); + for (auto& provider : trace.providers_) { + if (eventInfo->ProviderGuid == provider.get().guid_) { + provider.get().on_event(record, trace.context_); + return; + } } } + catch (krabs::could_not_find_schema&) {} if (trace.default_callback_ != nullptr) trace.default_callback_(record, trace.context_);