monlet/agent/internal/runner/runner.go

package runner

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os/exec"
	"strconv"
	"strings"
	"syscall"
	"time"
	"unicode/utf8"
)

var replacementRune = []byte("<22>")

const (
	MaxOutputBytes  = 8192
	truncMarkerFmt  = "...[truncated %d bytes]"
	limitFailPrefix = "__MONLET_RESOURCE_LIMIT_FAILED__:"
)

type Status string

const (
	StatusOK       Status = "ok"
	StatusWarning  Status = "warning"
	StatusCritical Status = "critical"
	StatusUnknown  Status = "unknown"
)

type Result struct {
	Status               Status
	ExitCode             int
	DurationMs           int64
	Output               string
	OutputTruncated      bool
	ResourceLimitFailure string
}

// ResourceLimits contains best-effort shell ulimit values for a check command.
type ResourceLimits struct {
	CPUTime     time.Duration
	MemoryBytes int64
	OpenFiles   int
}

func (l ResourceLimits) IsZero() bool {
	return l.CPUTime == 0 && l.MemoryBytes == 0 && l.OpenFiles == 0
}

// Run executes argv with timeout. Stdout+stderr are captured into a single
// buffer; both are truncated to 8 KiB UTF-8 bytes if exceeded.
func Run(ctx context.Context, argv []string, timeout time.Duration) Result {
	return runArgv(ctx, argv, timeout)
}

// RunCommand executes a shell command with optional resource limits.
func RunCommand(ctx context.Context, command string, timeout time.Duration, limits ResourceLimits) Result {
	argv := []string{"/bin/sh", "-c", command}
	if !limits.IsZero() {
		argv = limitedShellArgv(command, limits)
	}
	return runArgv(ctx, argv, timeout)
}

func runArgv(ctx context.Context, argv []string, timeout time.Duration) Result {
	start := time.Now()
	cctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	cmd := exec.CommandContext(cctx, argv[0], argv[1:]...)
	cmd.SysProcAttr = sysProcAttr()
	cmd.Cancel = func() error { return killGroup(cmd) }
	cmd.WaitDelay = 2 * time.Second

	var buf bytes.Buffer
	cmd.Stdout = &buf
	cmd.Stderr = &buf

	err := cmd.Run()
	dur := time.Since(start).Milliseconds()

	out, truncated := TruncateUTF8(buf.Bytes(), MaxOutputBytes)
	r := Result{
		DurationMs:      dur,
		Output:          out,
		OutputTruncated: truncated,
	}

	if cctx.Err() == context.DeadlineExceeded {
		r.Status = StatusCritical
		r.ExitCode = 2
		r.Output = fmt.Sprintf("critical: check timed out after %s", timeout)
		return r
	}
	if isExitCode(err, 125) {
		if resource, ok := detectLimitFailure(out); ok {
			r.Status = StatusUnknown
			r.ExitCode = -1
			r.Output = "unknown: failed to apply resource limit: " + resource
			r.OutputTruncated = false
			r.ResourceLimitFailure = resource
			return r
		}
	}
	if isSignal(err, syscall.SIGXCPU) {
		r.Status = StatusCritical
		r.ExitCode = 2
		r.Output = "critical: CPU time limit exceeded"
		r.OutputTruncated = false
		return r
	}
	if err != nil {
		var ee *exec.ExitError
		if errors.As(err, &ee) {
			r.ExitCode = ee.ExitCode()
		} else {
			r.Status = StatusUnknown
			r.ExitCode = -1
			// Launch failure (e.g. shell missing); surface exec error, never check output.
			r.Output, r.OutputTruncated = TruncateUTF8([]byte("unknown: failed to start check: "+err.Error()), MaxOutputBytes)
			return r
		}
	}
	r.Status = MapExit(r.ExitCode)
	return r
}

func isExitCode(err error, code int) bool {
	var ee *exec.ExitError
	return errors.As(err, &ee) && ee.ExitCode() == code
}

