ThothII/tools/tht/internal/compose/process_windows.go

//go:build windows

package compose

import (
	"errors"
	"os"
	"os/exec"
	"syscall"
	"time"
	"unsafe"

	"golang.org/x/sys/windows"
)

const finalTerminationBound = 2 * time.Second
const createNewProcessGroup = 0x00000200
const createSuspendedProcess = 0x00000004
const jobTerminationExitCode = 1

func configureProcess(command *exec.Cmd) {
	command.SysProcAttr = &syscall.SysProcAttr{
		CreationFlags: createNewProcessGroup | createSuspendedProcess,
	}
}

type windowsProcessJob struct {
	handle windows.Handle
}

func newWindowsProcessJob() (*windowsProcessJob, error) {
	handle, err := windows.CreateJobObject(nil, nil)
	if err != nil {
		return nil, err
	}
	limits := windows.JOBOBJECT_EXTENDED_LIMIT_INFORMATION{}
	limits.BasicLimitInformation.LimitFlags = windows.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
	if _, err := windows.SetInformationJobObject(
		handle,
		windows.JobObjectExtendedLimitInformation,
		uintptr(unsafe.Pointer(&limits)),
		uint32(unsafe.Sizeof(limits)),
	); err != nil {
		return nil, errors.Join(err, windows.CloseHandle(handle))
	}
	return &windowsProcessJob{handle: handle}, nil
}

func (job *windowsProcessJob) close() error {
	if job == nil || job.handle == 0 {
		return nil
	}
	handle := job.handle
	job.handle = 0
	return windows.CloseHandle(handle)
}

type windowsProcessLifetime struct {
	command *exec.Cmd
	job     *windowsProcessJob
}

// startProcess creates the Docker process suspended, assigns it to a kill-on-close Job Object,
// then resumes its one primary thread. The suspension closes the startup window in which a direct
// child could otherwise create descendants before the Job Object begins to own them.
func startProcess(command *exec.Cmd) (processLifetime, error) {
	job, err := newWindowsProcessJob()
	if err != nil {
		return nil, err
	}
	if err := command.Start(); err != nil {
		return nil, errors.Join(err, job.close())
	}
	if command.Process == nil {
		return nil, errors.Join(ErrProcessReap, abortWindowsStart(command, job))
	}
	if err := assignWindowsProcessToJob(job, command.Process.Pid); err != nil {
		return nil, errors.Join(err, abortWindowsStart(command, job))
	}
	if err := resumeWindowsProcess(command.Process.Pid); err != nil {
		return nil, errors.Join(err, abortWindowsStart(command, job))
	}
	return &windowsProcessLifetime{command: command, job: job}, nil
}

func (process *windowsProcessLifetime) terminate(done <-chan error) error {
	if process == nil || process.command == nil || process.command.Process == nil || process.job == nil || process.job.handle == 0 {
		return ErrProcessReap
	}
	return terminateProcessTree(
		done,
		func() error { return windows.TerminateJobObject(process.job.handle, jobTerminationExitCode) },
		process.command.Process.Kill,
		finalTerminationBound,
	)
}

func (process *windowsProcessLifetime) close() error {
	if process == nil {
		return nil
	}
	return process.job.close()
}

func assignWindowsProcessToJob(job *windowsProcessJob, pid int) error {
	process, err := windows.OpenProcess(windows.PROCESS_SET_QUOTA|windows.PROCESS_TERMINATE, false, uint32(pid))
	if err != nil {
		return err
	}
	return errors.Join(
		windows.AssignProcessToJobObject(job.handle, process),
		windows.CloseHandle(process),
	)
}

func resumeWindowsProcess(pid int) error {
	snapshot, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPTHREAD, 0)
	if err != nil {
		return err
	}
	threadID, threadErr := suspendedPrimaryThreadID(snapshot, uint32(pid))
	closeSnapshotErr := windows.CloseHandle(snapshot)
	if threadErr != nil || closeSnapshotErr != nil {
		return errors.Join(threadErr, closeSnapshotErr)
	}
	thread, err := windows.OpenThread(windows.THREAD_SUSPEND_RESUME, false, threadID)
	if err != nil {
		return err
	}
	_, resumeErr := windows.ResumeThread(thread)
	closeThreadErr := windows.CloseHandle(thread)
	return errors.Join(resumeErr, closeThreadErr)
}

func suspendedPrimaryThreadID(snapshot windows.Handle, pid uint32) (uint32, error) {
	entry := windows.ThreadEntry32{Size: uint32(unsafe.Sizeof(windows.ThreadEntry32{}))}
	if err := windows.Thread32First(snapshot, &entry); err != nil {
		return 0, err
	}
	threadIDs := make([]uint32, 0, 1)
	for {
		if entry.OwnerProcessID == pid {
			threadIDs = append(threadIDs, entry.ThreadID)
		}
		entry.Size = uint32(unsafe.Sizeof(windows.ThreadEntry32{}))
		err := windows.Thread32Next(snapshot, &entry)
		if errors.Is(err, windows.ERROR_NO_MORE_FILES) {
			break
		}
		if err != nil {
			return 0, err
		}
	}
	if len(threadIDs) != 1 {
		return 0, ErrProcessReap
	}
	return threadIDs[0], nil
}

// abortWindowsStart runs only while the direct process is still suspended or immediately after
// it was resumed unsuccessfully. It waits exactly once, so the normal runner path never races a
// second Wait or leaves a waiter goroutine behind.
func abortWindowsStart(command *exec.Cmd, job *windowsProcessJob) error {
	var cleanupErr error
	if job != nil && job.handle != 0 {
		cleanupErr = errors.Join(cleanupErr, windows.TerminateJobObject(job.handle, jobTerminationExitCode))
	}
	if command.Process != nil {
		if err := command.Process.Kill(); err != nil && !errors.Is(err, os.ErrProcessDone) {
			cleanupErr = errors.Join(cleanupErr, err)
		}
		if err := command.Wait(); err != nil {
			var exitError *exec.ExitError
			if !errors.As(err, &exitError) {
				cleanupErr = errors.Join(cleanupErr, err)
			}
		}
	}
	if job != nil {
		cleanupErr = errors.Join(cleanupErr, job.close())
	}
	return cleanupErr
}