Add new process package

Signed-off-by: Kevin Parsons <kevpar@microsoft.com>

go.mod

@@ -5,5 +5,5 @@ go 1.12
 require (
 	github.com/pkg/errors v0.8.1
 	github.com/sirupsen/logrus v1.4.1
-	golang.org/x/sys v0.0.0-20190507160741-ecd444e8653b
+	golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3
 )

go.sum

@@ -14,3 +14,5 @@ github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXf
 golang.org/x/sys v0.0.0-20180905080454-ebe1bf3edb33/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190507160741-ecd444e8653b h1:ag/x1USPSsqHud38I9BAC88qdNLDHHtQ4mlgQIZPPNA=
 golang.org/x/sys v0.0.0-20190507160741-ecd444e8653b/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3 h1:7TYNF4UdlohbFwpNH04CoPMp1cHUZgO1Ebq5r2hIjfo=
+golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

pkg/process/process.go

@@ -0,0 +1,86 @@
+package process
+
+import (
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+// EnumProcesses returns a slice containing the process IDs of all processes
+// currently running on the system.
+func EnumProcesses() ([]uint32, error) {
+	count := 256
+	uint32Size := unsafe.Sizeof(uint32(0))
+	for {
+		buf := make([]uint32, count)
+		bufferSize := uint32(len(buf) * int(uint32Size))
+		retBufferSize := uint32(0)
+		if err := enumProcesses(&buf[0], bufferSize, &retBufferSize); err != nil {
+			return nil, err
+		}
+		if retBufferSize == bufferSize {
+			count = count * 2
+			continue
+		}
+		actualCount := retBufferSize / uint32(uint32Size)
+		return buf[:actualCount], nil
+	}
+}
+
+// ProcessMemoryCountersEx is the PROCESS_MEMORY_COUNTERS_EX struct from
+// Windows:
+// https://docs.microsoft.com/en-us/windows/win32/api/psapi/ns-psapi-process_memory_counters_ex
+type ProcessMemoryCountersEx struct {
+	Cb                         uint32
+	PageFaultCount             uint32
+	PeakWorkingSetSize         uint
+	WorkingSetSize             uint
+	QuotaPeakPagedPoolUsage    uint
+	QuotaPagedPoolUsage        uint
+	QuotaPeakNonPagedPoolUsage uint
+	QuotaNonPagedPoolUsage     uint
+	PagefileUsage              uint
+	PeakPagefileUsage          uint
+	PrivateUsage               uint
+}
+
+// GetProcessMemoryInfo returns the memory usage information for the given
+// process. The process handle must have the PROCESS_QUERY_INFORMATION or
+// PROCESS_QUERY_LIMITED_INFORMATION, and the PROCESS_VM_READ access rights.
+func GetProcessMemoryInfo(process windows.Handle) (*ProcessMemoryCountersEx, error) {
+	memCounters := &ProcessMemoryCountersEx{}
+	size := unsafe.Sizeof(*memCounters)
+	if err := getProcessMemoryInfo(process, memCounters, uint32(size)); err != nil {
+		return nil, err
+	}
+	return memCounters, nil
+}
+
+// These constants are used with QueryFullProcessImageName's flags.
+const (
+	// ImageNameFormatWin32Path indicates to format the name as a Win32 path.
+	ImageNameFormatWin32Path = iota
+	// ImageNameFormatNTPath indicates to format the name as a NT path.
+	ImageNameFormatNTPath
+)
+
+// QueryFullProcessImageName returns the full process image name for the given
+// process. The process handle must have the PROCESS_QUERY_INFORMATION or
+// PROCESS_QUERY_LIMITED_INFORMATION access right. The flags can be either
+// `ImageNameFormatWin32Path` or `ImageNameFormatNTPath`.
+func QueryFullProcessImageName(process windows.Handle, flags uint32) (string, error) {
+	bufferSize := uint32(256)
+	for {
+		b := make([]uint16, bufferSize)
+		err := queryFullProcessImageName(process, flags, &b[0], &bufferSize)
+		if err == windows.ERROR_INSUFFICIENT_BUFFER {
+			bufferSize = bufferSize * 2
+			continue
+		}
+		if err != nil {
+			return "", err
+		}
+		return windows.UTF16ToString(b[:bufferSize]), nil
+	}
+
+}

pkg/process/syscall.go

@@ -0,0 +1,13 @@
+package process
+
+import (
+	"golang.org/x/sys/windows"
+)
+
+//go:generate go run $GOROOT/src/syscall/mksyscall_windows.go -output zsyscall_windows.go syscall.go
+
+//sys enumProcesses(pids *uint32, bufferSize uint32, retBufferSize *uint32) (err error) = kernel32.K32EnumProcesses
+//sys getProcessMemoryInfo(process handle, memCounters *ProcessMemoryCountersEx, size uint32) (err error) = kernel32.K32GetProcessMemoryInfo
+//sys queryFullProcessImageName(process handle, flags uint32, buffer *uint16, bufferSize *uint32) (err error) = kernel32.QueryFullProcessImageNameW
+
+type handle = windows.Handle

pkg/process/zsyscall_windows.go

@@ -0,0 +1,81 @@
+// Code generated by 'go generate'; DO NOT EDIT.
+
+package process
+
+import (
+	"syscall"
+	"unsafe"
+
+	"golang.org/x/sys/windows"
+)
+
+var _ unsafe.Pointer
+
+// Do the interface allocations only once for common
+// Errno values.
+const (
+	errnoERROR_IO_PENDING = 997
+)
+
+var (
+	errERROR_IO_PENDING error = syscall.Errno(errnoERROR_IO_PENDING)
+)
+
+// errnoErr returns common boxed Errno values, to prevent
+// allocations at runtime.
+func errnoErr(e syscall.Errno) error {
+	switch e {
+	case 0:
+		return nil
+	case errnoERROR_IO_PENDING:
+		return errERROR_IO_PENDING
+	}
+	// TODO: add more here, after collecting data on the common
+	// error values see on Windows. (perhaps when running
+	// all.bat?)
+	return e
+}
+
+var (
+	modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
+
+	procK32EnumProcesses           = modkernel32.NewProc("K32EnumProcesses")
+	procK32GetProcessMemoryInfo    = modkernel32.NewProc("K32GetProcessMemoryInfo")
+	procQueryFullProcessImageNameW = modkernel32.NewProc("QueryFullProcessImageNameW")
+)
+
+func enumProcesses(pids *uint32, bufferSize uint32, retBufferSize *uint32) (err error) {
+	r1, _, e1 := syscall.Syscall(procK32EnumProcesses.Addr(), 3, uintptr(unsafe.Pointer(pids)), uintptr(bufferSize), uintptr(unsafe.Pointer(retBufferSize)))
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func getProcessMemoryInfo(process handle, memCounters *ProcessMemoryCountersEx, size uint32) (err error) {
+	r1, _, e1 := syscall.Syscall(procK32GetProcessMemoryInfo.Addr(), 3, uintptr(process), uintptr(unsafe.Pointer(memCounters)), uintptr(size))
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}
+
+func queryFullProcessImageName(process handle, flags uint32, buffer *uint16, bufferSize *uint32) (err error) {
+	r1, _, e1 := syscall.Syscall6(procQueryFullProcessImageNameW.Addr(), 4, uintptr(process), uintptr(flags), uintptr(unsafe.Pointer(buffer)), uintptr(unsafe.Pointer(bufferSize)), 0, 0)
+	if r1 == 0 {
+		if e1 != 0 {
+			err = errnoErr(e1)
+		} else {
+			err = syscall.EINVAL
+		}
+	}
+	return
+}