108 lines
2.9 KiB
Go
108 lines
2.9 KiB
Go
//go:build windows
|
|
|
|
package compose
|
|
|
|
import (
|
|
"errors"
|
|
"os/exec"
|
|
"sync"
|
|
"unsafe"
|
|
|
|
"golang.org/x/sys/windows"
|
|
)
|
|
|
|
var ownedJobs = struct {
|
|
sync.Mutex
|
|
m map[*exec.Cmd]windows.Handle
|
|
}{m: make(map[*exec.Cmd]windows.Handle)}
|
|
|
|
// CREATE_SUSPENDED closes the registration race: no child code can create a
|
|
// descendant until the real process HANDLE has been assigned to the job.
|
|
func configureOwnedProcess(c *exec.Cmd) {
|
|
c.SysProcAttr = &windows.SysProcAttr{CreationFlags: windows.CREATE_NEW_PROCESS_GROUP | windows.CREATE_SUSPENDED}
|
|
}
|
|
|
|
func registerOwnedProcess(c *exec.Cmd) error {
|
|
if c.Process == nil {
|
|
return errors.New("owned process has no process handle")
|
|
}
|
|
job, err := windows.CreateJobObject(nil, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
closeJob := true
|
|
defer func() {
|
|
if closeJob {
|
|
windows.CloseHandle(job)
|
|
}
|
|
}()
|
|
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 {
|
|
return err
|
|
}
|
|
// c.Process.Pid is used only to obtain a genuine process HANDLE; a PID is
|
|
// never passed to AssignProcessToJobObject.
|
|
process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE|windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(c.Process.Pid))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer windows.CloseHandle(process)
|
|
if err = windows.AssignProcessToJobObject(job, process); err != nil {
|
|
return err
|
|
}
|
|
thread, err := suspendedPrimaryThread(uint32(c.Process.Pid))
|
|
if err != nil {
|
|
_ = windows.TerminateJobObject(job, 1)
|
|
return err
|
|
}
|
|
_, resumeErr := windows.ResumeThread(thread)
|
|
windows.CloseHandle(thread)
|
|
if resumeErr != nil {
|
|
_ = windows.TerminateJobObject(job, 1)
|
|
return resumeErr
|
|
}
|
|
ownedJobs.Lock()
|
|
ownedJobs.m[c] = job
|
|
ownedJobs.Unlock()
|
|
closeJob = false
|
|
return nil
|
|
}
|
|
|
|
func suspendedPrimaryThread(pid uint32) (windows.Handle, error) {
|
|
snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPTHREAD, 0)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer windows.CloseHandle(snapshot)
|
|
entry := windows.ThreadEntry32{Size: uint32(unsafe.Sizeof(windows.ThreadEntry32{}))}
|
|
err = windows.Thread32First(snapshot, &entry)
|
|
for err == nil {
|
|
if entry.OwnerProcessID == pid {
|
|
return windows.OpenThread(windows.THREAD_SUSPEND_RESUME, false, entry.ThreadID)
|
|
}
|
|
err = windows.Thread32Next(snapshot, &entry)
|
|
}
|
|
return 0, err
|
|
}
|
|
|
|
func releaseOwnedProcess(c *exec.Cmd) {
|
|
ownedJobs.Lock()
|
|
h := ownedJobs.m[c]
|
|
delete(ownedJobs.m, c)
|
|
ownedJobs.Unlock()
|
|
if h != 0 {
|
|
windows.CloseHandle(h)
|
|
}
|
|
}
|
|
func terminateOwnedProcess(c *exec.Cmd) {
|
|
ownedJobs.Lock()
|
|
h := ownedJobs.m[c]
|
|
ownedJobs.Unlock()
|
|
if h != 0 {
|
|
_ = windows.TerminateJobObject(h, 1)
|
|
} else if c.Process != nil {
|
|
_ = c.Process.Kill()
|
|
}
|
|
}
|