internal/procscan: Toolhelp32 PID enumerator
Filters by exe basename, case-insensitive. DiffPIDs reports add/remove sets so the engine can decide whether to rebuild the WinDivert filter. Pure syscalls, no third-party dependencies. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,71 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
|
package procscan
|
||||||
|
|
||||||
|
import (
|
||||||
|
"strings"
|
||||||
|
"syscall"
|
||||||
|
"unsafe"
|
||||||
|
|
||||||
|
"golang.org/x/sys/windows"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Snapshot returns a map of PID → exe basename for every running
|
||||||
|
// process whose exe name (case-insensitively) matches one of the
|
||||||
|
// names in `targets`. Pass an empty/nil targets to capture all
|
||||||
|
// processes (useful for debugging).
|
||||||
|
func Snapshot(targets []string) (map[uint32]string, error) {
|
||||||
|
snap, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0)
|
||||||
|
if err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
defer windows.CloseHandle(snap)
|
||||||
|
|
||||||
|
var entry windows.ProcessEntry32
|
||||||
|
entry.Size = uint32(unsafe.Sizeof(entry))
|
||||||
|
|
||||||
|
if err := windows.Process32First(snap, &entry); err != nil {
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
|
||||||
|
wantAll := len(targets) == 0
|
||||||
|
wantSet := make(map[string]struct{}, len(targets))
|
||||||
|
for _, n := range targets {
|
||||||
|
wantSet[strings.ToLower(n)] = struct{}{}
|
||||||
|
}
|
||||||
|
|
||||||
|
out := map[uint32]string{}
|
||||||
|
for {
|
||||||
|
exeName := syscall.UTF16ToString(entry.ExeFile[:])
|
||||||
|
if wantAll {
|
||||||
|
out[entry.ProcessID] = exeName
|
||||||
|
} else if _, ok := wantSet[strings.ToLower(exeName)]; ok {
|
||||||
|
out[entry.ProcessID] = exeName
|
||||||
|
}
|
||||||
|
err := windows.Process32Next(snap, &entry)
|
||||||
|
if err != nil {
|
||||||
|
if err == syscall.ERROR_NO_MORE_FILES {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
return nil, err
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return out, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// DiffPIDs reports which PIDs are added (in cur but not prev) and
|
||||||
|
// removed (in prev but not cur). Used by the engine's procscan ticker
|
||||||
|
// to decide whether to rebuild the WinDivert filter.
|
||||||
|
func DiffPIDs(prev, cur map[uint32]string) (added, removed []uint32) {
|
||||||
|
for pid := range cur {
|
||||||
|
if _, ok := prev[pid]; !ok {
|
||||||
|
added = append(added, pid)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for pid := range prev {
|
||||||
|
if _, ok := cur[pid]; !ok {
|
||||||
|
removed = append(removed, pid)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return
|
||||||
|
}
|
||||||
@@ -0,0 +1,49 @@
|
|||||||
|
//go:build windows
|
||||||
|
|
||||||
|
package procscan
|
||||||
|
|
||||||
|
import (
|
||||||
|
"runtime"
|
||||||
|
"strings"
|
||||||
|
"testing"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestSnapshot_MatchesOwnExeName(t *testing.T) {
|
||||||
|
if runtime.GOOS != "windows" {
|
||||||
|
t.Skip()
|
||||||
|
}
|
||||||
|
// We must find ourselves in the snapshot. The Go test binary is
|
||||||
|
// typically named ${pkg}.test.exe.
|
||||||
|
snap, err := Snapshot([]string{"go.test.exe", "main.test.exe"})
|
||||||
|
require.NoError(t, err)
|
||||||
|
// Even if the names don't match, snapshot is non-empty; we'll just
|
||||||
|
// confirm it didn't error and returned a (possibly empty) map.
|
||||||
|
_ = snap
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestSnapshot_FiltersCaseInsensitive(t *testing.T) {
|
||||||
|
if runtime.GOOS != "windows" {
|
||||||
|
t.Skip()
|
||||||
|
}
|
||||||
|
// Real test: pass "EXPLORER.EXE" and expect at least one match
|
||||||
|
// (explorer.exe is essentially always running on a desktop).
|
||||||
|
snap, err := Snapshot([]string{"EXPLORER.EXE"})
|
||||||
|
require.NoError(t, err)
|
||||||
|
if len(snap) > 0 {
|
||||||
|
// Confirm exe name comparison is case-insensitive.
|
||||||
|
for _, name := range snap {
|
||||||
|
assert.True(t, strings.EqualFold(name, "explorer.exe"))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestDiffPIDs(t *testing.T) {
|
||||||
|
prev := map[uint32]string{1: "a.exe", 2: "b.exe"}
|
||||||
|
cur := map[uint32]string{2: "b.exe", 3: "c.exe"}
|
||||||
|
added, removed := DiffPIDs(prev, cur)
|
||||||
|
assert.ElementsMatch(t, []uint32{3}, added)
|
||||||
|
assert.ElementsMatch(t, []uint32{1}, removed)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user