main.go

// Copyright ©2023 Dan Kortschak. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// The runner executable is a dex module for running arbitrary programs
// using dex.
package main

import (
	"bytes"
	"context"
	"errors"
	"flag"
	"fmt"
	"log/slog"
	"net"
	"os"
	"os/exec"
	"slices"
	"strings"
	"sync"
	"sync/atomic"
	"time"
	"unique"

	"github.com/kortschak/jsonrpc2"

	runner "github.com/kortschak/dex/cmd/runner/api"
	"github.com/kortschak/dex/config"
	"github.com/kortschak/dex/internal/slogext"
	"github.com/kortschak/dex/internal/version"
	"github.com/kortschak/dex/rpc"
)

// Exit status codes.
const (
	success       = 0
	internalError = 1 << (iota - 1)
	invocationError
)

func main() { os.Exit(Main()) }

func Main() int {
	network := flag.String("network", "", "network for communication (unix or tcp)")
	addr := flag.String("addr", "", "address for communication")
	uid := flag.String("uid", "", "unique ID")
	logging := flag.String("log", "info", "logging level (debug, info, warn or error)")
	lines := flag.Bool("lines", false, "display source line details in logs")
	logStdout := flag.Bool("log_stdout", false, "log to stdout instead of stderr")
	v := flag.Bool("version", false, "print version and exit")
	flag.Parse()
	if *v {
		err := version.Print()
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(internalError)
		}
		os.Exit(success)
	}

	switch *network {
	case "unix", "tcp":
	default:
		flag.Usage()
		return invocationError
	}

	switch "" {
	case *addr, *uid:
		flag.Usage()
		return invocationError
	default:
	}

	var level slog.LevelVar
	err := level.UnmarshalText([]byte(*logging))
	if err != nil {
		flag.Usage()
		return invocationError
	}
	addSource := slogext.NewAtomicBool(*lines)
	logDst := os.Stderr
	if *logStdout {
		logDst = os.Stdout
	}
	log := slog.New(slogext.GoID{Handler: slogext.NewJSONHandler(logDst, &slogext.HandlerOptions{
		Level:     &level,
		AddSource: addSource,
	})}).With(
		slog.String("component", *uid),
	)

	ctx, cancel := context.WithCancel(context.Background())
	waitParent(func() {
		log.LogAttrs(ctx, slog.LevelError, "dex died")
		cancel()
	})

	h := newDaemon(*uid, log, &level, addSource, cancel)
	err = h.dial(ctx, *network, *addr, net.Dialer{})
	if err != nil {
		log.LogAttrs(ctx, slog.LevelError, err.Error())
		return internalError
	}
	defer h.close()

	log.LogAttrs(ctx, slog.LevelInfo, "start")
	<-ctx.Done()
	log.LogAttrs(ctx, slog.LevelInfo, "exit")

	return success
}

func waitParent(fn func()) {
	go func() {
		os.Stdin.Read([]byte{0})
		fn()
	}()
}

func newDaemon(uid string, log *slog.Logger, level *slog.LevelVar, addSource *atomic.Bool, cancel context.CancelFunc) *daemon {
	return &daemon{
		uid:       uid,
		log:       log,
		level:     level,
		addSource: addSource,
		cancel:    cancel,
		waiting:   make(map[*exec.Cmd]context.CancelFunc),
		buttons:   make(buttons),
	}
}

func (d *daemon) dial(ctx context.Context, network, addr string, dialer net.Dialer) error {
	d.log.LogAttrs(ctx, slog.LevelDebug, "dial", slog.String("network", network), slog.String("addr", addr))
	var err error
	d.conn, err = jsonrpc2.Dial(ctx, jsonrpc2.NetDialer(network, addr, dialer), d)
	if err != nil {
		return err
	}
	return nil
}

func (d *daemon) close() error {
	d.conn.Notify(context.Background(), rpc.Unregister, rpc.NewMessage(rpc.UID{Module: d.uid}, rpc.None{}))
	return d.conn.Close()
}

