diff --git a/go.mod b/go.mod index b384682..50b9d6e 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index babb4a7..209aa8c 100644 --- a/go.sum +++ b/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= diff --git a/pkg/process/process.go b/pkg/process/process.go new file mode 100644 index 0000000..ab828f6 --- /dev/null +++ b/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 + } + +} diff --git a/pkg/process/syscall.go b/pkg/process/syscall.go new file mode 100644 index 0000000..6732e5e --- /dev/null +++ b/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 diff --git a/pkg/process/zsyscall_windows.go b/pkg/process/zsyscall_windows.go new file mode 100644 index 0000000..a1f1a10 --- /dev/null +++ b/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 +}