go-shout/run.go

package main

import (
	"io"
	"os"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"syscall"
	"time"
)

const sentinelPrefix = "__SHOUT_SENTINEL_"

type RunOptions struct {
	CleanEnv  bool
	PathDirs  []string
	EnvVars   map[string]string
	Timeout   time.Duration
	Verbose   bool
	OnCommand func(Command)
}

func buildScript(commands []Command) string {
	var b strings.Builder
	b.WriteString("exec 2>&1\n")
	for i, cmd := range commands {
		b.WriteString(cmd.Cmd)
		b.WriteByte('\n')
		b.WriteString("printf '\\n" + sentinelPrefix + "%s_" + strconv.Itoa(i) + "__\\n' \"$?\"\n")
	}
	return b.String()
}

func parseSentinelOutput(raw string, commandCount int) (outputs [][]string, exitCodes []int) {
	re := regexp.MustCompile(regexp.QuoteMeta(sentinelPrefix) + `(\d+)_(\d+)__`)

	remaining := raw
	for i := 0; i < commandCount; i++ {
		loc := re.FindStringSubmatchIndex(remaining)
		if loc == nil {
			lines := strings.Split(remaining, "\n")
			if len(lines) > 0 && lines[0] == "" {
				lines = lines[1:]
			}
			lines = trimTrailingEmpty(lines)
			outputs = append(outputs, lines)
			exitCodes = append(exitCodes, 1)
			break
		}

		before := remaining[:loc[0]]
		exitCodeStr := remaining[loc[2]:loc[3]]
		ec, _ := strconv.Atoi(exitCodeStr)

		lines := strings.Split(before, "\n")
		if len(lines) > 0 && lines[0] == "" {
			lines = lines[1:]
		}
		lines = trimTrailingEmpty(lines)
		if len(lines) == 1 && lines[0] == "" {
			lines = nil
		}

		outputs = append(outputs, lines)
		exitCodes = append(exitCodes, ec)

		afterSentinel := remaining[loc[1]:]
		if strings.HasPrefix(afterSentinel, "\n") {
			afterSentinel = afterSentinel[1:]
		}
		remaining = afterSentinel
	}

	for len(outputs) < commandCount {
		outputs = append(outputs, nil)
		exitCodes = append(exitCodes, 1)
	}

	return outputs, exitCodes
}

func runFile(file ShoutFile, opts RunOptions) FileResult {
	tmpDir, err := os.MkdirTemp("", "shout-")
	if err != nil {
		return FileResult{File: file, TmpDir: "", Error: err.Error()}
	}

	if len(file.Commands) == 0 {
		return FileResult{File: file, TmpDir: tmpDir}
	}

	script := buildScript(file.Commands)

	// Build environment
	var envMap map[string]string
	if opts.CleanEnv {
		envMap = make(map[string]string)
	} else {
		envMap = make(map[string]string)
		for _, e := range os.Environ() {
			if k, v, ok := strings.Cut(e, "="); ok {
				envMap[k] = v
			}
		}
	}

	envMap["HOME"] = tmpDir
	envMap["SHOUT_DIR"] = tmpDir

	for k, v := range opts.EnvVars {
		envMap[k] = v
	}

	if len(opts.PathDirs) > 0 {
		existing := envMap["PATH"]
		envMap["PATH"] = strings.Join(opts.PathDirs, ":") + ":" + existing
	}

	envSlice := make([]string, 0, len(envMap))
	for k, v := range envMap {
		envSlice = append(envSlice, k+"="+v)
	}

	cmd := exec.Command("/bin/sh")
	cmd.Dir = tmpDir
	cmd.Env = envSlice
	cmd.Stderr = nil
	cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return FileResult{File: file, TmpDir: tmpDir, Error: err.Error()}
	}

	stdoutPipe, err := cmd.StdoutPipe()
	if err != nil {
		return FileResult{File: file, TmpDir: tmpDir, Error: err.Error()}
	}

	if err := cmd.Start(); err != nil {
		return FileResult{File: file, TmpDir: tmpDir, Error: err.Error()}
	}

	defer func() {
		if cmd.Process != nil {
			_ = syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
		}
	}()

	_, _ = io.WriteString(stdin, script)
	stdin.Close()

	// Read stdout with timeout
	totalTimeout := opts.Timeout * time.Duration(len(file.Commands))
	type readResult struct {
		data []byte
		err  error
	}
	ch := make(chan readResult, 1)
	go func() {
		data, err := io.ReadAll(stdoutPipe)
		ch <- readResult{data, err}
	}()

	var output []byte
	select {
	case r := <-ch:
		if r.err != nil {
			return FileResult{File: file, TmpDir: tmpDir, Error: r.err.Error()}
		}
		output = r.data
	case <-time.After(totalTimeout):
		return FileResult{File: file, TmpDir: tmpDir, Error: "Timeout reading output"}
	}

	_ = cmd.Wait()

	outputs, exitCodesList := parseSentinelOutput(string(output), len(file.Commands))

	results := make([]CommandResult, len(file.Commands))
	for i, c := range file.Commands {
		if opts.Verbose && opts.OnCommand != nil {
			opts.OnCommand(c)
		}
		actual := outputs[i]
		if actual == nil {
			actual = []string{}
		}
		results[i] = CommandResult{
			Command:  c,
			Actual:   actual,
			ExitCode: exitCodesList[i],
		}
	}

	return FileResult{File: file, Results: results, TmpDir: tmpDir}
}

func cleanupTmpDir(dir string) {
	os.RemoveAll(dir)
}