rvbox/internal/client/supervisor/windows/native_exec.go

package windows

// This file contains the platform-neutral process bookkeeping shared by the
// native Windows adapter and the deterministic non-Windows test adapter. The
// Windows build supplies the token/Job-backed start and signal operations in
// native_windows.go; keeping output and stdin semantics here prevents those
// paths from drifting apart.

import (
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
	"github.com/rvbox/rvbox/internal/client/supervisor"
	"github.com/rvbox/rvbox/internal/domain"
)

var (
	ErrProcessAlreadyRunning = errors.New("a process for this command is already running")
	ErrProcessNotFound       = errors.New("supervised process was not found")
	ErrProcessNotReady       = errors.New("supervised process is not ready")
)

// NativeOptions is the common policy input. Windows callers additionally get
// token/session selection and Job Object containment in the native build; the
// test adapter uses the same shell and output limits without OS handles.
type NativeOptions struct {
	Shells           ShellPaths
	WorkRoot         string
	JobProfiles      map[string]JobProfile
	MaxWrapperBytes  uint64
	MaxOutputChunk   uint64
	WindowsTermGrace time.Duration
	Now              func() time.Time
}

// JobProfile is the validated administrator policy for one protocol profile.
// The native Windows adapter applies only controls it can set and read back
// atomically on a Job Object; a requested unsupported required control rejects
// the command before any child is created.
type JobProfile struct {
	RequiredControls []string
	CPUPercent       uint64
	MemoryMaxBytes   uint64
	PIDsMax          uint64
	IOReadBPS        uint64
	IOWriteBPS       uint64
}

func (options NativeOptions) withDefaults() NativeOptions {
	if options.MaxWrapperBytes == 0 {
		options.MaxWrapperBytes = 10 << 20
	}
	if options.MaxOutputChunk == 0 || options.MaxOutputChunk > 64<<10 {
		options.MaxOutputChunk = 64 << 10
	}
	if options.WindowsTermGrace <= 0 {
		options.WindowsTermGrace = 10 * time.Second
	}
	if options.Now == nil {
		options.Now = func() time.Time { return time.Now().UTC() }
	}
	return options
}

type execSupervisor struct {
	options NativeOptions
	mu      sync.Mutex
	active  map[domain.UUID]*execProcess
}

func newExecSupervisor(options NativeOptions) *execSupervisor {
	options = options.withDefaults()
	return &execSupervisor{options: options, active: make(map[domain.UUID]*execProcess)}
}

type execProcess struct {
	issue      domain.UUID
	identity   supervisor.EffectiveIdentity
	cmd        *exec.Cmd
	pid        uint32
	stdin      io.WriteCloser
	stdout     io.ReadCloser
	stderr     io.ReadCloser
	outputs    chan outputResult
	done       chan struct{}
	started    time.Time
	waitFn     func() (int32, bool, error)
	killFn     func(uint32) error
	snapshotFn func() (supervisor.ResourceSnapshot, error)

	mu       sync.Mutex
	finished bool
	status   supervisor.ExitStatus
	waitErr  error
	closeIn  sync.Once
}

type outputResult struct {
	chunk supervisor.OutputChunk
	err   error
}

func (process *execProcess) IssueUUID() domain.UUID { return process.issue }

func (process *execProcess) Identity() supervisor.EffectiveIdentity { return process.identity }

func (process *execProcess) ReadOutput(ctx context.Context) (supervisor.OutputChunk, error) {
	if process == nil || process.outputs == nil {
		return supervisor.OutputChunk{}, ErrProcessNotReady
	}
	select {
	case <-ctx.Done():
		return supervisor.OutputChunk{}, ctx.Err()
	case result, ok := <-process.outputs:
		if !ok {
			return supervisor.OutputChunk{}, io.EOF
		}
		return result.chunk, result.err
	}
}

func (process *execProcess) Wait(ctx context.Context) (supervisor.ExitStatus, error) {
	if process == nil || process.done == nil {
		return supervisor.ExitStatus{}, ErrProcessNotReady
	}
	select {
	case <-ctx.Done():
		return supervisor.ExitStatus{}, ctx.Err()
	case <-process.done:
		process.mu.Lock()
		defer process.mu.Unlock()
		return process.status, process.waitErr
	}
}

func (process *execProcess) WriteStdin(ctx context.Context, data []byte, appendNewline bool) error {
	if process == nil || process.stdin == nil {
		return ErrProcessNotReady
	}
	if appendNewline {
		data = append(append([]byte(nil), data...), stdinLineEnding()...)
	}
	return writeWithContext(ctx, process.stdin, data)
}

func (process *execProcess) CloseStdin(_ context.Context) error {
	if process == nil || process.stdin == nil {
		return ErrProcessNotReady
	}
	var err error
	process.closeIn.Do(func() { err = process.stdin.Close() })
	return err
}

func writeWithContext(ctx context.Context, writer io.Writer, data []byte) error {
	if len(data) == 0 {
		return nil
	}
	result := make(chan error, 1)
	go func() {
		_, err := writer.Write(data)
		result <- err
	}()
	select {
	case <-ctx.Done():
		return ctx.Err()
	case err := <-result:
		return err
	}
}

