From a691cc98e7e939b91cc704f36cb39e0330b64c87 Mon Sep 17 00:00:00 2001
From: Nicolas De Loof <nicolas.deloof@gmail.com>
Date: Wed, 31 Mar 2021 14:24:11 +0200
Subject: [PATCH] make windows "kill children process" code simpler

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>
---
 cli/mobycli/job_windows.go | 85 +++++++-------------------------------
 go.mod                     |  1 +
 2 files changed, 17 insertions(+), 69 deletions(-)

diff --git a/cli/mobycli/job_windows.go b/cli/mobycli/job_windows.go
index 39f93916e..f930e18f8 100644
--- a/cli/mobycli/job_windows.go
+++ b/cli/mobycli/job_windows.go
@@ -18,9 +18,9 @@ package mobycli
 
 import (
 	"fmt"
-	"os"
-	"syscall"
 	"unsafe"
+
+	"golang.org/x/sys/windows"
 )
 
 func init() {
@@ -29,81 +29,28 @@ func init() {
 	}
 }
 
-var (
-	kernel32 = syscall.NewLazyDLL("kernel32.dll")
-)
-
-type jobObjectExtendedLimitInformation struct {
-	BasicLimitInformation struct {
-		PerProcessUserTimeLimit uint64
-		PerJobUserTimeLimit     uint64
-		LimitFlags              uint32
-		MinimumWorkingSetSize   uintptr
-		MaximumWorkingSetSize   uintptr
-		ActiveProcessLimit      uint32
-		Affinity                uintptr
-		PriorityClass           uint32
-		SchedulingClass         uint32
-	}
-	IoInfo struct {
-		ReadOperationCount  uint64
-		WriteOperationCount uint64
-		OtherOperationCount uint64
-		ReadTransferCount   uint64
-		WriteTransferCount  uint64
-		OtherTransferCount  uint64
-	}
-	ProcessMemoryLimit    uintptr
-	JobMemoryLimit        uintptr
-	PeakProcessMemoryUsed uintptr
-	PeakJobMemoryUsed     uintptr
-}
-
 // killSubProcessesOnClose will ensure on windows that all child processes of the current process are killed if parent is killed.
 func killSubProcessesOnClose() error {
-	job, err := createJobObject()
+	job, err := windows.CreateJobObject(nil, nil)
 	if err != nil {
 		return err
 	}
-	info := jobObjectExtendedLimitInformation{}
-	info.BasicLimitInformation.LimitFlags = 0x2000
-	if err := setInformationJobObject(job, info); err != nil {
-		_ = syscall.CloseHandle(job)
+	info := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{
+		BasicLimitInformation: windows.JOBOBJECT_BASIC_LIMIT_INFORMATION{
+			LimitFlags: windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE,
+		},
+	}
+	if _, err := windows.SetInformationJobObject(
+		job,
+		windows.JobObjectExtendedLimitInformation,
+		uintptr(unsafe.Pointer(&info)),
+		uint32(unsafe.Sizeof(info))); err != nil {
+		_ = windows.CloseHandle(job)
 		return err
 	}
-	proc, err := syscall.GetCurrentProcess()
-	if err != nil {
-		_ = syscall.CloseHandle(job)
-		return err
-	}
-	if err := assignProcessToJobObject(job, proc); err != nil {
-		_ = syscall.CloseHandle(job)
+	if err := windows.AssignProcessToJobObject(job, windows.CurrentProcess()); err != nil {
+		_ = windows.CloseHandle(job)
 		return err
 	}
 	return nil
 }
-
-func createJobObject() (syscall.Handle, error) {
-	res, _, err := kernel32.NewProc("CreateJobObjectW").Call(uintptr(unsafe.Pointer(nil)), uintptr(unsafe.Pointer(nil)))
-	if res == 0 {
-		return syscall.InvalidHandle, os.NewSyscallError("CreateJobObject", err)
-	}
-	return syscall.Handle(res), nil
-}
-
-func setInformationJobObject(job syscall.Handle, info jobObjectExtendedLimitInformation) error {
-	infoClass := uint32(9)
-	res, _, err := kernel32.NewProc("SetInformationJobObject").Call(uintptr(job), uintptr(infoClass), uintptr(unsafe.Pointer(&info)), uintptr(uint32(unsafe.Sizeof(info))))
-	if res == 0 {
-		return os.NewSyscallError("SetInformationJobObject", err)
-	}
-	return nil
-}
-
-func assignProcessToJobObject(job syscall.Handle, process syscall.Handle) error {
-	res, _, err := kernel32.NewProc("AssignProcessToJobObject").Call(uintptr(job), uintptr(process))
-	if res == 0 {
-		return os.NewSyscallError("AssignProcessToJobObject", err)
-	}
-	return nil
-}
diff --git a/go.mod b/go.mod
index 48f222c93..5cf72ef34 100644
--- a/go.mod
+++ b/go.mod
@@ -56,6 +56,7 @@ require (
 	github.com/valyala/fasttemplate v1.2.1 // indirect
 	golang.org/x/oauth2 v0.0.0-20201109201403-9fd604954f58
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
+	golang.org/x/sys v0.0.0-20210119212857-b64e53b001e4
 	google.golang.org/grpc v1.33.2
 	google.golang.org/protobuf v1.25.0
 	gopkg.in/ini.v1 v1.62.0