From cf9aa4a8ec20b054a533b388c1d1dd6b38880878 Mon Sep 17 00:00:00 2001
From: Lomanic <Lomanic@users.noreply.github.com>
Date: Tue, 18 Jun 2019 22:41:32 +0200
Subject: [PATCH] [process][windows] Use win32 API in process.Nice() instead of
 slow WMI call

Convert priority classes values to their WMI equivalent for backward
compatiblity.
---
 process/process_windows.go | 29 ++++++++++++++++++++++++++---
 1 file changed, 26 insertions(+), 3 deletions(-)

diff --git a/process/process_windows.go b/process/process_windows.go
index c580ac5..c84bb46 100644
--- a/process/process_windows.go
+++ b/process/process_windows.go
@@ -34,6 +34,7 @@ var (
 	procAdjustTokenPrivileges = advapi32.NewProc("AdjustTokenPrivileges")
 
 	procQueryFullProcessImageNameW = common.Modkernel32.NewProc("QueryFullProcessImageNameW")
+	procGetPriorityClass           = common.Modkernel32.NewProc("GetPriorityClass")
 )
 
 type SystemProcessInformation struct {
@@ -392,17 +393,39 @@ func (p *Process) TerminalWithContext(ctx context.Context) (string, error) {
 	return "", common.ErrNotImplementedError
 }
 
+// priorityClasses maps a win32 priority class to its WMI equivalent Win32_Process.Priority
+// https://docs.microsoft.com/en-us/windows/desktop/api/processthreadsapi/nf-processthreadsapi-getpriorityclass
+// https://docs.microsoft.com/en-us/windows/desktop/cimwin32prov/win32-process
+var priorityClasses = map[int]int32{
+	0x00008000: 10, // ABOVE_NORMAL_PRIORITY_CLASS
+	0x00004000: 6,  // BELOW_NORMAL_PRIORITY_CLASS
+	0x00000080: 13, // HIGH_PRIORITY_CLASS
+	0x00000040: 4,  // IDLE_PRIORITY_CLASS
+	0x00000020: 8,  // NORMAL_PRIORITY_CLASS
+	0x00000100: 24, // REALTIME_PRIORITY_CLASS
+}
+
 // Nice returns priority in Windows
 func (p *Process) Nice() (int32, error) {
 	return p.NiceWithContext(context.Background())
 }
 
 func (p *Process) NiceWithContext(ctx context.Context) (int32, error) {
-	dst, err := GetWin32ProcWithContext(ctx, p.Pid)
+	// 0x1000 is PROCESS_QUERY_LIMITED_INFORMATION
+	c, err := syscall.OpenProcess(0x1000, false, uint32(p.Pid))
 	if err != nil {
-		return 0, fmt.Errorf("could not get Priority: %s", err)
+		return 0, err
+	}
+	defer syscall.CloseHandle(c)
+	ret, _, err := procGetPriorityClass.Call(uintptr(c))
+	if ret == 0 {
+		return 0, err
+	}
+	priority, ok := priorityClasses[int(ret)]
+	if !ok {
+		return 0, fmt.Errorf("unknown priority class %s", ret)
 	}
-	return int32(dst[0].Priority), nil
+	return priority, nil
 }
 func (p *Process) IOnice() (int32, error) {
 	return p.IOniceWithContext(context.Background())