Replace wmi queries with win32 api calls

sigar_windows.go

@@ -12,7 +12,6 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/StackExchange/wmi"
 	"github.com/elastic/gosigar/sys/windows"
 	"github.com/pkg/errors"
 )
@@ -83,11 +82,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")
 		}
-		bootTime = &os.LastBootUpTime
+		var boot = time.Unix(int64(uptime), 0)
+		bootTime = &boot
 	}
 
 	self.Length = time.Since(*bootTime).Seconds()
@@ -251,7 +251,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 +371,33 @@ 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 err != nil {
+		return nil
+	}
+	userProcParams, err := windows.GetUserProcessParams(handle, pbi)
+	if err != nil {
+		return nil
+	}
+	var args []string
+	if argsW, err := windows.ReadProcessUnicodeString(handle, &userProcParams.CommandLine); err == nil {
+		args, err = windows.ByteSliceToStringSlice(argsW)
+		if err != nil {
+			args = nil
+		}
 	}
+	var process = Win32_Process{CommandLine: &args[0]}
 	if process.CommandLine != nil {
 		self.List = []string{*process.CommandLine}
 	}
-
 	return nil
 }
 
@@ -396,35 +414,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

@@ -22,6 +22,11 @@ const (
 	PROCESS_QUERY_LIMITED_INFORMATION uint32 = 0x1000
 	PROCESS_VM_READ                   uint32 = 0x0010
 )
+const (
+	// SizeOfRtlUserProcessParameters gives the size
+	// of the RtlUserProcessParameters struct.
+	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
@@ -43,6 +48,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 +466,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 +581,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

sys/windows/syscall_windows_test.go

@@ -1,10 +1,12 @@
 package windows
 
 import (
+	"encoding/binary"
 	"os"
 	"runtime"
 	"syscall"
 	"testing"
+	"unsafe"
 
 	"github.com/pkg/errors"
 	"github.com/stretchr/testify/assert"
@@ -181,3 +183,139 @@ func TestNtQueryProcessBasicInformation(t *testing.T) {
 
 	t.Logf("NtQueryProcessBasicInformation: %+v", info)
 }
+
+func TestGetTickCount64(t *testing.T) {
+	uptime, err := GetTickCount64()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NotZero(t, uptime)
+}
+
+func TestGetUserProcessParams(t *testing.T) {
+	h, err := syscall.OpenProcess(syscall.PROCESS_QUERY_INFORMATION|PROCESS_VM_READ, false, uint32(syscall.Getpid()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	info, err := NtQueryProcessBasicInformation(h)
+	if err != nil {
+		t.Fatal(err)
+	}
+	userProc, err := GetUserProcessParams(h, info)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NoError(t, err)
+	assert.EqualValues(t, os.Getpid(), info.UniqueProcessID)
+	assert.EqualValues(t, os.Getppid(), info.InheritedFromUniqueProcessID)
+	assert.NotEmpty(t, userProc.CommandLine)
+}
+func TestReadProcessUnicodeString(t *testing.T) {
+	h, err := syscall.OpenProcess(syscall.PROCESS_QUERY_INFORMATION|PROCESS_VM_READ, false, uint32(syscall.Getpid()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	info, err := NtQueryProcessBasicInformation(h)
+	if err != nil {
+		t.Fatal(err)
+	}
+	userProc, err := GetUserProcessParams(h, info)
+	if err != nil {
+		t.Fatal(err)
+	}
+	read, err := ReadProcessUnicodeString(h, &userProc.CommandLine)
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NoError(t, err)
+	assert.NotEmpty(t, read)
+}
+
+func TestByteSliceToStringSlice(t *testing.T) {
+	cmd := syscall.GetCommandLine()
+	b := make([]byte, unsafe.Sizeof(cmd))
+	binary.LittleEndian.PutUint16(b, *cmd)
+	hes, err := ByteSliceToStringSlice(b)
+	assert.NoError(t, err)
+	assert.NotEmpty(t, hes)
+}
+
+func TestByteSliceToStringSliceEmptyBytes(t *testing.T) {
+	b := make([]byte, 0)
+	hes, err := ByteSliceToStringSlice(b)
+	assert.NoError(t, err)
+	assert.Empty(t, hes)
+}
+
+func TestReadProcessMemory(t *testing.T) {
+	h, err := syscall.OpenProcess(syscall.PROCESS_QUERY_INFORMATION|PROCESS_VM_READ, false, uint32(syscall.Getpid()))
+	if err != nil {
+		t.Fatal(err)
+	}
+	info, err := NtQueryProcessBasicInformation(h)
+	if err != nil {
+		t.Fatal(err)
+	}
+	pebSize := 0x20 + 8
+	if unsafe.Sizeof(uintptr(0)) == 4 {
+		pebSize = 0x10 + 8
+	}
+	peb := make([]byte, pebSize)
+	nRead, err := ReadProcessMemory(h, info.PebBaseAddress, peb)
+	assert.NoError(t, err)
+	assert.NotEmpty(t, nRead)
+	assert.EqualValues(t, nRead, uintptr(pebSize))
+}
+
+// A zero-byte read is a no-op, no error is returned.
+func TestReadProcessMemoryZeroByteRead(t *testing.T) {
+	peb := make([]byte, 0)
+	var h syscall.Handle
+	var address uintptr
+	nRead, err := ReadProcessMemory(h, address, peb)
+	assert.NoError(t, err)
+	assert.Empty(t, nRead)
+}
+
+func TestReadProcessMemoryInvalidHandler(t *testing.T) {
+	peb := make([]byte, 10)
+	var h syscall.Handle
+	var address uintptr
+	nRead, err := ReadProcessMemory(h, address, peb)
+	assert.Error(t, err)
+	assert.EqualValues(t, err.Error(), "The handle is invalid.")
+	assert.Empty(t, nRead)
+}
+
+func TestGetAccessPaths(t *testing.T) {
+	paths, err := GetAccessPaths()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NotEmpty(t, paths)
+	assert.True(t, len(paths) >= 1)
+}
+
+func TestGetVolumes(t *testing.T) {
+	paths, err := GetVolumes()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NotEmpty(t, paths)
+	assert.True(t, len(paths) >= 1)
+}
+
+func TestGetVolumePathsForVolume(t *testing.T) {
+	volumes, err := GetVolumes()
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NotNil(t, volumes)
+	assert.True(t, len(volumes) >= 1)
+	volumePath, err := GetVolumePathsForVolume(volumes[0])
+	if err != nil {
+		t.Fatal(err)
+	}
+	assert.NotNil(t, volumePath)
+	assert.True(t, len(volumePath) >= 1)
+}