diff --git a/pkg/etw/provider.go b/pkg/etw/provider.go index 1853ce2..4490b08 100644 --- a/pkg/etw/provider.go +++ b/pkg/etw/provider.go @@ -3,6 +3,7 @@ package etw import ( "bytes" "encoding/binary" + "sync" "unsafe" "golang.org/x/sys/windows" @@ -20,8 +21,14 @@ const ( // name and ID (GUID), which should always have a 1:1 mapping to each other // (e.g. don't use multiple provider names with the same ID, or vice versa). type Provider struct { - handle providerHandle - metadata []byte + handle providerHandle + metadata []byte + callback EnableCallback + index uint + enabled bool + level Level + keywordAny uint64 + keywordAll uint64 } type providerHandle windows.Handle @@ -61,18 +68,87 @@ func (descriptor *eventDataDescriptor) set(dataType eventDataDescriptorType, buf descriptor.dataType = dataType } -// NewProvider creates and registers a new provider. -func NewProvider(name string, id *windows.GUID, callback EnableCallback) (*Provider, error) { - provider := &Provider{} +// Because the provider callback function needs to be able to access the +// provider data when it is invoked by ETW, we need to keep provider data stored +// in a global map based on an index. The index is passed as the callback +// context to ETW. +type providerMap struct { + m map[uint]*Provider + i uint + lock sync.Mutex +} - innerCallback := func(sourceID *windows.GUID, state ProviderState, level Level, matchAnyKeyword uint64, matchAllKeyword uint64, filterData uintptr, _ uintptr) uintptr { - if callback != nil { - callback(sourceID, state, level, matchAnyKeyword, matchAllKeyword, filterData) - } - return 0 +var providers = providerMap{ + m: make(map[uint]*Provider), +} + +func (p *providerMap) newProvider() *Provider { + p.lock.Lock() + defer p.lock.Unlock() + + i := p.i + p.i++ + + provider := &Provider{ + index: i, } - if err := eventRegister(id, windows.NewCallback(innerCallback), 0, &provider.handle); err != nil { + p.m[i] = provider + return provider +} + +func (p *providerMap) removeProvider(provider *Provider) { + p.lock.Lock() + defer p.lock.Unlock() + + delete(p.m, provider.index) +} + +func (p *providerMap) getProvider(index uint) *Provider { + p.lock.Lock() + defer p.lock.Unlock() + + return p.m[index] +} + +func providerCallback(sourceID *windows.GUID, state ProviderState, level Level, matchAnyKeyword uint64, matchAllKeyword uint64, filterData uintptr, i uintptr) { + provider := providers.getProvider(uint(i)) + + switch state { + case ProviderStateDisable: + provider.enabled = false + case ProviderStateEnable: + provider.enabled = true + provider.level = level + provider.keywordAny = matchAnyKeyword + provider.keywordAll = matchAllKeyword + } + + if provider.callback != nil { + provider.callback(sourceID, state, level, matchAnyKeyword, matchAllKeyword, filterData) + } +} + +// providerCallbackAdapter acts as the first-level callback from the C/ETW side +// for provider notifications. Because Go has trouble with callback arguments of +// different size, it has only pointer-sized arguments, which are then cast to +// the appropriate types when calling providerCallback. +func providerCallbackAdapter(sourceID *windows.GUID, state uintptr, level uintptr, matchAnyKeyword uintptr, matchAllKeyword uintptr, filterData uintptr, i uintptr) uintptr { + providerCallback(sourceID, ProviderState(state), Level(level), uint64(matchAnyKeyword), uint64(matchAllKeyword), filterData, i) + return 0 +} + +// NewProvider creates and registers a new provider. +func NewProvider(name string, id *windows.GUID, callback EnableCallback) (provider *Provider, err error) { + provider = providers.newProvider() + defer func() { + if err != nil { + providers.removeProvider(provider) + } + }() + provider.callback = callback + + if err := eventRegister(id, windows.NewCallback(providerCallbackAdapter), uintptr(provider.index), &provider.handle); err != nil { return nil, err } @@ -88,9 +164,45 @@ func NewProvider(name string, id *windows.GUID, callback EnableCallback) (*Provi // Close unregisters the provider. func (provider *Provider) Close() error { + providers.removeProvider(provider) return eventUnregister(provider.handle) } +// IsEnabled calls IsEnabledForLevelAndKeywords with LevelAlways and all +// keywords set. +func (provider *Provider) IsEnabled() bool { + return provider.IsEnabledForLevelAndKeywords(LevelAlways, ^uint64(0)) +} + +// IsEnabledForLevel calls IsEnabledForLevelAndKeywords with the specified level +// and all keywords set. +func (provider *Provider) IsEnabledForLevel(level Level) bool { + return provider.IsEnabledForLevelAndKeywords(level, ^uint64(0)) +} + +// IsEnabledForLevelAndKeywords allows event producer code to check if there are +// any event sessions that are interested in an event, based on the event level +// and keywords. Although this check happens automatically in the ETW +// infrastructure, it can be useful to check if an event will actually be +// consumed before doing expensive work to build the event data. +func (provider *Provider) IsEnabledForLevelAndKeywords(level Level, keywords uint64) bool { + if !provider.enabled { + return false + } + + // ETW automatically sets the level to 255 if it is specified as 0, so we + // don't need to worry about the level=0 (all events) case. + if level > provider.level { + return false + } + + if keywords != 0 && (keywords&provider.keywordAny == 0 || keywords&provider.keywordAll != provider.keywordAll) { + return false + } + + return true +} + // WriteEvent writes a single event to ETW, from this provider. func (provider *Provider) WriteEvent(event *Event) error { // Finalize the event metadata buffer by filling in the buffer length at the diff --git a/pkg/etwlogrus/hook.go b/pkg/etwlogrus/hook.go index e13747c..6037388 100644 --- a/pkg/etwlogrus/hook.go +++ b/pkg/etwlogrus/hook.go @@ -44,6 +44,10 @@ func (h *Hook) Levels() []logrus.Level { // Fire receives each Logrus entry as it is logged, and logs it to ETW. func (h *Hook) Fire(e *logrus.Entry) error { + if !h.provider.IsEnabledForLevel(etw.Level(e.Level)) { + return + } + descriptor := etw.NewEventDescriptor() // We could try to map Logrus levels to ETW levels, but we would lose some