Skip to content

Commit

Permalink
Fixed slow command line retrieval on Windows
Browse files Browse the repository at this point in the history
  • Loading branch information
@mxmauro
mxmauro committed May 1, 2020
1 parent 93a90cc commit 00957cf
Show file tree
Hide file tree
Showing 5 changed files with 328 additions and 90 deletions.
30 changes: 21 additions & 9 deletions internal/common/common_windows.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,21 +49,33 @@ const (
PDH_NO_DATA = 0x800007d5
)

const (
ProcessBasicInformation = 0
ProcessWow64Information = 26
)

var (
Modkernel32 = windows.NewLazySystemDLL("kernel32.dll")
ModNt = windows.NewLazySystemDLL("ntdll.dll")
ModPdh = windows.NewLazySystemDLL("pdh.dll")
ModPsapi = windows.NewLazySystemDLL("psapi.dll")

ProcGetSystemTimes = Modkernel32.NewProc("GetSystemTimes")
ProcNtQuerySystemInformation = ModNt.NewProc("NtQuerySystemInformation")
PdhOpenQuery = ModPdh.NewProc("PdhOpenQuery")
PdhAddCounter = ModPdh.NewProc("PdhAddCounterW")
PdhCollectQueryData = ModPdh.NewProc("PdhCollectQueryData")
PdhGetFormattedCounterValue = ModPdh.NewProc("PdhGetFormattedCounterValue")
PdhCloseQuery = ModPdh.NewProc("PdhCloseQuery")

procQueryDosDeviceW = Modkernel32.NewProc("QueryDosDeviceW")
ProcGetSystemTimes = Modkernel32.NewProc("GetSystemTimes")
ProcNtQuerySystemInformation = ModNt.NewProc("NtQuerySystemInformation")
ProcRtlGetNativeSystemInformation = ModNt.NewProc("RtlGetNativeSystemInformation")
ProcRtlNtStatusToDosError = ModNt.NewProc("RtlNtStatusToDosError")
ProcNtQueryInformationProcess = ModNt.NewProc("NtQueryInformationProcess")
ProcNtReadVirtualMemory = ModNt.NewProc("NtReadVirtualMemory")
ProcNtWow64QueryInformationProcess64 = ModNt.NewProc("NtWow64QueryInformationProcess64")
ProcNtWow64ReadVirtualMemory64 = ModNt.NewProc("NtWow64ReadVirtualMemory64")

PdhOpenQuery = ModPdh.NewProc("PdhOpenQuery")
PdhAddCounter = ModPdh.NewProc("PdhAddCounterW")
PdhCollectQueryData = ModPdh.NewProc("PdhCollectQueryData")
PdhGetFormattedCounterValue = ModPdh.NewProc("PdhGetFormattedCounterValue")
PdhCloseQuery = ModPdh.NewProc("PdhCloseQuery")

procQueryDosDeviceW = Modkernel32.NewProc("QueryDosDeviceW")
)

type FILETIME struct {
Expand Down
18 changes: 18 additions & 0 deletions process/process_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -623,3 +623,21 @@ func Test_IsRunning(t *testing.T) {
t.Fatalf("process should NOT be found running")
}
}

func Test_AllProcesses_cmdLine(t *testing.T) {
procs, err := Processes()
if err == nil {
for _, proc := range procs {
var exeName string
var cmdLine string

exeName, _ = proc.Exe()
cmdLine, err = proc.Cmdline()
if err != nil {
cmdLine = "Error: " + err.Error()
}

t.Logf("Process #%v: Name: %v / CmdLine: %v\n", proc.Pid, exeName, cmdLine)
}
}
}
226 changes: 145 additions & 81 deletions process/process_windows.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,13 @@ package process

