autoPriority/windows.go

//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
	}
}