From c211741a304a92e2f230f93fc840ab91729c10e8 Mon Sep 17 00:00:00 2001
From: PhuocNTB <phuocntb@mieusoft.com>
Date: Wed, 1 Jul 2026 08:46:38 +0700
Subject: [PATCH] protect client process tree from blocker

---
 client/internal/blocker/blocker.go | 44 ++++++++++++++++++++++--------
 1 file changed, 33 insertions(+), 11 deletions(-)

diff --git a/client/internal/blocker/blocker.go b/client/internal/blocker/blocker.go
index 3502bd8..ffe0a63 100644
--- a/client/internal/blocker/blocker.go
+++ b/client/internal/blocker/blocker.go
@@ -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
 		}