func (process *execProcess) setResult(status supervisor.ExitStatus, err error) {
	process.mu.Lock()
	process.status, process.waitErr, process.finished = status, err, true
	process.mu.Unlock()
	close(process.done)
}

func (process *execProcess) startReaders(maxChunk uint64, remove func()) {
	var readers sync.WaitGroup
	read := func(stream rvboxv1.StreamKind, source io.ReadCloser) {
		defer readers.Done()
		defer source.Close()
		limit := int(maxChunk)
		if limit <= 0 {
			limit = 64 << 10
		}
		buffer := make([]byte, limit)
		for {
			count, err := source.Read(buffer)
			if count > 0 {
				data := append([]byte(nil), buffer[:count]...)
				select {
				case process.outputs <- outputResult{chunk: supervisor.OutputChunk{Stream: stream, Data: data}}:
				default:
					// The output channel is bounded. A reader must never hold a
					// child pipe open while waiting for network I/O; dropping here
					// is surfaced as an explicit read error to the caller.
					select {
					case process.outputs <- outputResult{err: fmt.Errorf("%w: output channel full", supervisor.ErrUnsupported)}:
					case <-time.After(time.Second):
					}
				}
			}
			if err != nil {
				if !errors.Is(err, io.EOF) {
					select {
					case process.outputs <- outputResult{err: err}:
					default:
					}
				}
				return
			}
		}
	}
	readers.Add(2)
	go read(rvboxv1.StreamKind_STREAM_STDOUT, process.stdout)
	go read(rvboxv1.StreamKind_STREAM_STDERR, process.stderr)
	go func() {
		readers.Wait()
		close(process.outputs)
		if remove != nil {
			remove()
		}
	}()
}

func materializeWrapper(directory string, wrapper Wrapper, now time.Time) (string, func(), error) {
	if directory == "" {
		return "", nil, ErrInvalidWorkingDirectory
	}
	if err := os.MkdirAll(directory, 0o700); err != nil {
		return "", nil, err
	}
	temporary, err := os.CreateTemp(directory, ".rvbox-wrapper-*")
	if err != nil {
		return "", nil, err
	}
	temporaryName := temporary.Name()
	cleanup := func() {
		_ = temporary.Close()
		_ = os.Remove(temporaryName)
	}
	if err := temporary.Chmod(0o600); err != nil {
		cleanup()
		return "", nil, err
	}
	if _, err := temporary.Write(wrapper.Bytes); err != nil {
		cleanup()
		return "", nil, err
	}
	if err := temporary.Sync(); err != nil {
		cleanup()
		return "", nil, err
	}
	if err := temporary.Close(); err != nil {
		_ = os.Remove(temporaryName)
		return "", nil, err
	}
	finalName := temporaryName + wrapper.Extension
	if err := os.Rename(temporaryName, finalName); err != nil {
		_ = os.Remove(temporaryName)
		return "", nil, err
	}
	return finalName, func() { _ = os.Remove(finalName) }, nil
}

func parseEnvironment(values []string) map[string]string {
	result := make(map[string]string, len(values))
	for _, value := range values {
		index := strings.IndexByte(value, '=')
		if index <= 0 {
			continue
		}
		result[value[:index]] = value[index+1:]
	}
	return result
}

func (process *execProcess) terminate(code uint32) error {
	if process == nil {
		return ErrProcessNotReady
	}
	if process.killFn != nil {
		return process.killFn(code)
	}
	if process.cmd == nil || process.cmd.Process == nil {
		return ErrProcessNotReady
	}
	return process.cmd.Process.Kill()
}

func (manager *execSupervisor) remove(issue domain.UUID, process *execProcess) {
	manager.mu.Lock()
	if manager.active[issue] == process {
		delete(manager.active, issue)
	}
	manager.mu.Unlock()
}

func validateExecutionSource(spec *supervisor.StartSpec) error {
	if err := spec.Validate(); err != nil {
		return err
	}
	if spec.Execution.GetScript() != nil && len(spec.ScriptBody) == 0 && spec.Execution.GetScript().GetSizeBytes() != 0 {
		return ErrProcessNotReady
	}
	if len(spec.ExecutionProfiles) != 0 {
		return fmt.Errorf("%w: execution profiles are not supported by this adapter yet", supervisor.ErrUnsupported)
	}
	return nil
}

func sourceCommand(spec supervisor.StartSpec, wrapperPath string) (string, []string, error) {
	source := spec.Execution.GetCommandText()
	if wrapperPath != "" {
		source = wrapperPath
	}
	switch spec.Execution.GetShellType() {
	case rvboxv1.ShellType_SHELL_SH:
		if wrapperPath != "" {
			return "/bin/sh", []string{wrapperPath}, nil
		}
		return "/bin/sh", []string{"-c", source}, nil
	case rvboxv1.ShellType_SHELL_BASH:
		if wrapperPath != "" {
			return "/bin/bash", []string{wrapperPath}, nil
		}
		return "/bin/bash", []string{"-c", source}, nil
	default:
		return "", nil, ErrUnsupportedShell
	}
}

