123 lines
3.4 KiB
Go
123 lines
3.4 KiB
Go
//go:build windows
|
|
|
|
package subprocess
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"os/exec"
|
|
"syscall"
|
|
"unsafe"
|
|
|
|
"golang.org/x/sys/windows"
|
|
)
|
|
|
|
type windowsProcessTree struct {
|
|
job windows.Handle
|
|
}
|
|
|
|
func newOwnedProcessTree() (ownedProcessTree, error) {
|
|
return &windowsProcessTree{}, nil
|
|
}
|
|
|
|
func (tree *windowsProcessTree) Start(cmd *exec.Cmd) error {
|
|
job, err := windows.CreateJobObject(nil, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("create job object: %w", err)
|
|
}
|
|
|
|
limits := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
|
|
limits.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
|
|
if _, err := windows.SetInformationJobObject(job, windows.JobObjectExtendedLimitInformation, uintptr(unsafe.Pointer(&limits)), uint32(unsafe.Sizeof(limits))); err != nil {
|
|
_ = windows.CloseHandle(job)
|
|
return fmt.Errorf("configure job object: %w", err)
|
|
}
|
|
|
|
cmd.SysProcAttr = &syscall.SysProcAttr{CreationFlags: windows.CREATE_SUSPENDED}
|
|
if err := cmd.Start(); err != nil {
|
|
_ = windows.CloseHandle(job)
|
|
return err
|
|
}
|
|
|
|
process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(cmd.Process.Pid))
|
|
if err == nil {
|
|
err = windows.AssignProcessToJobObject(job, process)
|
|
_ = windows.CloseHandle(process)
|
|
}
|
|
if err == nil {
|
|
err = resumeInitialThread(uint32(cmd.Process.Pid))
|
|
}
|
|
if err != nil {
|
|
killErr := cmd.Process.Kill()
|
|
waitErr := cmd.Wait()
|
|
_ = windows.CloseHandle(job)
|
|
return joinErrors(fmt.Errorf("assign process to job object: %w", err), killErr, waitErr)
|
|
}
|
|
|
|
tree.job = job
|
|
return nil
|
|
}
|
|
|
|
func resumeInitialThread(processID uint32) error {
|
|
snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPTHREAD, 0)
|
|
if err != nil {
|
|
return fmt.Errorf("snapshot initial thread: %w", err)
|
|
}
|
|
defer func() { _ = windows.CloseHandle(snapshot) }()
|
|
|
|
entry := windows.ThreadEntry32{Size: uint32(unsafe.Sizeof(windows.ThreadEntry32{}))}
|
|
if err := windows.Thread32First(snapshot, &entry); err != nil {
|
|
return fmt.Errorf("find initial thread: %w", err)
|
|
}
|
|
for {
|
|
if entry.OwnerProcessID != processID {
|
|
// Keep enumerating until the suspended process's only initial thread
|
|
// is found.
|
|
} else {
|
|
thread, openErr := windows.OpenThread(windows.THREAD_SUSPEND_RESUME, false, entry.ThreadID)
|
|
if openErr != nil {
|
|
return fmt.Errorf("open initial thread: %w", openErr)
|
|
}
|
|
defer func() { _ = windows.CloseHandle(thread) }()
|
|
if _, resumeErr := windows.ResumeThread(thread); resumeErr != nil {
|
|
return fmt.Errorf("resume initial thread: %w", resumeErr)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
if err := windows.Thread32Next(snapshot, &entry); err != nil {
|
|
if errors.Is(err, windows.ERROR_NO_MORE_FILES) {
|
|
break
|
|
}
|
|
return fmt.Errorf("find initial thread: %w", err)
|
|
}
|
|
}
|
|
return fmt.Errorf("find initial thread: no thread found for process %d", processID)
|
|
}
|
|
|
|
func (tree *windowsProcessTree) TerminateGracefully() error {
|
|
// Windows jobs have no portable graceful signal. Terminating the owned job
|
|
// is the safe fallback and prevents a descendant from escaping cleanup.
|
|
return tree.terminate()
|
|
}
|
|
|
|
func (tree *windowsProcessTree) TerminateForcefully() error {
|
|
return tree.terminate()
|
|
}
|
|
|
|
func (tree *windowsProcessTree) Dispose() error {
|
|
if tree.job == 0 {
|
|
return nil
|
|
}
|
|
err := windows.CloseHandle(tree.job)
|
|
tree.job = 0
|
|
return err
|
|
}
|
|
|
|
func (tree *windowsProcessTree) terminate() error {
|
|
if tree.job == 0 {
|
|
return nil
|
|
}
|
|
return windows.TerminateJobObject(tree.job, 1)
|
|
}
|