import (
"context"
"errors"
"fmt"
"os"
"strings"
"syscall"
"time"
"unsafe"

"github.com/StackExchange/wmi"
cpu "github.com/shirou/gopsutil/cpu"
"github.com/shirou/gopsutil/internal/common"
net "github.com/shirou/gopsutil/net"
Expand Down Expand Up @@ -47,6 +46,14 @@ type SystemProcessInformation struct {
Reserved6 [6]uint64
}

type SystemProcessorInformation struct {
ProcessorArchitecture uint16
ProcessorLevel uint16
ProcessorRevision uint16
Reserved uint16
ProcessorFeatureBits uint16
}

// Memory_info_ex is different between OSes
type MemoryInfoExStat struct {
}
Expand All @@ -65,43 +72,22 @@ type ioCounters struct {
OtherTransferCount uint64
}

type Win32_Process struct {
Name string
ExecutablePath *string
CommandLine *string
Priority uint32
CreationDate *time.Time
ProcessID uint32
ThreadCount uint32
Status *string
ReadOperationCount uint64
ReadTransferCount uint64
WriteOperationCount uint64
WriteTransferCount uint64
CSCreationClassName string
CSName string
Caption *string
CreationClassName string
Description *string
ExecutionState *uint16
HandleCount uint32
KernelModeTime uint64
MaximumWorkingSetSize *uint32
MinimumWorkingSetSize *uint32
OSCreationClassName string
OSName string
OtherOperationCount uint64
OtherTransferCount uint64
PageFaults uint32
PageFileUsage uint32
ParentProcessID uint32
PeakPageFileUsage uint32
PeakVirtualSize uint64
PeakWorkingSetSize uint32
PrivatePageCount uint64
TerminationDate *time.Time
UserModeTime uint64
WorkingSetSize uint64
type PROCESS_BASIC_INFORMATION32 struct {
Reserved1 uint32
PebBaseAddress uint32
Reserved2 uint32
Reserved3 uint32
UniqueProcessId uint32
Reserved4 uint32
}

type PROCESS_BASIC_INFORMATION64 struct {
Reserved1 uint64
PebBaseAddress uint64
Reserved2 uint64
Reserved3 uint64
UniqueProcessId uint64
Reserved4 uint64
}

type winLUID struct {
Expand All @@ -125,8 +111,6 @@ type winLong int32
type winDWord uint32

func init() {
wmi.DefaultClient.AllowMissingFields = true

// enable SeDebugPrivilege https://github.com/midstar/proci/blob/6ec79f57b90ba3d9efa2a7b16ef9c9369d4be875/proci_windows.go#L80-L119
handle, err := syscall.GetCurrentProcess()
if err != nil {
Expand Down Expand Up @@ -237,26 +221,6 @@ func (p *Process) PpidWithContext(ctx context.Context) (int32, error) {
return ppid, nil
}

func GetWin32Proc(pid int32) ([]Win32_Process, error) {
return GetWin32ProcWithContext(context.Background(), pid)
}

func GetWin32ProcWithContext(ctx context.Context, pid int32) ([]Win32_Process, error) {
var dst []Win32_Process
query := fmt.Sprintf("WHERE ProcessId = %d", pid)
q := wmi.CreateQuery(&dst, query)
err := common.WMIQueryWithContext(ctx, q, &dst)
if err != nil {
return []Win32_Process{}, fmt.Errorf("could not get win32Proc: %s", err)
}

if len(dst) == 0 {
return []Win32_Process{}, fmt.Errorf("could not get win32Proc: empty")
}

return dst, nil
}

func (p *Process) Name() (string, error) {
return p.NameWithContext(context.Background())
}
Expand Down Expand Up @@ -309,11 +273,11 @@ func (p *Process) Cmdline() (string, error) {
}

func (p *Process) CmdlineWithContext(ctx context.Context) (string, error) {
dst, err := GetWin32ProcWithContext(ctx, p.Pid)
cmdline, err := getProcessCommandLine(ctx, p.Pid)
if err != nil {
return "", fmt.Errorf("could not get CommandLine: %s", err)
}
return *dst[0].CommandLine, nil
return cmdline, nil
}

// CmdlineSlice returns the command line arguments of the process as a slice with each
Expand Down Expand Up @@ -774,24 +738,6 @@ func ProcessesWithContext(ctx context.Context) ([]*Process, error) {
return out, nil
}

func getProcInfo(pid int32) (*SystemProcessInformation, error) {
initialBufferSize := uint64(0x4000)
bufferSize := initialBufferSize
buffer := make([]byte, bufferSize)

var sysProcInfo SystemProcessInformation
ret, _, _ := common.ProcNtQuerySystemInformation.Call(
uintptr(unsafe.Pointer(&sysProcInfo)),
uintptr(unsafe.Pointer(&buffer[0])),
uintptr(unsafe.Pointer(&bufferSize)),
uintptr(unsafe.Pointer(&bufferSize)))
if ret != 0 {
return nil, windows.GetLastError()
}

return &sysProcInfo, nil
}

func getRusage(pid int32) (*windows.Rusage, error) {
var CPU windows.Rusage

Expand Down Expand Up @@ -860,3 +806,121 @@ func getProcessCPUTimes(pid int32) (SYSTEM_TIMES, error) {

return times, err
}

func is32BitProcess(procHandle syscall.Handle) bool {
var wow64 uint

ret, _, _ := common.ProcNtQueryInformationProcess.Call(
uintptr(procHandle),
uintptr(common.ProcessWow64Information),
uintptr(unsafe.Pointer(&wow64)),
uintptr(unsafe.Sizeof(wow64)),
uintptr(0),
)
if int(ret) >= 0 {
if wow64 != 0 {
return true
}
} else {
//if the OS does not support the call, we fallback into the bitness of the app
if unsafe.Sizeof(wow64) == 4 {
return true
}
}
return false
}

func getProcessCommandLine(_ context.Context, pid int32) (string, error) {
h, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION | windows.PROCESS_VM_READ, false, uint32(pid))
if err == windows.ERROR_ACCESS_DENIED || err == windows.ERROR_INVALID_PARAMETER {
return "", nil
}
if err != nil {
return "", err
}
defer syscall.CloseHandle(syscall.Handle(h))

procIs32Bits := is32BitProcess(syscall.Handle(h))

pebAddress := queryPebAddress(syscall.Handle(h), procIs32Bits)
if pebAddress == 0 {
return "", errors.New("cannot locate process PEB")
}

if is32BitProcess(syscall.Handle(h)) {
buf := readProcessMemory(syscall.Handle(h), procIs32Bits, pebAddress + uint64(16), 4)
if len(buf) != 4 {
return "", errors.New("cannot locate process user parameters")
}
userProcParams := uint64(buf[0]) | (uint64(buf[1]) << 8) | (uint64(buf[2]) << 16) | (uint64(buf[3]) << 24)

//read CommandLine field from PRTL_USER_PROCESS_PARAMETERS
remoteCmdLine := readProcessMemory(syscall.Handle(h), procIs32Bits, userProcParams + uint64(64), 8)
if len(remoteCmdLine) != 8 {
return "", errors.New("cannot read cmdline field")
}

//remoteCmdLine is actually a UNICODE_STRING32
//the first two bytes has the length
cmdLineLength := uint(remoteCmdLine[0]) | (uint(remoteCmdLine[1]) << 8)
if cmdLineLength > 0 {
//and, at offset 4, is the pointer to the buffer
bufferAddress := uint32(remoteCmdLine[4]) | (uint32(remoteCmdLine[5]) << 8) |
(uint32(remoteCmdLine[6]) << 16) | (uint32(remoteCmdLine[7]) << 24)

cmdLine := readProcessMemory(syscall.Handle(h), procIs32Bits, uint64(bufferAddress), cmdLineLength)
if len(cmdLine) != int(cmdLineLength) {
return "", errors.New("cannot read cmdline")
}

return convertUTF16ToString(cmdLine), nil
}
} else {
buf := readProcessMemory(syscall.Handle(h), procIs32Bits, pebAddress + uint64(32), 8)
if len(buf) != 8 {
return "", errors.New("cannot locate process user parameters")
}
userProcParams := uint64(buf[0]) | (uint64(buf[1]) << 8) | (uint64(buf[2]) << 16) | (uint64(buf[3]) << 24) |
(uint64(buf[4]) << 32) | (uint64(buf[5]) << 40) | (uint64(buf[6]) << 48) | (uint64(buf[7]) << 56)

//read CommandLine field from PRTL_USER_PROCESS_PARAMETERS
remoteCmdLine := readProcessMemory(syscall.Handle(h), procIs32Bits, userProcParams + uint64(112), 16)
if len(remoteCmdLine) != 16 {
return "", errors.New("cannot read cmdline field")
}

//remoteCmdLine is actually a UNICODE_STRING64
//the first two bytes has the length
cmdLineLength := uint(remoteCmdLine[0]) | (uint(remoteCmdLine[1]) << 8)
if cmdLineLength > 0 {
//and, at offset 8, is the pointer to the buffer
bufferAddress := uint64(remoteCmdLine[8]) | (uint64(remoteCmdLine[9]) << 8) |
(uint64(remoteCmdLine[10]) << 16) | (uint64(remoteCmdLine[11]) << 24) |
(uint64(remoteCmdLine[12]) << 32) | (uint64(remoteCmdLine[13]) << 40) |
(uint64(remoteCmdLine[14]) << 48) | (uint64(remoteCmdLine[15]) << 56)

cmdLine := readProcessMemory(syscall.Handle(h), procIs32Bits, bufferAddress, cmdLineLength)
if len(cmdLine) != int(cmdLineLength) {
return "", errors.New("cannot read cmdline")
}

return convertUTF16ToString(cmdLine), nil
}
}

//if we reach here, we have no command line
return "", nil
}

func convertUTF16ToString(src []byte) string {
srcLen := len(src) / 2

codePoints := make([]uint16, srcLen)

srcIdx := 0
for i := 0; i < srcLen; i++ {
codePoints[i] = uint16(src[srcIdx]) | uint16(src[srcIdx + 1] << 8)
srcIdx += 2
}
return syscall.UTF16ToString(codePoints)
}
Loading

0 comments on commit 00957cf

Please sign in to comment.