Replace wmi queries with win32 api calls (#116)

* Remove wmi queries with win32 api calls

Fixes elastic/beats#11840

Addresses high CPU load on windows

* Add more tests

* Add more tests for coverage

* Add tests

* Add tests for coverage

* Remove WIn32Process related structs

* Fix test failing on osx

* Work on tests

* Clean up tests

* Fix on changelog file

CHANGELOG.md

@@ -7,6 +7,7 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Added
 
 ### Fixed
+- Replaced the WMI queries with win32 apis due to high CPU usage. #11840
 
 ### Changed
 

sigar_interface_test.go

@@ -100,7 +100,6 @@ func TestProcState(t *testing.T) {
 		assert.Equal(t, u.Username, state.Username)
 		assert.True(t, state.Ppid > 0, "ppid=%v is non-positive", state.Ppid)
 	}
-
 	assert.Error(t, state.Get(invalidPid))
 }
 
@@ -123,9 +122,12 @@ func TestProcTime(t *testing.T) {
 }
 
 func TestProcArgs(t *testing.T) {
-	args := ProcArgs{}
-	if assert.NoError(t, args.Get(os.Getppid())) {
-		assert.NotEmpty(t, args.List)
+	procArgs := ProcArgs{}
+	if assert.NoError(t, procArgs.Get(os.Getppid())) {
+		assert.NotEmpty(t, procArgs.List)
+	}
+	if runtime.GOOS != "darwin" {
+		assert.Error(t, procArgs.Get(invalidPid))
 	}
 }
 

sigar_windows.go

@@ -12,26 +12,10 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/StackExchange/wmi"
 	"github.com/elastic/gosigar/sys/windows"
 	"github.com/pkg/errors"
 )
 
