protect client process tree from blocker
All checks were successful
Deploy on Master Change / deploy (push) Successful in 37s
All checks were successful
Deploy on Master Change / deploy (push) Successful in 37s
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user