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{}),
|
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
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user