type daemon struct {
	uid string

	// conn is the connection to the kernel.
	conn *jsonrpc2.Connection

	log       *slog.Logger
	level     *slog.LevelVar
	addSource *atomic.Bool
	cancel    context.CancelFunc

	wMu     sync.Mutex
	waiting map[*exec.Cmd]context.CancelFunc
	buttons buttons

	hMu       sync.Mutex
	heartbeat time.Duration
	hStop     chan struct{}
}

// buttons is the set of running jobs, stratified on the service,
// the action that was invoked and from where, and their process IDs.
type buttons map[rpc.UID]map[buttonAction]map[int]running

type running struct {
	Start time.Time `json:"start"`
	cmd   *exec.Cmd
}

func (b buttons) add(svc rpc.UID, action buttonAction, pid int, cmd *exec.Cmd) {
	r := running{Start: time.Now(), cmd: cmd}
	if b[svc] == nil {
		b[svc] = make(map[buttonAction]map[int]running)
	}
	if b[svc][action] == nil {
		b[svc][action] = make(map[int]running)
	}
	b[svc][action][pid] = r
}

func (b buttons) delete(svc rpc.UID, action buttonAction, pid int) {
	delete(b[svc][action], pid)
	if len(b[svc][action]) == 0 {
		delete(b[svc], action)
	}
	if len(b[svc]) == 0 {
		delete(b, svc)
	}
}

// buttonAction is a command action/location map key.
type buttonAction struct {
	loc unique.Handle[rpc.Button]
	cmd string // cmd is the string representation of the command.
}

func (d *daemon) Bind(ctx context.Context, conn *jsonrpc2.Connection) jsonrpc2.ConnectionOptions {
	d.conn = conn
	d.log.LogAttrs(ctx, slog.LevelDebug, "bind")
	return jsonrpc2.ConnectionOptions{
		Handler: d,
	}
}

