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{}), 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) snapshot, err := syscall.CreateToolhelp32Snapshot(syscall.TH32CS_SNAPPROCESS, 0)
if err != nil { if err != nil {
return nil, err return nil, nil, err
} }
defer syscall.CloseHandle(snapshot) defer syscall.CloseHandle(snapshot)
@@ -95,20 +95,22 @@ func getProcessMap() (map[uint32]string, error) {
err = syscall.Process32First(snapshot, &pe) err = syscall.Process32First(snapshot, &pe)
if err != nil { if err != nil {
return nil, err return nil, nil, err
} }
pm := make(map[uint32]string) pm := make(map[uint32]string)
parents := make(map[uint32]uint32)
for { for {
name := syscall.UTF16ToString(pe.ExeFile[:]) name := syscall.UTF16ToString(pe.ExeFile[:])
pm[pe.ProcessID] = name pm[pe.ProcessID] = name
parents[pe.ProcessID] = pe.ParentProcessID
err = syscall.Process32Next(snapshot, &pe) err = syscall.Process32Next(snapshot, &pe)
if err != nil { if err != nil {
break break
} }
} }
return pm, nil return pm, parents, nil
} }
func getWindowText(hwnd uintptr) string { func getWindowText(hwnd uintptr) string {
@@ -125,6 +127,7 @@ var (
enumWindowsMutex sync.Mutex enumWindowsMutex sync.Mutex
enumWindowsList []WindowInfo enumWindowsList []WindowInfo
enumProcessMap map[uint32]string enumProcessMap map[uint32]string
enumParentMap map[uint32]uint32
) )
var enumWindowsCallback = syscall.NewCallback(func(hwnd uintptr, lParam uintptr) uintptr { 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 return 1
}) })
func EnumerateGUIWindows() ([]WindowInfo, error) { func EnumerateGUIWindows() ([]WindowInfo, map[uint32]uint32, error) {
pMap, err := getProcessMap() pMap, parentMap, err := getProcessMap()
if err != nil { if err != nil {
return nil, err return nil, nil, err
} }
enumWindowsMutex.Lock() enumWindowsMutex.Lock()
@@ -171,16 +174,18 @@ func EnumerateGUIWindows() ([]WindowInfo, error) {
enumWindowsList = make([]WindowInfo, 0, 100) enumWindowsList = make([]WindowInfo, 0, 100)
enumProcessMap = pMap enumProcessMap = pMap
enumParentMap = parentMap
procEnumWindows.Call(enumWindowsCallback, 0) procEnumWindows.Call(enumWindowsCallback, 0)
// Clean up map reference so GC can reclaim it // Clean up map reference so GC can reclaim it
enumProcessMap = nil enumProcessMap = nil
enumParentMap = nil
// Copy to a new slice to return safely // Copy to a new slice to return safely
res := make([]WindowInfo, len(enumWindowsList)) res := make([]WindowInfo, len(enumWindowsList))
copy(res, enumWindowsList) copy(res, enumWindowsList)
return res, nil return res, parentMap, nil
} }
func parseKeywordList(keywords string) []string { func parseKeywordList(keywords string) []string {
@@ -272,6 +277,21 @@ var systemAllowed = map[string]bool{
"bash.exe": true, // Bash "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() { func (b *Blocker) checkAndKill() {
b.mu.Lock() b.mu.Lock()
keywords := make([]string, len(b.allowedKeywords)) keywords := make([]string, len(b.allowedKeywords))
@@ -288,17 +308,19 @@ func (b *Blocker) checkAndKill() {
currentExec = strings.ToLower(filepath.Base(execPath)) currentExec = strings.ToLower(filepath.Base(execPath))
} }
windows, err := EnumerateGUIWindows() windows, parentMap, err := EnumerateGUIWindows()
if err != nil { if err != nil {
return return
} }
myPid := uint32(os.Getpid())
for _, w := range windows { for _, w := range windows {
pNameLower := strings.ToLower(w.ProcessName) pNameLower := strings.ToLower(w.ProcessName)
wTitleLower := strings.ToLower(w.Title) 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) // 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)
if w.PID == uint32(os.Getpid()) || (currentExec != "" && pNameLower == currentExec) || systemAllowed[pNameLower] || strings.Contains(pNameLower, "antigravity") || strings.Contains(wTitleLower, "antigravity") { 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 continue
} }