From 7dba96a33334bd57d51d7b7ebafdee258dfcbf90 Mon Sep 17 00:00:00 2001
From: Kevin Parsons <kevpar@ntdev.microsoft.com>
Date: Thu, 17 Jan 2019 10:29:07 -0800
Subject: [PATCH] Set ETW provider info so it's always treated as TraceLogging

Signed-off-by: Kevin Parsons <kevpar@ntdev.microsoft.com>
---
 internal/etw/etw.go              |  1 +
 internal/etw/eventopt.go         |  6 ++++++
 internal/etw/provider.go         | 19 +++++++++++++++++++
 internal/etw/zsyscall_windows.go | 15 ++++++++++++---
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/internal/etw/etw.go b/internal/etw/etw.go
index da905ec9..5bddae26 100644
--- a/internal/etw/etw.go
+++ b/internal/etw/etw.go
@@ -12,3 +12,4 @@ package etw
 //sys eventRegister(providerId *windows.GUID, callback uintptr, callbackContext uintptr, providerHandle *providerHandle) (win32err error) = advapi32.EventRegister
 //sys eventUnregister(providerHandle providerHandle) (win32err error) = advapi32.EventUnregister
 //sys eventWriteTransfer(providerHandle providerHandle, descriptor *EventDescriptor, activityID *windows.GUID, relatedActivityID *windows.GUID, dataDescriptorCount uint32, dataDescriptors *eventDataDescriptor) (win32err error) = advapi32.EventWriteTransfer
+//sys eventSetInformation(providerHandle providerHandle, class eventInfoClass, information uintptr, length uint32) (win32err error) = advapi32.EventSetInformation
diff --git a/internal/etw/eventopt.go b/internal/etw/eventopt.go
index 5c752dc0..447cec3e 100644
--- a/internal/etw/eventopt.go
+++ b/internal/etw/eventopt.go
@@ -36,6 +36,12 @@ func WithKeyword(keyword uint64) EventOpt {
 	}
 }
 
+func WithChannel(channel Channel) EventOpt {
+	return func(options *eventOptions) {
+		options.descriptor.Channel = channel
+	}
+}
+
 // WithTags specifies the tags of the event to be written. Tags is a 28-bit
 // value (top 4 bits are ignored) which are interpreted by the event consumer.
 func WithTags(newTags uint32) EventOpt {
diff --git a/internal/etw/provider.go b/internal/etw/provider.go
index 4773a808..b656ff2f 100644
--- a/internal/etw/provider.go
+++ b/internal/etw/provider.go
@@ -6,6 +6,7 @@ import (
 	"encoding/binary"
 	"strings"
 	"unicode/utf16"
+	"unsafe"
 
 	"golang.org/x/sys/windows"
 )
@@ -41,6 +42,15 @@ const (
 	ProviderStateCaptureState
 )
 
+type eventInfoClass uint32
+
+const (
+	eventInfoClassProviderBinaryTrackInfo eventInfoClass = iota
+	eventInfoClassProviderSetReserved1
+	eventInfoClassProviderSetTraits
+	eventInfoClassProviderUseDescriptorType
+)
+
 // EnableCallback is the form of the callback function that receives provider
 // enable/disable notifications from ETW.
 type EnableCallback func(*windows.GUID, ProviderState, Level, uint64, uint64, uintptr)
@@ -133,6 +143,15 @@ func NewProviderWithID(name string, id *windows.GUID, callback EnableCallback) (
 	binary.LittleEndian.PutUint16(metadata.Bytes(), uint16(metadata.Len())) // Update the size at the beginning of the buffer
 	provider.metadata = metadata.Bytes()
 
+	if err := eventSetInformation(
+		provider.handle,
+		eventInfoClassProviderSetTraits,
+		uintptr(unsafe.Pointer(&provider.metadata[0])),
+		uint32(len(provider.metadata))); err != nil {
+
+		return nil, err
+	}
+
 	return provider, nil
 }
 
diff --git a/internal/etw/zsyscall_windows.go b/internal/etw/zsyscall_windows.go
index 5b044ca4..489a0f99 100644
--- a/internal/etw/zsyscall_windows.go
+++ b/internal/etw/zsyscall_windows.go
@@ -39,9 +39,10 @@ func errnoErr(e syscall.Errno) error {
 var (
 	modadvapi32 = windows.NewLazySystemDLL("advapi32.dll")
 
-	procEventRegister      = modadvapi32.NewProc("EventRegister")
-	procEventUnregister    = modadvapi32.NewProc("EventUnregister")
-	procEventWriteTransfer = modadvapi32.NewProc("EventWriteTransfer")
+	procEventRegister       = modadvapi32.NewProc("EventRegister")
+	procEventUnregister     = modadvapi32.NewProc("EventUnregister")
+	procEventWriteTransfer  = modadvapi32.NewProc("EventWriteTransfer")
+	procEventSetInformation = modadvapi32.NewProc("EventSetInformation")
 )
 
 func eventRegister(providerId *windows.GUID, callback uintptr, callbackContext uintptr, providerHandle *providerHandle) (win32err error) {
@@ -67,3 +68,11 @@ func eventWriteTransfer(providerHandle providerHandle, descriptor *EventDescript
 	}
 	return
 }
+
+func eventSetInformation(providerHandle providerHandle, class eventInfoClass, information uintptr, length uint32) (win32err error) {
+	r0, _, _ := syscall.Syscall6(procEventSetInformation.Addr(), 4, uintptr(providerHandle), uintptr(class), uintptr(information), uintptr(length), 0, 0)
+	if r0 != 0 {
+		win32err = syscall.Errno(r0)
+	}
+	return
+}