GoSetProcessAffinity/process_win32.go at main · spddl/GoSetProcessAffinity · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
package main // https://github.com/jet/damon/blob/8b2f833924dcfa53fc7196ad85f99d977d947e45/win32/process_win32.go

import (
	"log"
	"os"
	"syscall"
	"unsafe"

	"golang.org/x/sys/windows"
)

var (
	kernel32DLL = windows.NewLazySystemDLL("kernel32.dll")
	ntdllDLL    = windows.NewLazySystemDLL("ntdll.dll")

	procSetProcessAffinityMask  = kernel32DLL.NewProc("SetProcessAffinityMask")
	procSetProcessPriorityBoost = kernel32DLL.NewProc("SetProcessPriorityBoost")
	procSetPriorityClass        = kernel32DLL.NewProc("SetPriorityClass")

	procNtSetInformationProcess   = ntdllDLL.NewProc("NtSetInformationProcess")
	procNtQueryInformationProcess = ntdllDLL.NewProc("NtQueryInformationProcess")
	procNtSuspendProcess          = ntdllDLL.NewProc("NtSuspendProcess")
	procNtResumeProcess           = ntdllDLL.NewProc("NtResumeProcess")
)

type ULONG uint32

const ( // IO_PRIORITY_HINT
	IoPriorityVeryLow  = iota // Defragging, content indexing and other background I/Os.
	IoPriorityLow             // Prefetching for applications.
	IoPriorityNormal          // Normal I/Os.
	IoPriorityHigh            // Used by filesystems for checkpoint I/O.
	IoPriorityCritical        // Used by memory manager. Not available for applications.
)

const (
	ProcessAllAccess       = 0x1F0FFF
	ProcessSetIinformation = 0x0200
	ProcessIoPriority      = 0x21 // https://www.pinvoke.net/default.aspx/ntdll/PROCESSINFOCLASS.html
	ProcessPagePriority    = 0x27
)

// https://docs.microsoft.com/en-us/windows/desktop/api/winbase/nf-winbase-setprocessaffinitymask
func SetProcessAffinityMask(hProcess windows.Handle, dwProcessAffinityMask uint64) error {
	r1, _, e1 := procSetProcessAffinityMask.Call(
		uintptr(hProcess),
		uintptr(dwProcessAffinityMask), // uintptr(unsafe.Pointer(&sam)),
	)

	if int(r1) == 0 {
		return os.NewSyscallError("GetProcessAffinityMask", e1)
	}
	return nil // testReturnCodeNonZero(ret, errno)
}

func SetProcessPriorityBoost(process windows.Handle, disable bool) (err error) {
	var _p0 uint32
	if disable {
		_p0 = 1
	}
	r1, _, e1 := procSetProcessPriorityBoost.Call(
		uintptr(process),
		uintptr(_p0))

	if int(r1) == 0 {
		err = os.NewSyscallError("SetProcessPriorityBoost", e1)
	}
	return
}

// The Processinfoclass constants have been derived from the PROCESSINFOCLASS enum definition.
type Processinfoclass uint32

// https://github.com/hillu/go-ntdll/blob/a6f426aa8d92e55860a843a12b13d16824a082ad/process_generated.go
func NtSetInformationProcess(
	processHandle windows.Handle,
	processInformationClass Processinfoclass,
	processInformation *uint32,
	processInformationLength uint32,
) error {
	r1, _, e1 := procNtSetInformationProcess.Call(
		uintptr(processHandle),
		uintptr(processInformationClass),
		uintptr(unsafe.Pointer(processInformation)),
		uintptr(processInformationLength))
	if int(r1) == 0 {
		return os.NewSyscallError("SetInformationProcess", e1)
	}
	return nil
}

// NtQueryInformationProcess is a wrapper for ntdll.NtQueryInformationProcess.
// The handle must have the PROCESS_QUERY_INFORMATION access right.
// Returns an error of type NTStatus.
func NtQueryInformationProcess(processHandle windows.Handle, processInformationClass int32, processInformation windows.Pointer, processInformationLength uint32, returnLength *uint32) error {
	r1, _, err := procNtQueryInformationProcess.Call(
		uintptr(processHandle),
		uintptr(processInformationClass),
		uintptr(unsafe.Pointer(processInformation)),
		uintptr(processInformationLength),
		uintptr(unsafe.Pointer(returnLength)))
	if int(r1) < 0 {
		return os.NewSyscallError("NtQueryInformationProcess", err)
	}
	return nil
}

func SetPriorityClass(process syscall.Handle, priorityClass uint32) (err error) {
	r1, r2, e1 := syscall.Syscall(procSetPriorityClass.Addr(), 2, uintptr(process), uintptr(priorityClass), 0)
	log.Println(r1, r2, e1)
	if r1 == 0 {
		if e1 != 0 {
			err = e1
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

func NtSuspendProcess(process windows.Handle) (err error) {
	r1, _, e1 := procNtSuspendProcess.Call(
		uintptr(process))

	if int(r1) == 0 {
		err = os.NewSyscallError("NtSuspendProcess", e1)
	}
	return
}

func NtResumeProcess(process windows.Handle) (err error) {
	r1, _, e1 := procNtResumeProcess.Call(
		uintptr(process))

	if int(r1) == 0 {
		err = os.NewSyscallError("NtResumeProcess", e1)
	}
	return
}