-// Win32_Process represents a process on the Windows operating system. If
-// additional fields are added here (that match the Windows struct) they will
-// automatically be populated when calling getWin32Process.
-// https://msdn.microsoft.com/en-us/library/windows/desktop/aa394372(v=vs.85).aspx
-type Win32_Process struct {
-	CommandLine *string
-}
-
-// Win32_OperatingSystem WMI class represents a Windows-based operating system
-// installed on a computer.
-// https://msdn.microsoft.com/en-us/library/windows/desktop/aa394239(v=vs.85).aspx
-type Win32_OperatingSystem struct {
-	LastBootUpTime time.Time
-}
-
 var (
 	// version is Windows version of the host OS.
 	version = windows.GetWindowsVersion()
@@ -83,11 +67,12 @@ func (self *Uptime) Get() error {
 	bootTimeLock.Lock()
 	defer bootTimeLock.Unlock()
 	if bootTime == nil {
-		os, err := getWin32OperatingSystem()
+		uptime, err := windows.GetTickCount64()
 		if err != nil {
-			return errors.Wrap(err, "failed to get boot time using WMI")
+			return errors.Wrap(err, "failed to get boot time using win32 api")
 		}
-		bootTime = &os.LastBootUpTime
+		var boot = time.Unix(int64(uptime), 0)
+		bootTime = &boot
 	}
 
 	self.Length = time.Since(*bootTime).Seconds()
@@ -251,7 +236,7 @@ func getProcStatus(pid int) (RunState, error) {
 	var exitCode uint32
 	err = syscall.GetExitCodeProcess(handle, &exitCode)
 	if err != nil {
-		return RunStateUnknown, errors.Wrapf(err, "GetExitCodeProcess failed for pid=%v")
+		return RunStateUnknown, errors.Wrapf(err, "GetExitCodeProcess failed for pid=%v", pid)
 	}
 
 	if exitCode == 259 { //still active
@@ -371,15 +356,28 @@ func (self *ProcArgs) Get(pid int) error {
 	if !version.IsWindowsVistaOrGreater() {
 		return ErrNotImplemented{runtime.GOOS}
 	}
-
-	process, err := getWin32Process(int32(pid))
+	handle, err := syscall.OpenProcess(processQueryLimitedInfoAccess|windows.PROCESS_VM_READ, false, uint32(pid))
+	if err != nil {
+		return errors.Wrapf(err, "OpenProcess failed for pid=%v", pid)
+	}
+	defer syscall.CloseHandle(handle)
+	pbi, err := windows.NtQueryProcessBasicInformation(handle)
 	if err != nil {
-		return errors.Wrapf(err, "ProcArgs failed for pid=%v", pid)
+		return errors.Wrapf(err, "NtQueryProcessBasicInformation failed for pid=%v", pid)
 	}
-	if process.CommandLine != nil {
-		self.List = []string{*process.CommandLine}
+	if err != nil {
+		return nil
+	}
+	userProcParams, err := windows.GetUserProcessParams(handle, pbi)
+	if err != nil {
+		return nil
+	}
+	if argsW, err := windows.ReadProcessUnicodeString(handle, &userProcParams.CommandLine); err == nil {
+		self.List, err = windows.ByteSliceToStringSlice(argsW)
+		if err != nil {
+			return err
+		}
 	}
-
 	return nil
 }
 
@@ -396,35 +394,6 @@ func (self *FileSystemUsage) Get(path string) error {
 	return nil
 }
 
-// getWin32Process gets information about the process with the given process ID.
-// It uses a WMI query to get the information from the local system.
-func getWin32Process(pid int32) (Win32_Process, error) {
-	var dst []Win32_Process
-	query := fmt.Sprintf("WHERE ProcessId = %d", pid)
-	q := wmi.CreateQuery(&dst, query)
-	err := wmi.Query(q, &dst)
-	if err != nil {
-		return Win32_Process{}, fmt.Errorf("could not get Win32_Process %s: %v", query, err)
-	}
-	if len(dst) < 1 {
-		return Win32_Process{}, fmt.Errorf("could not get Win32_Process %s: Process not found", query)
-	}
-	return dst[0], nil
-}
-
-func getWin32OperatingSystem() (Win32_OperatingSystem, error) {
-	var dst []Win32_OperatingSystem
-	q := wmi.CreateQuery(&dst, "")
-	err := wmi.Query(q, &dst)
-	if err != nil {
-		return Win32_OperatingSystem{}, errors.Wrap(err, "wmi query for Win32_OperatingSystem failed")
-	}
-	if len(dst) != 1 {
-		return Win32_OperatingSystem{}, errors.New("wmi query for Win32_OperatingSystem failed")
-	}
-	return dst[0], nil
-}
-
 func (self *Rusage) Get(who int) error {
 	if who != 0 {
 		return ErrNotImplemented{runtime.GOOS}

sys/windows/syscall_windows.go

@@ -23,6 +23,10 @@ const (
 	PROCESS_VM_READ                   uint32 = 0x0010
 )
 
+// SizeOfRtlUserProcessParameters gives the size
+// of the RtlUserProcessParameters struct.
+const SizeOfRtlUserProcessParameters = unsafe.Sizeof(RtlUserProcessParameters{})
+
 // MAX_PATH is the maximum length for a path in Windows.
 // https://msdn.microsoft.com/en-us/library/windows/desktop/aa365247(v=vs.85).aspx
 const MAX_PATH = 260
@@ -43,6 +47,26 @@ const (
 	DRIVE_RAMDISK
 )
 
+// UnicodeString is Go's equivalent for the _UNICODE_STRING struct.
+type UnicodeString struct {
+	Size          uint16
+	MaximumLength uint16
+	Buffer        uintptr
+}
+
+// RtlUserProcessParameters is Go's equivalent for the
+// _RTL_USER_PROCESS_PARAMETERS struct.
+// A few undocumented fields are exposed.
+type RtlUserProcessParameters struct {
+	Reserved1              [16]byte
+	Reserved2              [5]uintptr
+	CurrentDirectoryPath   UnicodeString
+	CurrentDirectoryHandle uintptr
+	DllPath                UnicodeString
+	ImagePathName          UnicodeString
+	CommandLine            UnicodeString
+}
+
 func (dt DriveType) String() string {
 	names := map[DriveType]string{
 		DRIVE_UNKNOWN:     "unknown",
@@ -441,6 +465,95 @@ func UTF16SliceToStringSlice(buffer []uint16) []string {
 	return result
 }
 
+func GetUserProcessParams(handle syscall.Handle, pbi ProcessBasicInformation) (params RtlUserProcessParameters, err error) {
+	const is32bitProc = unsafe.Sizeof(uintptr(0)) == 4
+
+	// Offset of params field within PEB structure.
+	// This structure is different in 32 and 64 bit.
+	paramsOffset := 0x20
+	if is32bitProc {
+		paramsOffset = 0x10
+	}
+
+	// Read the PEB from the target process memory
+	pebSize := paramsOffset + 8
+	peb := make([]byte, pebSize)
+	nRead, err := ReadProcessMemory(handle, pbi.PebBaseAddress, peb)
+	if err != nil {
+		return params, err
+	}
+	if nRead != uintptr(pebSize) {
+		return params, errors.Errorf("PEB: short read (%d/%d)", nRead, pebSize)
+	}
+
+	// Get the RTL_USER_PROCESS_PARAMETERS struct pointer from the PEB
+	paramsAddr := *(*uintptr)(unsafe.Pointer(&peb[paramsOffset]))
+
+	// Read the RTL_USER_PROCESS_PARAMETERS from the target process memory
+	paramsBuf := make([]byte, SizeOfRtlUserProcessParameters)
+	nRead, err = ReadProcessMemory(handle, paramsAddr, paramsBuf)
+	if err != nil {
+		return params, err
+	}
+	if nRead != uintptr(SizeOfRtlUserProcessParameters) {
+		return params, errors.Errorf("RTL_USER_PROCESS_PARAMETERS: short read (%d/%d)", nRead, SizeOfRtlUserProcessParameters)
+	}
+
+	params = *(*RtlUserProcessParameters)(unsafe.Pointer(&paramsBuf[0]))
+	return params, nil
+}
+
+func ReadProcessUnicodeString(handle syscall.Handle, s *UnicodeString) ([]byte, error) {
+	buf := make([]byte, s.Size)
+	nRead, err := ReadProcessMemory(handle, s.Buffer, buf)
+	if err != nil {
+		return nil, err
+	}
+	if nRead != uintptr(s.Size) {
+		return nil, errors.Errorf("unicode string: short read: (%d/%d)", nRead, s.Size)
+	}
+	return buf, nil
+}
+
+// Use Windows' CommandLineToArgv API to split an UTF-16 command line string
+// into a list of parameters.
+func ByteSliceToStringSlice(utf16 []byte) ([]string, error) {
+	if len(utf16) == 0 {
+		return nil, nil
+	}
+	var numArgs int32
+	argsWide, err := syscall.CommandLineToArgv((*uint16)(unsafe.Pointer(&utf16[0])), &numArgs)
+	if err != nil {
+		return nil, err
+	}
+	args := make([]string, numArgs)
+	for idx := range args {
+		args[idx] = syscall.UTF16ToString(argsWide[idx][:])
+	}
+	return args, nil
+}
+
+// ReadProcessMemory reads from another process memory. The Handle needs to have
+// the PROCESS_VM_READ right.
+// A zero-byte read is a no-op, no error is returned.
+func ReadProcessMemory(handle syscall.Handle, baseAddress uintptr, dest []byte) (numRead uintptr, err error) {
+	n := len(dest)
+	if n == 0 {
+		return 0, nil
+	}
+	if err = _ReadProcessMemory(handle, baseAddress, uintptr(unsafe.Pointer(&dest[0])), uintptr(n), &numRead); err != nil {
+		return 0, err
+	}
+	return numRead, nil
+}
+
+func GetTickCount64() (uptime uint64, err error) {
+	if uptime, err = _GetTickCount64(); err != nil {
+		return 0, err
+	}
+	return uptime, nil
+}
+
 // Use "GOOS=windows go generate -v -x ." to generate the source.
 
 // Add -trace to enable debug prints around syscalls.
@@ -467,3 +580,5 @@ func UTF16SliceToStringSlice(buffer []uint16) []string {
 //sys  _FindNextVolume(handle syscall.Handle, volumeName *uint16, size uint32) (err error) = kernel32.FindNextVolumeW
 //sys  _FindVolumeClose(handle syscall.Handle) (err error) = kernel32.FindVolumeClose
 //sys  _GetVolumePathNamesForVolumeName(volumeName string, buffer *uint16, bufferSize uint32, length *uint32) (err error) = kernel32.GetVolumePathNamesForVolumeNameW
+//sys  _ReadProcessMemory(handle syscall.Handle, baseAddress uintptr, buffer uintptr, size uintptr, numRead *uintptr) (err error) = kernel32.ReadProcessMemory
+//sys  _GetTickCount64() (uptime uint64, err error) = kernel32.GetTickCount64