feat: add tray control, priority restoration and session logging
This commit is contained in:
lzrdblzzrd committed
1 parent
4f8f6218ad
commit
4880aa46dc
25 files changed
+3072
-676
No files matched your search
+295
@@ -0,0 +1,295 @@
|
||||
//go:build windows
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"syscall"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
const (
|
||||
CreateToolhelp32SnapshotProcess = 0x00000002
|
||||
ProcessSetInformation = 0x00000200
|
||||
ProcessQueryLimitedInformation = 0x00001000
|
||||
|
||||
ErrorNoMoreFiles = 18
|
||||
ErrorAlreadyExists = 183
|
||||
)
|
||||
|
||||
type processEntry32 struct {
|
||||
Size uint32
|
||||
CntUsage uint32
|
||||
PID uint32
|
||||
DefaultHeapID uintptr
|
||||
ModuleID uint32
|
||||
CntThreads uint32
|
||||
ParentPID uint32
|
||||
PrioClass int32
|
||||
Flags uint32
|
||||
ExeFile [260]uint16
|
||||
}
|
||||
|
||||
type processMemoryCounters struct {
|
||||
CBM uint32
|
||||
PageFaultCount uint32
|
||||
PeakWorkingSetSize uintptr
|
||||
WorkingSetSize uintptr
|
||||
QuotaPeakPagedPoolUsage uintptr
|
||||
QuotaPagedPoolUsage uintptr
|
||||
QuotaPeakNonPagedPoolUsage uintptr
|
||||
QuotaNonPagedPoolUsage uintptr
|
||||
PagefileUsage uintptr
|
||||
PeakPagefileUsage uintptr
|
||||
PrivateUsage uintptr
|
||||
}
|
||||
|
||||
type processKey struct {
|
||||
PID uint32
|
||||
Created uint64
|
||||
}
|
||||
|
||||
type procInfo struct {
|
||||
PID uint32
|
||||
Name string
|
||||
RSS uint64
|
||||
Key processKey
|
||||
Prio uint32
|
||||
Handle syscall.Handle
|
||||
Waitable bool
|
||||
ParentPID uint32
|
||||
}
|
||||
|
||||
var (
|
||||
k32 = syscall.NewLazyDLL("kernel32.dll")
|
||||
|
||||
procCreateSnap = k32.NewProc("CreateToolhelp32Snapshot")
|
||||
procProcess32First = k32.NewProc("Process32FirstW")
|
||||
procProcess32Next = k32.NewProc("Process32NextW")
|
||||
procGetMemInfo = k32.NewProc("K32GetProcessMemoryInfo")
|
||||
procSetPriority = k32.NewProc("SetPriorityClass")
|
||||
procGetPriority = k32.NewProc("GetPriorityClass")
|
||||
procOpenProcess = k32.NewProc("OpenProcess")
|
||||
procCloseHandle = k32.NewProc("CloseHandle")
|
||||
procGetProcessTimes = k32.NewProc("GetProcessTimes")
|
||||
procQueryImageName = k32.NewProc("QueryFullProcessImageNameW")
|
||||
procCreateMutex = k32.NewProc("CreateMutexW")
|
||||
)
|
||||
|
||||
func closeH(h syscall.Handle) {
|
||||
_, _, _ = procCloseHandle.Call(uintptr(h))
|
||||
}
|
||||
|
||||
func openProc(pid uint32, acc uint32) (syscall.Handle, error) {
|
||||
r, _, e := procOpenProcess.Call(uintptr(acc), 0, uintptr(pid))
|
||||
if r == 0 {
|
||||
return 0, fmt.Errorf("OpenProcess(%d) failed: %w", pid, e)
|
||||
}
|
||||
return syscall.Handle(r), nil
|
||||
}
|
||||
|
||||
func setPrio(h syscall.Handle, pid uint32, cls uint32) error {
|
||||
r, _, e := procSetPriority.Call(uintptr(h), uintptr(cls))
|
||||
if r == 0 {
|
||||
return fmt.Errorf("SetPriorityClass(%d) failed: %w", pid, e)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func processCreated(h syscall.Handle) (uint64, error) {
|
||||
var created, exited, kernel, user syscall.Filetime
|
||||
r, _, e := procGetProcessTimes.Call(
|
||||
uintptr(h),
|
||||
uintptr(unsafe.Pointer(&created)),
|
||||
uintptr(unsafe.Pointer(&exited)),
|
||||
uintptr(unsafe.Pointer(&kernel)),
|
||||
uintptr(unsafe.Pointer(&user)),
|
||||
)
|
||||
if r == 0 {
|
||||
return 0, fmt.Errorf("GetProcessTimes failed: %w", e)
|
||||
}
|
||||
return uint64(created.HighDateTime)<<32 | uint64(created.LowDateTime), nil
|
||||
}
|
||||
|
||||
func processName(h syscall.Handle, buf []uint16) (string, error) {
|
||||
size := uint32(len(buf))
|
||||
r, _, e := procQueryImageName.Call(uintptr(h), 0, uintptr(unsafe.Pointer(&buf[0])), uintptr(unsafe.Pointer(&size)))
|
||||
if r == 0 {
|
||||
return "", fmt.Errorf("QueryFullProcessImageName failed: %w", e)
|
||||
}
|
||||
return filepath.Base(syscall.UTF16ToString(buf[:size])), nil
|
||||
}
|
||||
|
||||
func processAlive(h syscall.Handle) (bool, error) {
|
||||
r, _, err := procWaitForSingleObject.Call(uintptr(h), 0)
|
||||
return decodeWaitResult(r, err)
|
||||
}
|
||||
func decodeWaitResult(r uintptr, err error) (bool, error) {
|
||||
switch r {
|
||||
case 0:
|
||||
return false, nil
|
||||
case 0x102:
|
||||
return true, nil
|
||||
default:
|
||||
return false, fmt.Errorf("WaitForSingleObject returned %#x: %v", r, err)
|
||||
}
|
||||
}
|
||||
|
||||
func allProcs() ([]procInfo, error) {
|
||||
snap, _, e := procCreateSnap.Call(CreateToolhelp32SnapshotProcess, 0)
|
||||
if snap == ^uintptr(0) {
|
||||
return nil, fmt.Errorf("CreateToolhelp32Snapshot failed: %w", e)
|
||||
}
|
||||
defer closeH(syscall.Handle(snap))
|
||||
|
||||
pe := processEntry32{Size: uint32(unsafe.Sizeof(processEntry32{}))}
|
||||
r, _, e := procProcess32First.Call(snap, uintptr(unsafe.Pointer(&pe)))
|
||||
if r == 0 {
|
||||
if e == syscall.Errno(ErrorNoMoreFiles) {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, fmt.Errorf("Process32First failed: %w", e)
|
||||
}
|
||||
|
||||
var out []procInfo
|
||||
for {
|
||||
name := syscall.UTF16ToString(pe.ExeFile[:])
|
||||
out = append(out, procInfo{PID: pe.PID, ParentPID: pe.ParentPID, Name: name})
|
||||
pe.Size = uint32(unsafe.Sizeof(processEntry32{}))
|
||||
r, _, e = procProcess32Next.Call(snap, uintptr(unsafe.Pointer(&pe)))
|
||||
if r == 0 {
|
||||
if e == syscall.Errno(ErrorNoMoreFiles) {
|
||||
break
|
||||
}
|
||||
return nil, fmt.Errorf("Process32Next failed: %w", e)
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func singleInstance() (syscall.Handle, bool, error) {
|
||||
name, err := syscall.UTF16PtrFromString(`Local\autoPriority`)
|
||||
if err != nil {
|
||||
return 0, false, err
|
||||
}
|
||||
r, _, e := procCreateMutex.Call(0, 0, uintptr(unsafe.Pointer(name)))
|
||||
if r == 0 {
|
||||
return 0, false, fmt.Errorf("CreateMutex failed: %w", e)
|
||||
}
|
||||
return syscall.Handle(r), e == syscall.Errno(ErrorAlreadyExists), nil
|
||||
}
|
||||
|
||||
const processSynchronize = 0x00100000
|
||||
|
||||
var procWaitForSingleObject = k32.NewProc("WaitForSingleObject")
|
||||
|
||||
type windowsBackend struct {
|
||||
list func() ([]procInfo, error)
|
||||
self uint32
|
||||
nameBuffer [32768]uint16
|
||||
unreadable map[uint32]procInfo
|
||||
}
|
||||
|
||||
func newWindowsBackend() *windowsBackend {
|
||||
return &windowsBackend{list: allProcs, self: uint32(os.Getpid()), unreadable: make(map[uint32]procInfo)}
|
||||
}
|
||||
func (w *windowsBackend) inspect(p procInfo) (procInfo, error) {
|
||||
h, err := openProc(p.PID, ProcessQueryLimitedInformation|processSynchronize)
|
||||
p.Waitable = err == nil
|
||||
if err != nil {
|
||||
h, err = openProc(p.PID, ProcessQueryLimitedInformation)
|
||||
}
|
||||
if err != nil {
|
||||
return procInfo{}, err
|
||||
}
|
||||
good := false
|
||||
defer func() {
|
||||
if !good {
|
||||
closeH(h)
|
||||
}
|
||||
}()
|
||||
created, err := processCreated(h)
|
||||
if err != nil {
|
||||
return procInfo{}, err
|
||||
}
|
||||
if name, e := processName(h, w.nameBuffer[:]); e == nil {
|
||||
p.Name = name
|
||||
}
|
||||
p.Handle = h
|
||||
p.Key = processKey{PID: p.PID, Created: created}
|
||||
var m processMemoryCounters
|
||||
m.CBM = uint32(unsafe.Sizeof(m))
|
||||
r, _, e := procGetMemInfo.Call(uintptr(h), uintptr(unsafe.Pointer(&m)), uintptr(unsafe.Sizeof(m)))
|
||||
if r == 0 {
|
||||
return procInfo{}, fmt.Errorf("K32GetProcessMemoryInfo: %w", e)
|
||||
}
|
||||
p.RSS = uint64(m.WorkingSetSize)
|
||||
p.Prio, err = w.Priority(&p)
|
||||
if err != nil {
|
||||
return procInfo{}, err
|
||||
}
|
||||
good = true
|
||||
return p, nil
|
||||
}
|
||||
func (w *windowsBackend) Snapshot() (processSnapshot, error) {
|
||||
entries, err := w.list()
|
||||
if err != nil {
|
||||
return processSnapshot{}, err
|
||||
}
|
||||
snap := processSnapshot{PIDs: make(map[uint32]bool, len(entries)), Processes: make([]procInfo, 0, len(entries))}
|
||||
for _, entry := range entries {
|
||||
snap.PIDs[entry.PID] = true
|
||||
if entry.PID == 0 || entry.PID == w.self {
|
||||
continue
|
||||
}
|
||||
if blocked, ok := w.unreadable[entry.PID]; ok && blocked.Name == entry.Name && blocked.ParentPID == entry.ParentPID {
|
||||
continue
|
||||
}
|
||||
info, err := w.inspect(entry)
|
||||
if err != nil {
|
||||
snap.Errors = append(snap.Errors, processError{Name: entry.Name, PID: entry.PID, Description: fmt.Sprintf("чтение: %v; внесён в блоклист", err)})
|
||||
w.unreadable[entry.PID] = entry
|
||||
continue
|
||||
}
|
||||
delete(w.unreadable, entry.PID)
|
||||
snap.Processes = append(snap.Processes, info)
|
||||
}
|
||||
for pid := range w.unreadable {
|
||||
if !snap.PIDs[pid] {
|
||||
delete(w.unreadable, pid)
|
||||
}
|
||||
}
|
||||
return snap, nil
|
||||
}
|
||||
func (w *windowsBackend) Alive(p *procInfo) (bool, error) { return processAlive(p.Handle) }
|
||||
func (w *windowsBackend) Priority(p *procInfo) (uint32, error) {
|
||||
r, _, err := procGetPriority.Call(uintptr(p.Handle))
|
||||
if r == 0 {
|
||||
return 0, fmt.Errorf("GetPriorityClass: %w", err)
|
||||
}
|
||||
return uint32(r), nil
|
||||
}
|
||||
func (w *windowsBackend) SetPriority(p *procInfo, target uint32) error {
|
||||
// Request write access only for an actual change, never while enumerating.
|
||||
h, err := openProc(p.PID, ProcessQueryLimitedInformation|ProcessSetInformation)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer closeH(h)
|
||||
created, err := processCreated(h)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if created != p.Key.Created {
|
||||
return fmt.Errorf("process instance changed; refusing to write")
|
||||
}
|
||||
return setPrio(h, p.PID, target)
|
||||
}
|
||||
func (w *windowsBackend) Close(p *procInfo) {
|
||||
if p.Handle != 0 {
|
||||
closeH(p.Handle)
|
||||
p.Handle = 0
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user