osmedeus/internal/runner/runner.go

package runner

import (
	"context"
	"fmt"
	"strings"
	"sync"

	"github.com/j3ssie/osmedeus/v5/internal/core"
)

// MaxOutputSize limits the combined output size to prevent memory issues
// with very large command outputs (10MB default)
const MaxOutputSize = 10 * 1024 * 1024

// MaxStderrSize limits stderr capture (1MB) since stderr is typically small
const MaxStderrSize = 1 * 1024 * 1024

// LimitedBuffer wraps bytes.Buffer and silently discards writes after maxSize.
// This prevents unbounded memory growth when a tool produces GBs of output.
type LimitedBuffer struct {
	buf      []byte
	maxSize  int
	overflow bool
	mu       sync.Mutex
}

// NewLimitedBuffer creates a LimitedBuffer with the given max size.
func NewLimitedBuffer(maxSize int) *LimitedBuffer {
	return &LimitedBuffer{
		maxSize: maxSize,
	}
}

// Write implements io.Writer. Writes beyond maxSize are silently discarded.
func (lb *LimitedBuffer) Write(p []byte) (n int, err error) {
	lb.mu.Lock()
	defer lb.mu.Unlock()

	if len(lb.buf) >= lb.maxSize {
		lb.overflow = true
		return len(p), nil // pretend we wrote it all
	}

	remaining := lb.maxSize - len(lb.buf)
	if len(p) > remaining {
		lb.buf = append(lb.buf, p[:remaining]...)
		lb.overflow = true
		return len(p), nil
	}

	lb.buf = append(lb.buf, p...)
	return len(p), nil
}

// Bytes returns the buffered bytes.
func (lb *LimitedBuffer) Bytes() []byte {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	return lb.buf
}

// Len returns the number of buffered bytes.
func (lb *LimitedBuffer) Len() int {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	return len(lb.buf)
}

// Overflow returns true if any writes were discarded.
func (lb *LimitedBuffer) Overflow() bool {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	return lb.overflow
}

// combineOutput efficiently combines stdout and stderr from LimitedBuffers.
// Appends "[output truncated]" if either buffer overflowed.
func combineOutput(stdout, stderr *LimitedBuffer) string {
	stdoutBytes := stdout.Bytes()
	stderrBytes := stderr.Bytes()
	totalLen := len(stdoutBytes) + len(stderrBytes)

	if totalLen == 0 {
		return ""
	}

	truncated := stdout.Overflow() || stderr.Overflow()

	var sb strings.Builder
	if truncated {
		sb.Grow(totalLen + 20)
	} else {
		sb.Grow(totalLen)
	}
	sb.Write(stdoutBytes)
	sb.Write(stderrBytes)
	if truncated {
		sb.WriteString("\n[output truncated]")
	}
	return sb.String()
}

// CommandResult holds the output of a command execution
type CommandResult struct {
	Output   string // Combined stdout and stderr
	ExitCode int    // Exit code of the command
	Error    error  // Error if execution failed
}

// PIDCallback is called when a process starts or ends
type PIDCallback func(pid int)

// Runner interface for executing commands in different environments
type Runner interface {
	// Execute runs a command and returns the result
	Execute(ctx context.Context, command string) (*CommandResult, error)

	// Setup prepares the runner (e.g., copy binary, start container, establish SSH)
	Setup(ctx context.Context) error

	// Cleanup tears down resources (e.g., stop container, close SSH connection)
	Cleanup(ctx context.Context) error

	// Type returns the runner type
	Type() core.RunnerType

	// IsRemote returns true if commands run on a remote machine
	IsRemote() bool

	// CopyFromRemote copies a file from the remote environment to the local host
	// For Docker: uses docker cp, for SSH: uses rsync
	CopyFromRemote(ctx context.Context, remotePath, localPath string) error

	// SetPIDCallbacks sets callbacks for process lifecycle events.
	// onStart is called when a process starts (with the PID)
	// onEnd is called when the process ends (with the PID)
	// This enables tracking of running processes for cancellation support.
	SetPIDCallbacks(onStart, onEnd PIDCallback)
}

// NewRunner creates a runner based on workflow configuration
func NewRunner(workflow *core.Workflow, binaryPath string) (Runner, error) {
	runnerType := workflow.Runner
	if runnerType == "" {
		runnerType = core.RunnerTypeHost
	}

	config := workflow.RunnerConfig
	if config == nil {
		config = &core.RunnerConfig{}
	}

	switch runnerType {
	case core.RunnerTypeHost:
		return NewHostRunner(binaryPath), nil
	case core.RunnerTypeDocker:
		return NewDockerRunner(config, binaryPath)
	case core.RunnerTypeSSH:
		return NewSSHRunner(config, binaryPath)
	default:
		return nil, fmt.Errorf("unknown runner type: %s", runnerType)
	}
}

// NewRunnerFromType creates a runner from type and config directly
func NewRunnerFromType(runnerType core.RunnerType, config *core.RunnerConfig, binaryPath string) (Runner, error) {
	if config == nil {
		config = &core.RunnerConfig{}
	}

	switch runnerType {
	case core.RunnerTypeHost, "":
		return NewHostRunner(binaryPath), nil
	case core.RunnerTypeDocker:
		return NewDockerRunner(config, binaryPath)
	case core.RunnerTypeSSH:
		return NewSSHRunner(config, binaryPath)
	default:
		return nil, fmt.Errorf("unknown runner type: %s", runnerType)
	}
}