155 lines
3.2 KiB
Go
155 lines
3.2 KiB
Go
package runner
|
||
|
||
import (
|
||
"bytes"
|
||
"context"
|
||
"errors"
|
||
"fmt"
|
||
"os/exec"
|
||
"syscall"
|
||
"time"
|
||
"unicode/utf8"
|
||
)
|
||
|
||
var replacementRune = []byte("<22>")
|
||
|
||
const (
|
||
MaxOutputBytes = 8192
|
||
truncMarkerFmt = "...[truncated %d bytes]"
|
||
)
|
||
|
||
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
|
||
}
|
||
|
||
// 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 {
|
||
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 = StatusUnknown
|
||
r.ExitCode = -1
|
||
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
|
||
return r
|
||
}
|
||
}
|
||
r.Status = MapExit(r.ExitCode)
|
||
return r
|
||
}
|
||
|
||
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()
|
||
}
|