| 1 | //go:build windows |
| 2 | |
| 3 | package proc |
| 4 | |
| 5 | import ( |
| 6 | "os/exec" |
| 7 | "sort" |
| 8 | "strings" |
| 9 | "sync" |
| 10 | "time" |
| 11 | "unsafe" |
| 12 | |
| 13 | "golang.org/x/sys/windows" |
| 14 | ) |
| 15 | |
| 16 | // TreeTracker records a process tree while a command is running. Windows Job |
| 17 | // Objects should own normal children, but Git Bash/MSYS launch chains can briefly |
| 18 | // expose grandchildren before or outside taskkill's live tree walk. Recording |
| 19 | // descendants gives cancellation a second chance to terminate those escapees. |
| 20 | type TreeTracker struct { |
| 21 | root uint32 |
| 22 | done chan struct{} |
| 23 | once sync.Once |
| 24 | |
| 25 | mu sync.Mutex |
| 26 | records map[uint32]processRecord |
| 27 | } |
| 28 | |
| 29 | type processRecord struct { |
| 30 | pid uint32 |
| 31 | parent uint32 |
| 32 | exe string |
| 33 | created windows.Filetime |
| 34 | hasTimes bool |
| 35 | } |
| 36 | |
| 37 | func TrackTree(cmd *exec.Cmd) *TreeTracker { |
| 38 | if cmd == nil || cmd.Process == nil { |
| 39 | return nil |
| 40 | } |
| 41 | t := &TreeTracker{ |
| 42 | root: uint32(cmd.Process.Pid), |
| 43 | done: make(chan struct{}), |
| 44 | records: map[uint32]processRecord{}, |
| 45 | } |
| 46 | t.record() |
| 47 | go t.loop() |
| 48 | return t |
| 49 | } |
| 50 | |
| 51 | func (t *TreeTracker) Stop() { |
| 52 | if t == nil { |
| 53 | return |
| 54 | } |
| 55 | t.once.Do(func() { close(t.done) }) |
| 56 | } |
| 57 | |
| 58 | func (t *TreeTracker) Kill() int { |
| 59 | if t == nil { |
| 60 | return 0 |
| 61 | } |
| 62 | t.record() |
| 63 | records := t.snapshot() |
| 64 | killed := 0 |
| 65 | for _, rec := range records { |
| 66 | if rec.pid != t.root { |
| 67 | killed += terminateRecord(rec) |
| 68 | } |
| 69 | } |
| 70 | for _, rec := range records { |
| 71 | if rec.pid == t.root { |
| 72 | killed += terminateRecord(rec) |
| 73 | break |
| 74 | } |
| 75 | } |
| 76 | return killed |
| 77 | } |
| 78 | |
| 79 | func (t *TreeTracker) loop() { |
| 80 | ticker := time.NewTicker(100 * time.Millisecond) |
| 81 | defer ticker.Stop() |
| 82 | for { |
| 83 | select { |
| 84 | case <-ticker.C: |
| 85 | t.record() |
| 86 | case <-t.done: |
| 87 | return |
| 88 | } |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | func (t *TreeTracker) record() { |
| 93 | if t == nil || t.root == 0 { |
| 94 | return |
| 95 | } |
| 96 | records := processSnapshot() |
| 97 | t.mu.Lock() |
| 98 | if root, ok := records[t.root]; ok { |
| 99 | t.records[t.root] = root |
| 100 | } |
| 101 | for _, rec := range descendantRecords(t.root, records) { |
| 102 | t.records[rec.pid] = rec |
| 103 | } |
| 104 | t.mu.Unlock() |
| 105 | } |
| 106 | |
| 107 | func (t *TreeTracker) snapshot() []processRecord { |
| 108 | t.mu.Lock() |
| 109 | defer t.mu.Unlock() |
| 110 | out := make([]processRecord, 0, len(t.records)) |
| 111 | for _, rec := range t.records { |
| 112 | out = append(out, rec) |
| 113 | } |
| 114 | sort.Slice(out, func(i, j int) bool { return out[i].pid < out[j].pid }) |
| 115 | return out |
| 116 | } |
| 117 | |
| 118 | func descendantRecords(root uint32, records map[uint32]processRecord) []processRecord { |
| 119 | if root == 0 { |
| 120 | return nil |
| 121 | } |
| 122 | children := map[uint32][]uint32{} |
| 123 | for _, rec := range records { |
| 124 | children[rec.parent] = append(children[rec.parent], rec.pid) |
| 125 | } |
| 126 | |
| 127 | var out []processRecord |
| 128 | seen := map[uint32]bool{root: true} |
| 129 | var walk func(uint32) |
| 130 | walk = func(pid uint32) { |
| 131 | for _, child := range children[pid] { |
| 132 | if child == 0 || seen[child] { |
| 133 | continue |
| 134 | } |
| 135 | seen[child] = true |
| 136 | if rec, ok := records[child]; ok { |
| 137 | out = append(out, rec) |
| 138 | } |
| 139 | walk(child) |
| 140 | } |
| 141 | } |
| 142 | walk(root) |
| 143 | return out |
| 144 | } |
| 145 | |
| 146 | func processSnapshot() map[uint32]processRecord { |
| 147 | snap, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, 0) |
| 148 | if err != nil { |
| 149 | return nil |
| 150 | } |
| 151 | defer func() { _ = windows.CloseHandle(snap) }() |
| 152 | |
| 153 | records := map[uint32]processRecord{} |
| 154 | var pe windows.ProcessEntry32 |
| 155 | pe.Size = uint32(unsafe.Sizeof(pe)) |
| 156 | for err := windows.Process32First(snap, &pe); err == nil; err = windows.Process32Next(snap, &pe) { |
| 157 | rec := processRecord{ |
| 158 | pid: pe.ProcessID, |
| 159 | parent: pe.ParentProcessID, |
| 160 | exe: strings.ToLower(windows.UTF16ToString(pe.ExeFile[:])), |
| 161 | } |
| 162 | rec.created, rec.hasTimes = processCreationTime(pe.ProcessID) |
| 163 | records[rec.pid] = rec |
| 164 | } |
| 165 | return records |
| 166 | } |
| 167 | |
| 168 | func processCreationTime(pid uint32) (windows.Filetime, bool) { |
| 169 | h, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid) |
| 170 | if err != nil { |
| 171 | return windows.Filetime{}, false |
| 172 | } |
| 173 | defer func() { _ = windows.CloseHandle(h) }() |
| 174 | var created, exited, kernel, user windows.Filetime |
| 175 | if err := windows.GetProcessTimes(h, &created, &exited, &kernel, &user); err != nil { |
| 176 | return windows.Filetime{}, false |
| 177 | } |
| 178 | return created, true |
| 179 | } |
| 180 | |
| 181 | func terminateRecord(rec processRecord) int { |
| 182 | if rec.pid == 0 { |
| 183 | return 0 |
| 184 | } |
| 185 | current, ok := processSnapshot()[rec.pid] |
| 186 | if !ok || !sameProcessIdentity(rec, current) { |
| 187 | return 0 |
| 188 | } |
| 189 | h, err := windows.OpenProcess(windows.PROCESS_TERMINATE, false, rec.pid) |
| 190 | if err != nil { |
| 191 | return 0 |
| 192 | } |
| 193 | defer func() { _ = windows.CloseHandle(h) }() |
| 194 | _ = windows.TerminateProcess(h, 1) |
| 195 | return 1 |
| 196 | } |
| 197 | |
| 198 | func sameProcessIdentity(recorded, current processRecord) bool { |
| 199 | if recorded.pid != current.pid { |
| 200 | return false |
| 201 | } |
| 202 | if recorded.hasTimes && current.hasTimes { |
| 203 | return recorded.created == current.created |
| 204 | } |
| 205 | if recorded.exe != "" && current.exe != "" { |
| 206 | return strings.EqualFold(recorded.exe, current.exe) |
| 207 | } |
| 208 | return true |
| 209 | } |
| 210 |