// startPortable launches the same bounded pipe contract on Unix. It is kept
// private because v1 does not advertise a Unix client; tests use it to prove
// runner/supervisor ordering without a Windows host.
func (manager *execSupervisor) startPortable(ctx context.Context, spec supervisor.StartSpec) (supervisor.Process, error) {
	if err := validateExecutionSource(&spec); err != nil {
		return nil, err
	}
	manager.mu.Lock()
	if _, exists := manager.active[spec.IssueUUID]; exists {
		manager.mu.Unlock()
		return nil, ErrProcessAlreadyRunning
	}
	manager.mu.Unlock()
	var wrapperPath string
	var cleanup func()
	if spec.Execution.GetScript() != nil {
		wrapper, err := BuildWrapper(spec.Execution, spec.ScriptBody, manager.options.MaxWrapperBytes)
		if err != nil {
			return nil, err
		}
		wrapperPath, cleanup, err = materializeWrapper(spec.WorkingDirectory, wrapper, manager.options.Now())
		if err != nil {
			return nil, err
		}
	}
	program, arguments, err := sourceCommand(spec, wrapperPath)
	if err != nil {
		if cleanup != nil {
			cleanup()
		}
		return nil, err
	}
	return manager.startCommand(ctx, spec, program, arguments, supervisor.EffectiveIdentity{Context: "CURRENT_PROCESS", Elevated: false}, cleanup, nil)
}

func (manager *execSupervisor) startCommand(ctx context.Context, spec supervisor.StartSpec, program string, arguments []string, identity supervisor.EffectiveIdentity, cleanup func(), configure func(*exec.Cmd) error) (supervisor.Process, error) {
	command := exec.CommandContext(ctx, program, arguments...)
	if configure != nil {
		if err := configure(command); err != nil {
			if cleanup != nil {
				cleanup()
			}
			return nil, err
		}
	}
	command.Dir = filepath.Clean(spec.WorkingDirectory)
	base := parseEnvironment(os.Environ())
	entries, err := MergeEnvironment(base, spec.Environment)
	if err != nil {
		if cleanup != nil {
			cleanup()
		}
		return nil, err
	}
	command.Env = make([]string, 0, len(entries))
	for _, entry := range entries {
		command.Env = append(command.Env, entry.Key+"="+entry.Value)
	}
	stdinRead, stdin, err := os.Pipe()
	if err != nil {
		if cleanup != nil {
			cleanup()
		}
		return nil, err
	}
	stdout, stdoutWrite, err := os.Pipe()
	if err != nil {
		_ = stdinRead.Close()
		_ = stdin.Close()
		if cleanup != nil {
			cleanup()
		}
		return nil, err
	}
	stderr, stderrWrite, err := os.Pipe()
	if err != nil {
		_ = stdinRead.Close()
		_ = stdin.Close()
		_ = stdout.Close()
		_ = stdoutWrite.Close()
		if cleanup != nil {
			cleanup()
		}
		return nil, err
	}
	command.Stdin = stdinRead
	command.Stdout = stdoutWrite
	command.Stderr = stderrWrite
	if err := command.Start(); err != nil {
		_ = stdinRead.Close()
		_ = stdin.Close()
		_ = stdout.Close()
		_ = stdoutWrite.Close()
		_ = stderr.Close()
		_ = stderrWrite.Close()
		if cleanup != nil {
			cleanup()
		}
		return nil, err
	}
	_ = stdinRead.Close()
	_ = stdoutWrite.Close()
	_ = stderrWrite.Close()
	started := manager.options.Now()
	waitFn := func() (int32, bool, error) {
		err := command.Wait()
		var code int32
		if command.ProcessState != nil {
			code = int32(command.ProcessState.ExitCode())
		}
		return code, false, err
	}
	return manager.registerProcess(spec.IssueUUID, identity, command, stdin, stdout, stderr, started, waitFn, func(uint32) error { return command.Process.Kill() }, cleanup), nil
}

func (manager *execSupervisor) registerProcess(issue domain.UUID, identity supervisor.EffectiveIdentity, command *exec.Cmd, stdin io.WriteCloser, stdout, stderr io.ReadCloser, started time.Time, waitFn func() (int32, bool, error), killFn func(uint32) error, cleanup func()) *execProcess {
	var pid uint32
	if command != nil && command.Process != nil && command.Process.Pid > 0 {
		pid = uint32(command.Process.Pid)
	}
	process := &execProcess{issue: issue, identity: identity, cmd: command, pid: pid, stdin: stdin, stdout: stdout, stderr: stderr, outputs: make(chan outputResult, 32), done: make(chan struct{}), started: started, waitFn: waitFn, killFn: killFn}
	manager.mu.Lock()
	manager.active[issue] = process
	manager.mu.Unlock()
	process.startReaders(manager.options.MaxOutputChunk, cleanup)
	go func() {
		code, signaled, err := process.waitFn()
		finished := manager.options.Now()
		status := supervisor.ExitStatus{StartedAt: started, FinishedAt: finished, OutputDrained: false, Code: code, Signaled: signaled}
		process.setResult(status, err)
		manager.remove(issue, process)
	}()
	return process
}