func (d *daemon) Handle(ctx context.Context, req *jsonrpc2.Request) (any, error) {
	if d.log.Enabled(ctx, slog.LevelDebug-1) {
		d.log.LogAttrs(ctx, slog.LevelDebug-1, "handle", slog.Any("req", slogext.Request{Request: req}))
	} else {
		d.log.LogAttrs(ctx, slog.LevelDebug, "handle", slog.Any("req", slogext.RequestRedactPrivate{Request: req}))
	}

	switch req.Method {
	case rpc.Who:
		version, err := version.String()
		if err != nil {
			version = err.Error()
		}
		return rpc.NewMessage(rpc.UID{Module: d.uid}, version), nil

	case runner.State:
		var m rpc.Message[runner.StateRequest]
		err := rpc.UnmarshalMessage(req.Params, &m)
		if err != nil {
			d.log.LogAttrs(ctx, slog.LevelError, "state", slog.Any("req", slogext.RequestRedactPrivate{Request: req}), slog.Any("error", err))
			return nil, err
		}
		d.wMu.Lock()
		var running runner.StateResponse
		if m.Body.Service != nil {
			running = collectRunning(*m.Body.Service, m.Body.Buttons, d.buttons[*m.Body.Service])
		} else {
			for svc, buttons := range d.buttons {
				running = append(running, collectRunning(svc, m.Body.Buttons, buttons)...)
			}
		}
		d.wMu.Unlock()
		d.log.LogAttrs(ctx, slog.LevelDebug, "state", slog.Any("result", running))
		return rpc.NewMessage(m.UID, running), nil

	case runner.Run:
		var m rpc.Message[runner.Params]
		err := rpc.UnmarshalMessage(req.Params, &m)
		if err != nil {
			d.log.LogAttrs(ctx, slog.LevelError, "run", slog.Any("req", slogext.RequestRedactPrivate{Request: req}), slog.Any("error", err))
			return nil, err
		}
		typ := "notify"
		if req.IsCall() {
			typ = "call"
		}
		d.log.LogAttrs(ctx, slog.LevelDebug, "run", slog.Any("type", typ), slog.Any("details", m))

		p := m.Body
		if p.Path == "" {
			err = rpc.NewError(rpc.ErrCodeInvalidMessage,
				"missing executable path",
				map[string]any{
					"type": rpc.ErrCodeParameters,
				},
			)
			d.log.LogAttrs(ctx, slog.LevelError, err.Error())
			return nil, err
		}

		// Currently the parent context is cancelled for a notify
		// when the handle call completes, so we need a separate
		// context to control the child process.
		// We keep a list of waiting running notify-started children
		// to cancel when runner exits.
		cctx := context.Background()
		// Cannot defer cancel Timeout, otherwise notifies will be
		// cancelled on handle completion. Make sure it is called
		// in all paths and in the case of a notify, that it is
		// called after cmd.Wait returns.
		var cancel context.CancelFunc
		if p.Timeout > 0 {
			cctx, cancel = context.WithTimeout(cctx, p.Timeout)
		} else {
			cctx, cancel = context.WithCancel(cctx)
		}

		cmd := exec.CommandContext(cctx, p.Path, p.Args...)
		cmd.Env = p.Env
		cmd.Dir = p.Dir
		cmd.WaitDelay = p.WaitDelay
		if p.Stdin != "" {
			cmd.Stdin = strings.NewReader(p.Stdin)
		}
		d.log.LogAttrs(ctx, slog.LevelDebug, "run", slog.Any("command", cmd.String()))
		if !req.IsCall() {
			err = cmd.Start()
			if err != nil {
				cancel()
				return nil, rpc.NewError(rpc.ErrCodeInternal,
					err.Error(),
					map[string]any{
						"type": rpc.ErrCodeProcess,
					},
				)
			}
			action := buttonAction{
				cmd: cmd.String(),
			}
			if m.Button != nil {
				action.loc = unique.Make(*m.Button)
			}
			pid := cmd.Process.Pid
			go func() {
				d.wMu.Lock()
				d.waiting[cmd] = cancel
				d.buttons.add(m.UID, action, pid, cmd)
				d.wMu.Unlock()

				cmd.Wait()

				d.wMu.Lock()
				delete(d.waiting, cmd)
				d.buttons.delete(m.UID, action, pid)
				d.wMu.Unlock()

				cancel()
			}()
			return nil, nil
		}
		defer cancel() // Mop up remaining cases.
		var stdout, stderr bytes.Buffer
		cmd.Stdout = &stdout
		cmd.Stderr = &stderr
		err = cmd.Run()
		var exitErr *exec.ExitError
		if errors.As(err, &exitErr) {
			d.log.LogAttrs(ctx, slog.LevelError, err.Error(), slog.Any("cmd", p))
			return nil, rpc.NewError(rpc.ErrCodeInternal,
				err.Error(),
				map[string]any{
					"type":      rpc.ErrCodeProcess,
					"exit_code": exitErr.ExitCode(),
					"stderr":    stderr.String(),
				},
			)
		}
		resp := runner.Return{
			Stdout: stdout.String(),
			Stderr: stderr.String(),
		}
		if err != nil {
			resp.Err = err.Error()
		}
		d.log.LogAttrs(ctx, slog.LevelDebug, "run", slog.Any("cmd", p), slog.Any("result", resp))
		return rpc.NewMessage(m.UID, resp), nil

	case rpc.Configure:
		isService, err := config.IsService(req)
		if err != nil {
			d.log.LogAttrs(ctx, slog.LevelError, "configure", slog.Any("req", slogext.RequestRedactPrivate{Request: req}), slog.Any("error", err))
			return nil, err
		}
		var uid rpc.UID
		if isService {
			// This is currently only for config validation since service
			// configs are handled by the kernel for runner, being only
			// button actions. This may change.
			var m rpc.Message[runner.Service]
			err = rpc.UnmarshalMessage(req.Params, &m)
			if err != nil {
				d.log.LogAttrs(ctx, slog.LevelError, "configure", slog.Any("error", err))
				return nil, err
			}
			uid = m.UID
		} else {
			var m rpc.Message[runner.Config]
			err := rpc.UnmarshalMessage(req.Params, &m)
			if err != nil {
				d.log.LogAttrs(ctx, slog.LevelError, "configure", slog.Any("error", err))
				return nil, err
			}
			uid = m.UID

			if m.Body.LogLevel != nil {
				d.level.Set(*m.Body.LogLevel)
			}
			if m.Body.AddSource != nil {
				d.addSource.Store(*m.Body.AddSource)
			}
			d.log.LogAttrs(ctx, slog.LevelDebug, "configure", slog.Any("details", m))

			if m.Body.Options.Heartbeat != nil {
				d.beat(ctx, m.Body.Options.Heartbeat.Duration)
			}
		}

		if !req.ID.IsValid() {
			return nil, nil
		}
		return rpc.NewMessage(uid, "done"), nil

	case rpc.Stop:
		d.log.LogAttrs(ctx, slog.LevelInfo, "stop")

		// Clean up any still running children.
		d.wMu.Lock()
		for _, cancel := range d.waiting {
			cancel()
		}
		d.wMu.Unlock()

		d.cancel()
		return nil, nil

	default:
		return nil, jsonrpc2.ErrNotHandled
	}
}

