suspend process tree by pid (windows)

suspend process tree by pid (windows).go

// This is a script to suspend the child processes in a process tree
// Tn this example a minecraft process is used as example
// (windows specific)

package main

import (
	"bufio"
	"fmt"
	"os/exec"
	"path/filepath"
	"syscall"
	"time"
	"unsafe"
)

var (
	// DebugActiveProcess function (debugapi.h)
	// https://docs.microsoft.com/en-us/windows/win32/api/debugapi/nf-debugapi-debugactiveprocess
	modkernel32                = syscall.NewLazyDLL("kernel32.dll")
	procDebugActiveProcess     = modkernel32.NewProc("DebugActiveProcess")
	procDebugActiveProcessStop = modkernel32.NewProc("DebugActiveProcessStop")
)

func main() {
	dir := "C:\\Users\\Administrator\\Desktop\\minecraft server"
	execName := "server.jar"
	cmd := exec.Command("java", "-Xmx1024M", "-Xms1024M", "-jar", filepath.Join(dir, execName), "nogui")
	cmd.Dir = dir

	cmdprinter(cmd)

	cmd.Start()

	time.Sleep(time.Duration(20) * time.Second)

	fmt.Print("--- STOPPING ---")
	tree, _ := getTreePids(uint32(cmd.Process.Pid))
	for _, childPid := range tree {
		// https://github.com/go-delve/delve/blob/master/pkg/proc/native/zsyscall_windows.go#L177-L199
		r1, _, _ := syscall.Syscall(procDebugActiveProcess.Addr(), 1, uintptr(childPid), 0, 0)
		if r1 == 0 {
			fmt.Println("error stopping process tree")
		}
	}
	fmt.Println(" done")

	for n := 0; n < 20; n++ {
		fmt.Println(n)
		time.Sleep(time.Second)
	}

	fmt.Print("--- STARTING ---")
	tree, _ = getTreePids(uint32(cmd.Process.Pid))
	for _, childPid := range tree {
		r1, _, _ := syscall.Syscall(procDebugActiveProcessStop.Addr(), 1, uintptr(childPid), 0, 0)
		if r1 == 0 {
			fmt.Println("error stopping process tree")
		}
	}
	fmt.Println(" done")

	cmd.Wait()
	fmt.Println("exiting")
}

// getTreePids will return a list of pids that represent the tree of process pids originating from the specified one.
// (they are ordered: [parent, 1 gen child, 2 gen child, ...])
func getTreePids(rootPid uint32) ([]uint32, error) {
	contains := func(list []uint32, e uint32) bool {
		for _, l := range list {
			if l == e {
				return true
			}
		}
		return false
	}

	// https://docs.microsoft.com/en-us/windows/win32/api/tlhelp32/ns-tlhelp32-processentry32
	procEntry := syscall.ProcessEntry32{}
	parentLayer := []uint32{rootPid}
	treePids := parentLayer
	foundRootPid := false

	snapshot, err := syscall.CreateToolhelp32Snapshot(uint32(syscall.TH32CS_SNAPPROCESS), 0)
	if err != nil {
		return nil, err
	}
	defer syscall.CloseHandle(snapshot)

	procEntry.Size = uint32(unsafe.Sizeof(procEntry))

	for {
		// set procEntry to the first process in the snapshot
		err = syscall.Process32First(snapshot, &procEntry)
		if err != nil {
			return nil, err
		}

		// loop through the processes in the snapshot, if the parent pid of the analyzed process
		// is in in the parent layer, append the analyzed process pid in the child layer
		var childLayer []uint32
		for {
			if procEntry.ProcessID == rootPid {
				foundRootPid = true
			}

			if contains(parentLayer, procEntry.ParentProcessID) {
				// avoid adding a pid if it's already contained in treePids
				// useful for pid 0 whose ppid is 0 and would lead to recursion (windows)
				if !contains(treePids, procEntry.ProcessID) {
					childLayer = append(childLayer, procEntry.ProcessID)
				}
			}

			// advance to next process in snapshot
			err = syscall.Process32Next(snapshot, &procEntry)
			if err != nil {
				// if there aren't anymore processes to be analyzed, break out of the loop
				break
			}
		}

		// if the specified rootPid is not found, return error
		if !foundRootPid {
			return nil, fmt.Errorf("specified rootPid not found")
		}

		// fmt.Println(childLayer)

		// there are no more child processes, return the process tree
		if len(childLayer) == 0 {
			return treePids, nil
		}

		// append the child layer to the tree pids
		treePids = append(treePids, childLayer...)

		// to analyze the next layer, set the child layer to be the new parent layer
		parentLayer = childLayer
	}
}

func cmdprinter(cmd *exec.Cmd) {
	outP, _ := cmd.StdoutPipe()
	errP, _ := cmd.StderrPipe()
	go func() {
		scanner := bufio.NewScanner(outP)
		for scanner.Scan() {
			line := scanner.Text()
			fmt.Println(line)
		}
	}()
	go func() {
		scanner := bufio.NewScanner(errP)
		for scanner.Scan() {
			line := scanner.Text()
			fmt.Println(line)
		}
	}()
}