From 47323f9ad5780511d033d75bc55e894881d7bce2 Mon Sep 17 00:00:00 2001
From: Lomanic <Lomanic@users.noreply.github.com>
Date: Sun, 16 Jun 2019 22:59:03 +0200
Subject: [PATCH] [process][windows] Fix #586 use win32 API in process.Exe()
 instead of slow WMI call

This is faster by a factor of 100.

References:
https://github.com/giampaolo/psutil/blob/5f4287d17fc6aa2643c4c6e3589c12abd0f1ded9/psutil/_pswindows.py#L221
https://github.com/giampaolo/psutil/blob/921870d54091f399cd2b129db19530cc486b5700/psutil/_psutil_windows.c#L1211
https://github.com/giampaolo/psutil/blob/921870d54091f399cd2b129db19530cc486b5700/psutil/_psutil_windows.c#L626
---
 internal/common/common_windows.go | 25 ++++++++++++++++++++++
 process/process_windows.go        | 44 ++++++++++++++++++++++++---------------
 2 files changed, 52 insertions(+), 17 deletions(-)

diff --git a/internal/common/common_windows.go b/internal/common/common_windows.go
index 0997c9b..dbc8b67 100644
--- a/internal/common/common_windows.go
+++ b/internal/common/common_windows.go
@@ -4,6 +4,9 @@ package common
 
 import (
 	"context"
+	"path/filepath"
+	"strings"
+	"syscall"
 	"unsafe"
 
 	"github.com/StackExchange/wmi"
@@ -59,6 +62,8 @@ var (
 	PdhCollectQueryData          = ModPdh.NewProc("PdhCollectQueryData")
 	PdhGetFormattedCounterValue  = ModPdh.NewProc("PdhGetFormattedCounterValue")
 	PdhCloseQuery                = ModPdh.NewProc("PdhCloseQuery")
+
+	procQueryDosDeviceW = Modkernel32.NewProc("QueryDosDeviceW")
 )
 
 type FILETIME struct {
@@ -133,3 +138,23 @@ func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, con
 		return err
 	}
 }
+
+// Convert paths using native DOS format like:
+//   "\Device\HarddiskVolume1\Windows\systemew\file.txt"
+// into:
+//   "C:\Windows\systemew\file.txt"
+func ConvertDOSPath(p string) string {
+	rawDrive := strings.Join(strings.Split(p, `\`)[:3], `\`)
+
+	for d := 'A'; d <= 'Z'; d++ {
+		szDeviceName := string(d) + ":"
+		szTarget := make([]uint16, 512)
+		ret, _, _ := procQueryDosDeviceW.Call(uintptr(unsafe.Pointer(syscall.StringToUTF16Ptr(szDeviceName))),
+			uintptr(unsafe.Pointer(&szTarget[0])),
+			uintptr(len(szTarget)))
+		if ret != 0 && windows.UTF16ToString(szTarget[:]) == rawDrive {
+			return filepath.Join(szDeviceName, p[len(rawDrive):])
+		}
+	}
+	return p
+}
diff --git a/process/process_windows.go b/process/process_windows.go
index 2520958..c580ac5 100644
--- a/process/process_windows.go
+++ b/process/process_windows.go
@@ -25,12 +25,15 @@ const (
 )
 
 var (
-	modpsapi                 = windows.NewLazySystemDLL("psapi.dll")
-	procGetProcessMemoryInfo = modpsapi.NewProc("GetProcessMemoryInfo")
+	modpsapi                     = windows.NewLazySystemDLL("psapi.dll")
+	procGetProcessMemoryInfo     = modpsapi.NewProc("GetProcessMemoryInfo")
+	procGetProcessImageFileNameW = modpsapi.NewProc("GetProcessImageFileNameW")
 
 	advapi32                  = windows.NewLazySystemDLL("advapi32.dll")
 	procLookupPrivilegeValue  = advapi32.NewProc("LookupPrivilegeValueW")
 	procAdjustTokenPrivileges = advapi32.NewProc("AdjustTokenPrivileges")
+
+	procQueryFullProcessImageNameW = common.Modkernel32.NewProc("QueryFullProcessImageNameW")
 )
 
 type SystemProcessInformation struct {
@@ -234,24 +237,31 @@ func (p *Process) Exe() (string, error) {
 }
 
 func (p *Process) ExeWithContext(ctx context.Context) (string, error) {
-	if p.Pid != 0 { // 0 or null is the current process for CreateToolhelp32Snapshot
-		snap := w32.CreateToolhelp32Snapshot(w32.TH32CS_SNAPMODULE|w32.TH32CS_SNAPMODULE32, uint32(p.Pid))
-		if snap != 0 { // don't report errors here, fallback to WMI instead
-			defer w32.CloseHandle(snap)
-			var me32 w32.MODULEENTRY32
-			me32.Size = uint32(unsafe.Sizeof(me32))
-
-			if w32.Module32First(snap, &me32) {
-				szexepath := windows.UTF16ToString(me32.SzExePath[:])
-				return szexepath, nil
-			}
+	// 0x1000 is PROCESS_QUERY_LIMITED_INFORMATION
+	c, err := syscall.OpenProcess(0x1000, false, uint32(p.Pid))
+	if err != nil {
+		return "", err
+	}
+	defer syscall.CloseHandle(c)
+	buf := make([]uint16, syscall.MAX_LONG_PATH)
+	size := uint32(syscall.MAX_LONG_PATH)
+	if err := procQueryFullProcessImageNameW.Find(); err == nil { // Vista+
+		ret, _, err := procQueryFullProcessImageNameW.Call(
+			uintptr(c),
+			uintptr(0),
+			uintptr(unsafe.Pointer(&buf[0])),
+			uintptr(unsafe.Pointer(&size)))
+		if ret == 0 {
+			return "", err
 		}
+		return windows.UTF16ToString(buf[:]), nil
 	}
-	dst, err := GetWin32ProcWithContext(ctx, p.Pid)
-	if err != nil {
-		return "", fmt.Errorf("could not get ExecutablePath: %s", err)
+	// XP fallback
+	ret, _, err := procGetProcessImageFileNameW.Call(uintptr(c), uintptr(unsafe.Pointer(&buf[0])), uintptr(size))
+	if ret == 0 {
+		return "", err
 	}
-	return *dst[0].ExecutablePath, nil
+	return common.ConvertDOSPath(windows.UTF16ToString(buf[:])), nil
 }
 
 func (p *Process) Cmdline() (string, error) {