Merge pull request #192 from kylewo/add-provider-group

Add ability to set provider group

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.ID = opts.id
-	provider.callback = callback
+	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()
 

pkg/etw/provider.go

@@ -120,7 +120,7 @@ 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))
 }
 
 // Close unregisters the provider.

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.NewProviderWithOptions("TestProviderWithGroup", etw.WithGroup(group), etw.WithCallback(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
+	}
 }