narratio/internal/adapters/subprocess/process_tree_windows.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
	}

	var assignErr error
	withHandleErr := cmd.Process.WithHandle(func(handle uintptr) {
		assignErr = windows.AssignProcessToJobObject(job, windows.Handle(handle))
	})
	if withHandleErr != nil {
		return failWindowsProcessStart(cmd, job, fmt.Errorf("acquire stable process handle for job assignment: %w", withHandleErr))
	}
	if assignErr != nil {
		return failWindowsProcessStart(cmd, job, fmt.Errorf("assign process to job object: %w", assignErr))
	}
	if err := resumeInitialThread(uint32(cmd.Process.Pid)); err != nil {
		return failWindowsProcessStart(cmd, job, fmt.Errorf("resume initial process thread: %w", err))
	}

	tree.job = job
	return nil
}

func failWindowsProcessStart(cmd *exec.Cmd, job windows.Handle, setupErr error) error {
	killErr := cmd.Process.Kill()
	waitErr := cmd.Wait()
	var exitErr *exec.ExitError
	if killErr == nil && errors.As(waitErr, &exitErr) {
		// A killed suspended child normally reports a non-zero exit from Wait;
		// successful reaping is not an additional setup failure.
		waitErr = nil
	}
	return joinErrors(
		setupErr,
		wrapWindowsProcessCleanupError("kill suspended child", killErr),
		wrapWindowsProcessCleanupError("reap suspended child", waitErr),
		wrapWindowsProcessCleanupError("close unowned job object", windows.CloseHandle(job)),
	)
}

func wrapWindowsProcessCleanupError(operation string, err error) error {
	if err == nil {
		return nil
	}
	return fmt.Errorf("%s: %w", operation, err)
}

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)
}