func limitedShellArgv(command string, limits ResourceLimits) []string {
	cpu := ""
	if limits.CPUTime > 0 {
		cpu = strconv.FormatInt(int64((limits.CPUTime+time.Second-1)/time.Second), 10)
	}
	mem := ""
	if limits.MemoryBytes > 0 {
		mem = strconv.FormatInt((limits.MemoryBytes+1023)/1024, 10)
	}
	openFiles := ""
	if limits.OpenFiles > 0 {
		openFiles = strconv.Itoa(limits.OpenFiles)
	}
	script := `
cmd=$1
cpu=$2
mem=$3
open_files=$4
if [ -n "$cpu" ]; then ulimit -t "$cpu" 2>/dev/null || { printf '%s\n' '` + limitFailPrefix + `cpu_time'; exit 125; }; fi
if [ -n "$mem" ]; then ulimit -v "$mem" 2>/dev/null || { printf '%s\n' '` + limitFailPrefix + `memory'; exit 125; }; fi
if [ -n "$open_files" ]; then ulimit -n "$open_files" 2>/dev/null || { printf '%s\n' '` + limitFailPrefix + `open_files'; exit 125; }; fi
exec /bin/sh -c "$cmd"
`
	return []string{"/bin/sh", "-c", script, "monlet-limit-wrapper", command, cpu, mem, openFiles}
}

func detectLimitFailure(output string) (string, bool) {
	for _, line := range strings.Split(output, "\n") {
		line = strings.TrimSpace(line)
		if strings.HasPrefix(line, limitFailPrefix) {
			resource := strings.TrimPrefix(line, limitFailPrefix)
			if resource != "cpu_time" && resource != "memory" && resource != "open_files" {
				resource = "unknown"
			}
			return resource, true
		}
	}
	return "", false
}

func isSignal(err error, sig syscall.Signal) bool {
	var ee *exec.ExitError
	if !errors.As(err, &ee) {
		return false
	}
	ws, ok := ee.Sys().(syscall.WaitStatus)
	return ok && ws.Signaled() && ws.Signal() == sig
}

func MapExit(code int) Status {
	switch code {
	case 0:
		return StatusOK
	case 1:
		return StatusWarning
	case 2:
		return StatusCritical
	default:
		return StatusUnknown
	}
}

// TruncateUTF8 returns a string whose UTF-8 byte length is <= max and a
// truncation flag. If truncated, appends "...[truncated N bytes]" marker.
func TruncateUTF8(b []byte, max int) (string, bool) {
	// Sanitize invalid bytes first; otherwise json.Marshal replaces them with
	// U+FFFD (3 bytes each) after our size check and the wire payload can
	// exceed the contract limit.
	b = bytes.ToValidUTF8(b, replacementRune)
	if len(b) <= max {
		return string(b), false
	}
	budget := max - 32
	if budget < 0 {
		budget = 0
	}
	for i := 0; i < 5; i++ {
		cut := safeBoundary(b, budget)
		marker := fmt.Sprintf(truncMarkerFmt, len(b)-cut)
		if cut+len(marker) <= max {
			return string(b[:cut]) + marker, true
		}
		budget = max - len(marker) - 1
		if budget < 0 {
			budget = 0
		}
	}
	cut := safeBoundary(b, max/2)
	marker := fmt.Sprintf(truncMarkerFmt, len(b)-cut)
	return string(b[:cut]) + marker, true
}

func safeBoundary(b []byte, limit int) int {
	if limit >= len(b) {
		return len(b)
	}
	for i := limit; i > 0; i-- {
		if utf8.RuneStart(b[i]) {
			if utf8.Valid(b[:i]) {
				return i
			}
		}
	}
	return 0
}

func sysProcAttr() *syscall.SysProcAttr {
	return &syscall.SysProcAttr{Setpgid: true}
}

func killGroup(cmd *exec.Cmd) error {
	if cmd.Process == nil {
		return nil
	}
	pgid, err := syscall.Getpgid(cmd.Process.Pid)
	if err == nil {
		_ = syscall.Kill(-pgid, syscall.SIGKILL)
		return nil
	}
	return cmd.Process.Kill()
}