mesh-catalog/modules/lemurs/internal/desktop/run.go

package desktop

import (
	"bytes"
	"context"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"time"
)

// Bounds on a command a tool runs: well below the runtime's 30 s call limit, and an answer that
// fits in a tool's reply.
const (
	DefaultTimeout = 10 * time.Second
	MostOutput     = 256 << 10
)

// Result is what one command did.
type Result struct {
	Command   []string `json:"command"`
	Code      int      `json:"exit_code"`
	Stdout    string   `json:"stdout,omitempty"`
	Stderr    string   `json:"stderr,omitempty"`
	Truncated bool     `json:"truncated,omitempty"`
	TimedOut  bool     `json:"timed_out,omitempty"`
}

// OK is whether the command ran and exited 0.
func (r Result) OK() bool { return r.Code == 0 && !r.TimedOut }

// Err is the command's failure as an error naming it and what it said, or nil.
func (r Result) Err() error {
	if r.OK() {
		return nil
	}
	said := strings.TrimSpace(r.Stderr)
	if said == "" {
		said = strings.TrimSpace(r.Stdout)
	}
	if r.TimedOut {
		return fmt.Errorf("%s did not finish in time", strings.Join(r.Command, " "))
	}
	return fmt.Errorf("%s exited %d: %s", strings.Join(r.Command, " "), r.Code, said)
}

// Runner runs a command with an environment and answers what it did. Tools take one, so their
// tests replace the machine with a table of answers.
type Runner func(ctx context.Context, env []string, stdin []byte, name string, args ...string) Result

// Exec is the machine's Runner: the command in its own process group, ended with everything it
// started at the deadline, each stream cut at MostOutput.
func Exec(ctx context.Context, env []string, stdin []byte, name string, args ...string) Result {
	if _, ok := ctx.Deadline(); !ok {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, DefaultTimeout)
		defer cancel()
	}
	res := Result{Command: append([]string{name}, args...)}
	cmd := exec.Command(name, args...)
	cmd.Env = env
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
	if stdin != nil {
		cmd.Stdin = bytes.NewReader(stdin)
	}
	out, errb := &capped{}, &capped{}
	cmd.Stdout, cmd.Stderr = out, errb
	if err := cmd.Start(); err != nil {
		res.Code = 127
		res.Stderr = err.Error()
		return res
	}
	done := make(chan error, 1)
	go func() { done <- cmd.Wait() }()
	var err error
	select {
	case err = <-done:
	case <-ctx.Done():
		_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		err = <-done
		res.TimedOut = true
	}
	res.Stdout, res.Stderr = out.String(), errb.String()
	res.Truncated = out.cut || errb.cut
	var exit *exec.ExitError
	switch {
	case err == nil:
	case errors.As(err, &exit):
		res.Code = exit.ExitCode()
		if res.Code < 0 {
			res.Code = 128
		}
	default:
		res.Code = 1
		if res.Stderr == "" {
			res.Stderr = err.Error()
		}
	}
	return res
}

// capped keeps the first MostOutput bytes written to it. Its buffer is a field, not embedded: an
// embedded bytes.Buffer brings ReadFrom along, and io.Copy would use it and never call Write.
type capped struct {
	buf bytes.Buffer
	cut bool
}

func (c *capped) Write(p []byte) (int, error) {
	if room := MostOutput - c.buf.Len(); room < len(p) {
		if room > 0 {
			c.buf.Write(p[:room])
		}
		c.cut = true
		return len(p), nil
	}
	return c.buf.Write(p)
}

func (c *capped) String() string { return c.buf.String() }

// Desk is what a desktop tool needs: how to find the session, and how to run a command.
type Desk struct {
	Find func() (*Session, error)
	Run  Runner
	// Base is the environment a command starts from, before the session's words.
	Base []string
}

// Machine is the real Desk, preferring the named processes as the session's.
func Machine(prefer ...string) Desk {
	return Desk{
		Find: func() (*Session, error) { return Find(prefer...) },
		Run:  Exec,
		Base: os.Environ(),
	}
}

// InSession runs a command in the operator's session, or answers NoSession.
func (d Desk) InSession(ctx context.Context, name string, args ...string) (Result, *Session, error) {
	s, err := d.Find()
	if err != nil {
		return Result{}, nil, err
	}
	return d.Run(ctx, s.Env(d.Base), nil, name, args...), s, nil
}

// InSessionWith is InSession with standard input.
func (d Desk) InSessionWith(ctx context.Context, stdin []byte, name string, args ...string) (Result, *Session, error) {
	s, err := d.Find()
	if err != nil {
		return Result{}, nil, err
	}
	return d.Run(ctx, s.Env(d.Base), stdin, name, args...), s, nil
}

// Plain runs a command with the base environment: for what needs no session.
func (d Desk) Plain(ctx context.Context, name string, args ...string) Result {
	return d.Run(ctx, d.Base, nil, name, args...)
}

// AsUser runs a command with the account's own runtime directory and bus, and no display.
func (d Desk) AsUser(ctx context.Context, name string, args ...string) Result {
	return d.Run(ctx, UserEnv(d.Base, os.Getuid()), nil, name, args...)
}

// Launched is how a program was started in the session.
type Launched struct {
	Unit string `json:"unit,omitempty"`
	PID  int    `json:"pid,omitempty"`
	How  string `json:"how"`
}

// Launch starts a program in the operator's session that outlives the call and the runtime.
//
// **Not as a child of this process.** The runtime is a system service; everything it starts is in
// its control group, and the service manager ends that group whenever the runtime restarts — which
// is every push that changes it. So the program is handed to the account's own service manager as a
// transient unit (`systemd-run --user`), with the session's words set on it, and lives as long as the
// operator's user manager does. Without a user manager it is started detached as a last resort, and
// the answer says it will end with the runtime.
func (d Desk) Launch(ctx context.Context, s *Session, name string, argv ...string) (Launched, error) {
	if len(argv) == 0 {
		return Launched{}, errors.New("nothing to launch")
	}
	env := s.Env(d.Base)
	unit := "mesh-" + name + "-" + token()
	args := []string{"--user", "--collect", "--quiet", "--unit=" + unit}
	for _, w := range []string{"DISPLAY", "WAYLAND_DISPLAY", "XAUTHORITY", "XDG_SESSION_TYPE", "XDG_CURRENT_DESKTOP", "XDG_SESSION_DESKTOP", "I3SOCK", "SWAYSOCK"} {
		if v := lookup(env, w); v != "" {
			args = append(args, "--setenv="+w+"="+v)
		}
	}
	args = append(args, "--")
	args = append(args, argv...)
	res := d.Run(ctx, env, nil, "systemd-run", args...)
	if res.OK() {
		return Launched{Unit: unit, How: "a transient unit of the account's service manager; ends when it exits or when the operator logs out"}, nil
	}
	if s.Bus != "" {
		return Launched{}, res.Err()
	}
	cmd := exec.Command(argv[0], argv[1:]...)
	cmd.Env = env
	cmd.SysProcAttr = &syscall.SysProcAttr{Setsid: true}
	if err := cmd.Start(); err != nil {
		return Launched{}, err
	}
	pid := cmd.Process.Pid
	go func() { _ = cmd.Wait() }()
	return Launched{PID: pid, How: "detached from the runtime with no user manager to hand it to; it ends when the runtime restarts"}, nil
}

func lookup(env []string, name string) string {
	for i := len(env) - 1; i >= 0; i-- {
		if k, v, ok := strings.Cut(env[i], "="); ok && k == name {
			return v
		}
	}
	return ""
}

func token() string {
	b := make([]byte, 4)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}