ThothII/tools/thothctl/internal/compose/runner.go

// Package compose executes Docker Compose through a fixed executable and argument arrays.
package compose

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os"
	"os/exec"
	"sync"
)

var ErrOutputLimit = errors.New("compose output limit exceeded")

type CaptureLimits struct {
	StdoutBytes int64
	StderrBytes int64
}
type Result struct {
	Stdout   string
	Stderr   string
	ExitCode int
}
type Runner struct{ binary string }

func NewRunner(binary string) Runner {
	if binary == "" {
		binary = "docker"
	}
	return Runner{binary: binary}
}
func (r Runner) Run(ctx context.Context, args []string, stdin io.Reader) (Result, error) {
	c := exec.CommandContext(ctx, r.binary, args...)
	c.Stdin = stdin
	var out, er bytes.Buffer
	c.Stdout = &out
	c.Stderr = &er
	err := c.Run()
	res := Result{Stdout: out.String(), Stderr: er.String()}
	if err == nil {
		return res, nil
	}
	var x *exec.ExitError
	if errors.As(err, &x) {
		res.ExitCode = x.ExitCode()
		return res, err
	}
	if errors.Is(err, exec.ErrNotFound) || errors.Is(err, os.ErrNotExist) {
		res.ExitCode = 127
		return res, fmt.Errorf("%w: %w", exec.ErrNotFound, err)
	}
	return res, err
}

// RunBounded streams each pipe into a preallocated fixed-capacity collector. It owns the
// process group and tears it down on the first overflow or cancellation.
func (r Runner) RunBounded(ctx context.Context, args []string, stdin io.Reader, limits CaptureLimits) (Result, error) {
	if limits.StdoutBytes < 0 || limits.StderrBytes < 0 {
		return Result{}, ErrOutputLimit
	}
	c := exec.Command(r.binary, args...)
	c.Stdin = stdin
	op, err := c.StdoutPipe()
	if err != nil {
		return Result{}, err
	}
	ep, err := c.StderrPipe()
	if err != nil {
		return Result{}, err
	}
	configureOwnedProcess(c)
	if err = c.Start(); err != nil {
		return Result{ExitCode: 127}, err
	}
	if err = registerOwnedProcess(c); err != nil {
		terminateOwnedProcess(c)
		_ = c.Wait()
		return Result{}, err
	}
	defer releaseOwnedProcess(c)
	type collector struct {
		buf      []byte
		n        int64
		overflow bool
	}
	out := collector{buf: make([]byte, limits.StdoutBytes)}
	er := collector{buf: make([]byte, limits.StderrBytes)}
	overflow := make(chan struct{}, 1)
	var wg sync.WaitGroup
	wg.Add(2)
	read := func(rd io.Reader, max int64, col *collector) {
		defer wg.Done()
		chunk := make([]byte, 32*1024)
		for {
			n, e := rd.Read(chunk)
			if n > 0 {
				remain := max - col.n
				if remain > 0 {
					take := int64(n)
					if take > remain {
						take = remain
					}
					copy(col.buf[col.n:col.n+take], chunk[:take])
					col.n += take
				}
				if int64(n) > remain {
					col.overflow = true
					select {
					case overflow <- struct{}{}:
					default:
					}
					return
				}
			}
			if e != nil {
				return
			}
		}
	}
	go read(op, limits.StdoutBytes, &out)
	go read(ep, limits.StderrBytes, &er)
	waited := make(chan error, 1)
	go func() { waited <- c.Wait() }()
	var waitErr error
	var cancelled error
	select {
	case waitErr = <-waited:
		terminateOwnedProcess(c)
	case <-overflow:
		terminateOwnedProcess(c)
		waitErr = <-waited
	case <-ctx.Done():
		cancelled = ctx.Err()
		terminateOwnedProcess(c)
		waitErr = <-waited
	}
	wg.Wait()
	res := Result{Stdout: string(out.buf[:out.n]), Stderr: string(er.buf[:er.n])}
	if out.overflow || er.overflow {
		return res, ErrOutputLimit
	}
	if cancelled != nil {
		return res, cancelled
	}
	if waitErr == nil {
		return res, nil
	}
	var ex *exec.ExitError
	if errors.As(waitErr, &ex) {
		res.ExitCode = ex.ExitCode()
		return res, waitErr
	}
	if errors.Is(waitErr, exec.ErrNotFound) || errors.Is(waitErr, os.ErrNotExist) {
		res.ExitCode = 127
		return res, fmt.Errorf("%w: %w", exec.ErrNotFound, waitErr)
	}
	return res, waitErr
}