iop/apps/agent/internal/clientprocess/process.go

package clientprocess

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"

	"iop/packages/go/agentconfig"
)

type osProcessBackend struct{}

type commandProcess struct {
	command  *exec.Cmd
	identity ProcessIdentity
	waitOnce sync.Once
	waitErr  error
}

func (osProcessBackend) Start(
	ctx context.Context,
	_ ClientKind,
	spec agentconfig.ClientProcessSpec,
) (OwnedProcess, error) {
	command := exec.CommandContext(ctx, spec.Executable, spec.Args...)
	command.Dir = spec.WorkingDirectory
	if err := command.Start(); err != nil {
		return nil, fmt.Errorf("clientprocess: start client: %w", err)
	}
	identity, err := currentProcessIdentity(command.Process.Pid)
	if err != nil {
		_ = command.Process.Kill()
		_ = command.Wait()
		return nil, fmt.Errorf("clientprocess: establish start identity: %w", err)
	}
	return &commandProcess{command: command, identity: identity}, nil
}

func (p *commandProcess) Identity() ProcessIdentity {
	return p.identity
}

func (p *commandProcess) Wait() error {
	p.waitOnce.Do(func() {
		p.waitErr = p.command.Wait()
	})
	return p.waitErr
}

func (p *commandProcess) Abort() error {
	killErr := p.command.Process.Kill()
	waitErr := p.Wait()
	if errors.Is(killErr, os.ErrProcessDone) {
		killErr = nil
	}
	return errors.Join(killErr, waitErr)
}

func (osProcessBackend) Inspect(
	_ context.Context,
	identity ProcessIdentity,
) (IdentityObservation, error) {
	current, err := currentProcessIdentity(identity.PID)
	if errors.Is(err, os.ErrProcessDone) {
		return IdentityObservation{State: IdentityExited}, nil
	}
	if err != nil {
		return IdentityObservation{State: IdentityAmbiguous}, err
	}
	if current.StartToken != identity.StartToken {
		return IdentityObservation{
			State:    IdentityStale,
			Identity: current,
		}, nil
	}
	return IdentityObservation{
		State:    IdentityLive,
		Identity: current,
	}, nil
}

func (backend osProcessBackend) Signal(
	ctx context.Context,
	identity ProcessIdentity,
	signal os.Signal,
) error {
	observation, err := backend.Inspect(ctx, identity)
	if err != nil || observation.State == IdentityAmbiguous {
		return ErrIdentityAmbiguous
	}
	if observation.State != IdentityLive {
		return os.ErrProcessDone
	}
	process, err := os.FindProcess(identity.PID)
	if err != nil {
		return err
	}
	return process.Signal(signal)
}

func (backend osProcessBackend) Kill(
	ctx context.Context,
	identity ProcessIdentity,
) error {
	return backend.Signal(ctx, identity, os.Kill)
}

func (osProcessBackend) Focus(
	ctx context.Context,
	kind ClientKind,
	spec agentconfig.ClientProcessSpec,
) error {
	if kind != ClientFlutter || len(spec.FocusArgs) == 0 {
		return ErrFocusUnsupported
	}
	command := exec.CommandContext(ctx, spec.Executable, spec.FocusArgs...)
	command.Dir = spec.WorkingDirectory
	if err := command.Start(); err != nil {
		return fmt.Errorf("clientprocess: start focus command: %w", err)
	}
	if err := command.Wait(); err != nil {
		return fmt.Errorf("clientprocess: focus command: %w", err)
	}
	return nil
}

func currentProcessIdentity(pid int) (ProcessIdentity, error) {
	if pid <= 0 {
		return ProcessIdentity{}, ErrIdentityAmbiguous
	}
	process, err := os.FindProcess(pid)
	if err != nil {
		return ProcessIdentity{}, err
	}
	if err := process.Signal(syscall.Signal(0)); err != nil {
		if errors.Is(err, os.ErrProcessDone) || errors.Is(err, syscall.ESRCH) {
			return ProcessIdentity{}, os.ErrProcessDone
		}
		return ProcessIdentity{}, err
	}
	token, err := processStartToken(pid)
	if err != nil {
		return ProcessIdentity{}, err
	}
	return ProcessIdentity{PID: pid, StartToken: token}, nil
}

func processStartToken(pid int) (string, error) {
	switch runtime.GOOS {
	case "linux":
		payload, err := os.ReadFile("/proc/" + strconv.Itoa(pid) + "/stat")
		if errors.Is(err, os.ErrNotExist) {
			return "", os.ErrProcessDone
		}
		if err != nil {
			return "", err
		}
		closeIndex := strings.LastIndexByte(string(payload), ')')
		if closeIndex < 0 || closeIndex+2 >= len(payload) {
			return "", ErrIdentityAmbiguous
		}
		fields := strings.Fields(string(payload[closeIndex+2:]))
		// The tail begins at field 3; field 22 (starttime) is index 19.
		if len(fields) <= 19 || fields[19] == "" {
			return "", ErrIdentityAmbiguous
		}
		return "linux:" + fields[19], nil
	case "darwin":
		output, err := exec.Command(
			"/bin/ps",
			"-o",
			"lstart=",
			"-p",
			strconv.Itoa(pid),
		).Output()
		if err != nil {
			if errors.Is(err, os.ErrProcessDone) {
				return "", os.ErrProcessDone
			}
			return "", err
		}
		value := strings.TrimSpace(string(output))
		if value == "" {
			return "", os.ErrProcessDone
		}
		sum := sha256.Sum256([]byte(value))
		return "darwin:" + hex.EncodeToString(sum[:]), nil
	default:
		return "", fmt.Errorf(
			"%w: process identity is unsupported on %s",
			ErrIdentityAmbiguous,
			runtime.GOOS,
		)
	}
}