From 30844eed7547cc370d53f0ddf12ded209a445b78 Mon Sep 17 00:00:00 2001 From: Kyle Wojtaszek Date: Wed, 27 Jan 2021 13:12:37 -0800 Subject: [PATCH] Add ability to set provider group --- pkg/etw/newprovider.go | 63 ++++++++++++++++++++++++++++++++++++++-- pkg/etw/provider.go | 8 ++++- pkg/etw/sample/sample.go | 43 +++++++++++++++++++++++++++ 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/pkg/etw/newprovider.go b/pkg/etw/newprovider.go index f344fb6..eb99cea 100644 --- a/pkg/etw/newprovider.go +++ b/pkg/etw/newprovider.go @@ -11,11 +11,60 @@ import ( "golang.org/x/sys/windows" ) +type providerOpts struct { + callback EnableCallback + id guid.GUID + group guid.GUID +} + +// ProviderOpt allows the caller to specify provider options to +// NewProviderWithOptions +type ProviderOpt func(*providerOpts) + +// WithCallback is used to provide a callback option to NewProviderWithOptions +func WithCallback(callback EnableCallback) ProviderOpt { + return func(opts *providerOpts) { + opts.callback = callback + } +} + +// WithID is used to provide a provider ID option to NewProviderWithOptions +func WithID(id guid.GUID) ProviderOpt { + return func(opts *providerOpts) { + opts.id = id + } +} + +// WithGroup is used to provide a provider group option to +// NewProviderWithOptions +func WithGroup(group guid.GUID) ProviderOpt { + return func(opts *providerOpts) { + opts.group = group + } +} + // NewProviderWithID creates and registers a new ETW provider, allowing the // provider ID to be manually specified. This is most useful when there is an // existing provider ID that must be used to conform to existing diagnostic // infrastructure. func NewProviderWithID(name string, id guid.GUID, callback EnableCallback) (provider *Provider, err error) { + return NewProviderWithOptions(name, WithID(id), WithCallback(callback)) +} + +// NewProviderWithOptions creates and registers a new ETW provider, allowing +// the provider ID and Group to be manually specified. This is most useful when +// there is an existing provider ID that must be used to conform to existing +// diagnostic infrastructure. +func NewProviderWithOptions(name string, options ...ProviderOpt) (provider *Provider, err error) { + var opts providerOpts + for _, opt := range options { + opt(&opts) + } + + if opts.id == (guid.GUID{}) { + opts.id = providerIDFromName(name) + } + providerCallbackOnce.Do(func() { globalProviderCallback = windows.NewCallback(providerCallbackAdapter) }) @@ -26,17 +75,27 @@ func NewProviderWithID(name string, id guid.GUID, callback EnableCallback) (prov providers.removeProvider(provider) } }(provider) - provider.ID = id - provider.callback = callback + provider.ID = opts.id + provider.callback = opts.callback if err := eventRegister((*windows.GUID)(&provider.ID), globalProviderCallback, uintptr(provider.index), &provider.handle); err != nil { return nil, err } + trait := &bytes.Buffer{} + if opts.group != (guid.GUID{}) { + binary.Write(trait, binary.LittleEndian, uint16(0)) // Write empty size for buffer (update later) + binary.Write(trait, binary.LittleEndian, uint8(1)) // EtwProviderTraitTypeGroup + traitArray := opts.group.ToWindowsArray() // Append group guid + trait.Write(traitArray[:]) + binary.LittleEndian.PutUint16(trait.Bytes(), uint16(trait.Len())) // Update size + } + metadata := &bytes.Buffer{} binary.Write(metadata, binary.LittleEndian, uint16(0)) // Write empty size for buffer (to update later) metadata.WriteString(name) metadata.WriteByte(0) // Null terminator for name + trait.WriteTo(metadata) // Add traits if applicable binary.LittleEndian.PutUint16(metadata.Bytes(), uint16(metadata.Len())) // Update the size at the beginning of the buffer provider.metadata = metadata.Bytes() diff --git a/pkg/etw/provider.go b/pkg/etw/provider.go index 2369601..285038a 100644 --- a/pkg/etw/provider.go +++ b/pkg/etw/provider.go @@ -120,7 +120,13 @@ func providerIDFromName(name string) guid.GUID { // NewProvider creates and registers a new ETW provider. The provider ID is // generated based on the provider name. func NewProvider(name string, callback EnableCallback) (provider *Provider, err error) { - return NewProviderWithID(name, providerIDFromName(name), callback) + return NewProviderWithOptions(name, WithCallback(callback)) +} + +// NewProviderWithGroup creates and registers a new ETW provider with a +// specificed group. The provider ID is generated based on the provider name. +func NewProviderWithGroup(name string, group guid.GUID, callback EnableCallback) (provider *Provider, err error) { + return NewProviderWithOptions(name, WithGroup(group), WithCallback(callback)) } // Close unregisters the provider. diff --git a/pkg/etw/sample/sample.go b/pkg/etw/sample/sample.go index 0f3bebe..9eb2962 100644 --- a/pkg/etw/sample/sample.go +++ b/pkg/etw/sample/sample.go @@ -16,6 +16,12 @@ func callback(sourceID guid.GUID, state etw.ProviderState, level etw.Level, matc } func main() { + group, err := guid.FromString("12341234-abcd-abcd-abcd-123412341234") + if err != nil { + logrus.Error(err) + return + } + provider, err := etw.NewProvider("TestProvider", callback) if err != nil { @@ -28,7 +34,20 @@ func main() { } }() + providerWithGroup, err := etw.NewProviderWithGroup("TestProviderWithGroup", group, callback) + + if err != nil { + logrus.Error(err) + return + } + defer func() { + if err := providerWithGroup.Close(); err != nil { + logrus.Error(err) + } + }() + fmt.Printf("Provider ID: %s\n", provider) + fmt.Printf("Provider w/ Group ID: %s\n", providerWithGroup) reader := bufio.NewReader(os.Stdin) @@ -58,4 +77,28 @@ func main() { logrus.Error(err) return } + + if err := providerWithGroup.WriteEvent( + "TestEventWithGroup", + etw.WithEventOpts( + etw.WithLevel(etw.LevelInfo), + etw.WithKeyword(0x140), + ), + etw.WithFields( + etw.StringField("TestField", "Foo"), + etw.StringField("TestField2", "Bar"), + etw.Struct("TestStruct", + etw.StringField("Field1", "Value1"), + etw.StringField("Field2", "Value2")), + etw.StringArray("TestArray", []string{ + "Item1", + "Item2", + "Item3", + "Item4", + "Item5", + })), + ); err != nil { + logrus.Error(err) + return + } }