From dc06ae791a7244d0d244b6c34b1e7930bc5841b0 Mon Sep 17 00:00:00 2001 From: John Uhlmann Date: Wed, 1 Apr 2026 15:24:53 +0800 Subject: [PATCH] PR feedback --- O365.Security.Native.ETW.Debug.nuspec | 2 +- O365.Security.Native.ETW.nuspec | 2 +- krabs/krabs/parser.hpp | 38 +++-------- krabs/krabs/schema.hpp | 10 ++- krabs/krabs/schema_locator.hpp | 91 ++++++++++++++------------- krabsetw.nuspec | 2 +- 6 files changed, 66 insertions(+), 79 deletions(-) diff --git a/O365.Security.Native.ETW.Debug.nuspec b/O365.Security.Native.ETW.Debug.nuspec index 4961ca7..4942958 100644 --- a/O365.Security.Native.ETW.Debug.nuspec +++ b/O365.Security.Native.ETW.Debug.nuspec @@ -2,7 +2,7 @@ Microsoft.O365.Security.Native.ETW.Debug - 4.4.8 + 4.4.9 Microsoft.O365.Security.Native.ETW Debug - managed wrappers for krabsetw Microsoft Microsoft diff --git a/O365.Security.Native.ETW.nuspec b/O365.Security.Native.ETW.nuspec index 7d0945a..f129c8a 100644 --- a/O365.Security.Native.ETW.nuspec +++ b/O365.Security.Native.ETW.nuspec @@ -2,7 +2,7 @@ Microsoft.O365.Security.Native.ETW - 4.4.8 + 4.4.9 Microsoft.O365.Security.Native.ETW - managed wrappers for krabsetw Microsoft Microsoft diff --git a/krabs/krabs/parser.hpp b/krabs/krabs/parser.hpp index 2736857..03468d6 100644 --- a/krabs/krabs/parser.hpp +++ b/krabs/krabs/parser.hpp @@ -127,40 +127,17 @@ namespace krabs { // A schema contains a collection of properties that are keyed by name. // These properties are stored in a blob of bytes that needs to be // interpreted according to information that is packaged up in the - // schema and that can be retrieved using the Tdh* APIs. This format - // requires a linear traversal over the blob, incrementing according to - // the contents within it. This is janky, so our strategy is to - // minimize this as much as possible via caching. + // schema and that can be retrieved using the Tdh* APIs. + // We have pre-calculated and cached a name to index map. - const ULONG totalPropCount = schema_.pSchema_->PropertyCount; - - // Resolve property name to index. - ULONG index = totalPropCount; // sentinel = not found - if (pPropertyNames_) { - // Fast path: use the persistent name to index map shared across - // all events of the same type. - auto it = pPropertyNames_->find(name); - if (it != pPropertyNames_->end()) { - index = it->second; - } - } else { - // Fallback: linear scan of property names in the schema. - for (ULONG i = 0; i < totalPropCount; ++i) { - auto &propInfo = schema_.pSchema_->EventPropertyInfoArray[i]; - const wchar_t *pName = reinterpret_cast( - reinterpret_cast(schema_.pSchema_) + - propInfo.NameOffset); - if (name == pName) { - index = i; - break; - } - } - } - - if (index >= totalPropCount) { + // Resolve property name to index via map. + const auto it = pPropertyNames_->find(name); + if (it == pPropertyNames_->end()) { // not found return property_info(); } + const ULONG index = it->second; + // The first step is to use our cache for the property to see if we've // discovered it already. if (index < lastPropertyIndex_) { @@ -172,6 +149,7 @@ namespace krabs { // accept that last property can be omitted from buffer. this happens if last property // is string but empty and the provider strips the null terminator + const ULONG totalPropCount = schema_.pSchema_->PropertyCount; assert((pBufferIndex_ == pEndBuffer_ ? ((totalPropCount - lastPropertyIndex_) <= 1) : true) && "invariant: if we've exhausted our buffer, then we must've" diff --git a/krabs/krabs/schema.hpp b/krabs/krabs/schema.hpp index 5212e8c..192500d 100644 --- a/krabs/krabs/schema.hpp +++ b/krabs/krabs/schema.hpp @@ -306,6 +306,8 @@ namespace krabs { ULONG64 process_start_key() const; private: + schema(const EVENT_RECORD &, std::pair); + const EVENT_RECORD &record_; const TRACE_EVENT_INFO *pSchema_; // Persistent name to index map, owned by schema_locator. May be nullptr. @@ -337,9 +339,13 @@ namespace krabs { // ------------------------------------------------------------------------ inline schema::schema(const EVENT_RECORD &record, const krabs::schema_locator &schema_locator) + : schema(record, schema_locator.get_schema_and_names(record)) + { } + + inline schema::schema(const EVENT_RECORD &record, std::pair p) : record_(record) - , pSchema_(schema_locator.get_event_schema(record)) - , pPropertyNames_(schema_locator.get_property_names(pSchema_)) + , pSchema_(p.first) + , pPropertyNames_(p.second) { } inline schema::schema(const EVENT_RECORD &record, const PTRACE_EVENT_INFO pSchema) diff --git a/krabs/krabs/schema_locator.hpp b/krabs/krabs/schema_locator.hpp index db00d56..fbc9114 100644 --- a/krabs/krabs/schema_locator.hpp +++ b/krabs/krabs/schema_locator.hpp @@ -217,19 +217,24 @@ namespace krabs { /** * - * Returns the persistent property name to index map for a schema. + * Retrieves the event schema and its persistent property name to index + * map in a single cache lookup. Throws on error. * The map is built when the schema is first cached. - * Returns nullptr if pSchema is null or not in the cache. * */ - const property_name_map* get_property_names(const TRACE_EVENT_INFO* pSchema) const; + std::pair + get_schema_and_names(const EVENT_RECORD& record) const; private: - void build_property_names(const TRACE_EVENT_INFO* pSchema) const; + struct schema_entry { + std::unique_ptr buffer; + property_name_map name_map; + }; - mutable std::unordered_map, TDHSTATUS>> cache_; - // Persistent property name to index maps, keyed by schema pointer. - mutable std::unordered_map property_name_cache_; + std::pair + get_schema_and_names_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const; + + mutable std::unordered_map> cache_; }; // Implementation @@ -301,14 +306,26 @@ namespace krabs { } inline const PTRACE_EVENT_INFO schema_locator::get_event_schema(const EVENT_RECORD &record) const + { + return get_schema_and_names(record).first; + } + + inline const PTRACE_EVENT_INFO schema_locator::get_event_schema_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const + { + return get_schema_and_names_no_throw(record, status).first; + } + + inline std::pair + schema_locator::get_schema_and_names(const EVENT_RECORD& record) const { TDHSTATUS status = ERROR_SUCCESS; - auto buffer = get_event_schema_no_throw(record, status); + auto result = get_schema_and_names_no_throw(record, status); error_check_common_conditions(status, record); - return buffer; + return result; } - inline const PTRACE_EVENT_INFO schema_locator::get_event_schema_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const + inline std::pair + schema_locator::get_schema_and_names_no_throw(const EVENT_RECORD& record, TDHSTATUS& status) const { status = ERROR_SUCCESS; @@ -321,24 +338,33 @@ namespace krabs { auto& value = it->second; if (std::holds_alternative(value)) { status = std::get(value); - return nullptr; + return {nullptr, nullptr}; } - return (PTRACE_EVENT_INFO)std::get>(value).get(); + auto& entry = std::get(value); + return {(PTRACE_EVENT_INFO)entry.buffer.get(), &entry.name_map}; } - // Cache miss. Fetch the schema... + // Cache miss. Fetch the schema from TDH. + // NB: key's 'internalize_name' gets called by the cctor during emplace. auto buffer = get_event_schema_from_tdh_no_throw(record, status); - auto returnVal = (PTRACE_EVENT_INFO)buffer.get(); + auto pSchema = (PTRACE_EVENT_INFO)buffer.get(); - // Add the new instance to the cache. - // NB: key's 'internalize_name' gets called by the cctor here. if (status == ERROR_SUCCESS) { - cache_.emplace(key, std::move(buffer)); - build_property_names(returnVal); - } else + schema_entry entry; + entry.buffer = std::move(buffer); + for (ULONG i = 0; i < pSchema->PropertyCount; ++i) { + const wchar_t* pName = reinterpret_cast( + reinterpret_cast(pSchema) + + pSchema->EventPropertyInfoArray[i].NameOffset); + entry.name_map.emplace(std::wstring_view(pName), i); + } + auto result = cache_.emplace(key, std::move(entry)); + auto& inserted = std::get(result.first->second); + return {(PTRACE_EVENT_INFO)inserted.buffer.get(), &inserted.name_map}; + } else { cache_.emplace(key, status); - - return returnVal; + return {nullptr, nullptr}; + } } inline bool schema_locator::has_event_schema(const EVENT_RECORD& record) const @@ -348,29 +374,6 @@ namespace krabs { return status == ERROR_SUCCESS; } - inline void schema_locator::build_property_names(const TRACE_EVENT_INFO* pSchema) const - { - property_name_map names; - for (ULONG i = 0; i < pSchema->PropertyCount; ++i) { - const wchar_t* pName = reinterpret_cast( - reinterpret_cast(pSchema) + - pSchema->EventPropertyInfoArray[i].NameOffset); - names.emplace(std::wstring_view(pName), i); - } - property_name_cache_.emplace(pSchema, std::move(names)); - } - - inline const property_name_map* schema_locator::get_property_names(const TRACE_EVENT_INFO* pSchema) const - { - if (!pSchema) return nullptr; - - auto it = property_name_cache_.find(pSchema); - if (it != property_name_cache_.end()) { - return &it->second; - } - return nullptr; - } - inline std::unique_ptr get_event_schema_from_tdh(const EVENT_RECORD &record) { TDHSTATUS status = ERROR_SUCCESS; diff --git a/krabsetw.nuspec b/krabsetw.nuspec index f1ce3b5..506ecca 100644 --- a/krabsetw.nuspec +++ b/krabsetw.nuspec @@ -2,7 +2,7 @@ Microsoft.O365.Security.Krabsetw - 4.4.8 + 4.4.9 Krabs ETW Wrappers Microsoft Microsoft