protect client process tree from blocker
All checks were successful
Deploy on Master Change / deploy (push) Successful in 37s

This commit is contained in:
2026-07-01 08:46:38 +07:00
parent 393a972e58
commit c211741a30

View File

@@ -83,10 +83,10 @@ var Instance = &Blocker{
stopChan: make(chan struct{}),
}
func getProcessMap() (map[uint32]string, error) {
func getProcessMap() (map[uint32]string, map[uint32]uint32, error) {
snapshot, err := syscall.CreateToolhelp32Snapshot(syscall.TH32CS_SNAPPROCESS, 0)
if err != nil {
return nil, err
return nil, nil, err
}
defer syscall.CloseHandle(snapshot)
@@ -95,20 +95,22 @@ func getProcessMap() (map[uint32]string, error) {
err = syscall.Process32First(snapshot, &pe)
if err != nil {
return nil, err
return nil, nil, err
}
pm := make(map[uint32]string)
parents := make(map[uint32]uint32)
for {
name := syscall.UTF16ToString(pe.ExeFile[:])
pm[pe.ProcessID] = name
parents[pe.ProcessID] = pe.ParentProcessID
err = syscall.Process32Next(snapshot, &pe)
if err != nil {
break
}
}
return pm, nil
return pm, parents, nil
}
func getWindowText(hwnd uintptr) string {
@@ -125,6 +127,7 @@ var (
enumWindowsMutex sync.Mutex
enumWindowsList []WindowInfo
enumProcessMap map[uint32]string
enumParentMap map[uint32]uint32
)
var enumWindowsCallback = syscall.NewCallback(func(hwnd uintptr, lParam uintptr) uintptr {
@@ -160,10 +163,10 @@ var enumWindowsCallback = syscall.NewCallback(func(hwnd uintptr, lParam uintptr)
return 1
})
func EnumerateGUIWindows() ([]WindowInfo, error) {
pMap, err := getProcessMap()
func EnumerateGUIWindows() ([]WindowInfo, map[uint32]uint32, error) {
pMap, parentMap, err := getProcessMap()
if err != nil {
return nil, err
return nil, nil, err
}
enumWindowsMutex.Lock()
@@ -171,16 +174,18 @@ func EnumerateGUIWindows() ([]WindowInfo, error) {
enumWindowsList = make([]WindowInfo, 0, 100)
enumProcessMap = pMap
enumParentMap = parentMap
procEnumWindows.Call(enumWindowsCallback, 0)
// Clean up map reference so GC can reclaim it
enumProcessMap = nil
enumParentMap = nil
// Copy to a new slice to return safely
res := make([]WindowInfo, len(enumWindowsList))
copy(res, enumWindowsList)
return res, nil
return res, parentMap, nil
}
func parseKeywordList(keywords string) []string {
@@ -272,6 +277,21 @@ var systemAllowed = map[string]bool{
"bash.exe": true, // Bash
}
func isDescendant(pid, targetPid uint32, parentMap map[uint32]uint32) bool {
curr := pid
for i := 0; i < 16; i++ {
parent, ok := parentMap[curr]
if !ok || parent == 0 {
return false
}
if parent == targetPid {
return true
}
curr = parent
}
return false
}
func (b *Blocker) checkAndKill() {
b.mu.Lock()
keywords := make([]string, len(b.allowedKeywords))
@@ -288,17 +308,19 @@ func (b *Blocker) checkAndKill() {
currentExec = strings.ToLower(filepath.Base(execPath))
}
windows, err := EnumerateGUIWindows()
windows, parentMap, err := EnumerateGUIWindows()
if err != nil {
return
}
myPid := uint32(os.Getpid())
for _, w := range windows {
pNameLower := strings.ToLower(w.ProcessName)
wTitleLower := strings.ToLower(w.Title)
// 1. Luôn cho phép hệ thống/app cốt lõi hoặc chính tiến trình này (bao gồm tiến trình đổi tên)
if w.PID == uint32(os.Getpid()) || (currentExec != "" && pNameLower == currentExec) || systemAllowed[pNameLower] || strings.Contains(pNameLower, "antigravity") || strings.Contains(wTitleLower, "antigravity") {
// 1. Luôn cho phép hệ thống/app cốt lõi hoặc chính tiến trình này (bao gồm tiến trình con/cháu, và đổi tên)
isOurApp := w.PID == myPid || isDescendant(w.PID, myPid, parentMap)
if isOurApp || (currentExec != "" && pNameLower == currentExec) || systemAllowed[pNameLower] || strings.Contains(pNameLower, "antigravity") || strings.Contains(wTitleLower, "antigravity") {
continue
}