func collectRunning(svc rpc.UID, buttons []rpc.Button, state map[buttonAction]map[int]running) []runner.Running {
	var running []runner.Running
	for action, proc := range state {
		loc := action.loc.Value()
		if buttons != nil && !slices.Contains(buttons, loc) {
			continue
		}
		for pid, r := range proc {
			running = append(running, runner.Running{
				Start:   r.Start,
				Service: svc,
				Button:  loc,
				Command: action.cmd,
				PID:     pid,
			})
		}
	}
	return running
}

func (d *daemon) beat(ctx context.Context, p time.Duration) {
	select {
	case <-ctx.Done():
		return
	default:
	}

	d.hMu.Lock()
	defer d.hMu.Unlock()

	if p == d.heartbeat {
		return
	}

	if d.hStop != nil {
		close(d.hStop)
		d.hStop = nil
	}

	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	switch {
	case p == 0:
		d.log.LogAttrs(ctx, slog.LevelInfo, "stop heartbeat")
		err := d.conn.Notify(ctx, "heartbeat",
			rpc.NewMessage(rpc.UID{Module: d.uid}, rpc.Deadline{}),
		)
		if err != nil {
			d.log.LogAttrs(ctx, slog.LevelError, "heartbeat", slog.Any("error", err))
		}
		d.heartbeat = p

	case p > 0:
		d.log.LogAttrs(ctx, slog.LevelInfo, "update heartbeat", slog.Duration("old", d.heartbeat), slog.Duration("new", p))
		stop := make(chan struct{})
		d.hStop = stop
		go func() {
			ticker := time.NewTicker(p)
			ctx, cancel := context.WithCancel(context.Background())
			defer cancel()
			for {
				select {
				case <-ctx.Done():
					ticker.Stop()
					return
				case <-stop:
					ticker.Stop()
					return
				case t := <-ticker.C:
					d.log.LogAttrs(ctx, slog.LevelDebug, "heartbeat", slog.Any("tick", t))
					deadline := t.Add(2 * p)
					err := d.conn.Notify(ctx, "heartbeat",
						rpc.Message[rpc.Deadline]{
							Time: t, UID: rpc.UID{Module: d.uid},
							Body: rpc.Deadline{Deadline: &deadline},
						},
					)
					if err != nil {
						d.log.LogAttrs(ctx, slog.LevelError, "heartbeat", slog.Any("error", err))
					}
				}
			}
		}()
		d.heartbeat = p

	default:
		d.log.LogAttrs(ctx, slog.LevelError, "update heartbeat", slog.Duration("invalid duration", p))
	}
}