spark-mcp/internal/executor/spark_submit.go

// Package executor runs the local spark-submit binary for the spark_submit Tool.
//
// It is the single place in the codebase that shells out to user-configured
// binaries. To prevent command injection the binary is taken verbatim from
// cluster.SparkSubmitExecuteBin and arguments are passed via exec.Command's
// slice form — never as a single string, never via "sh -c".
package executor

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"os/exec"
	"regexp"
	"strings"
	"time"
)

// stdoutCap caps spark-submit stdout/stderr to keep the per-Tool log file
// from blowing up on verbose jobs. 10 MiB is generous for the "Submitted
// application application_xxx" line and a few progress lines.
const stdoutCap = 10 << 20

// appIDPattern matches the YARN line printed on successful submission:
//
//	"Submitted application application_12345_0001"
var appIDPattern = regexp.MustCompile(`application_\d+_\d+`)

// SparkSubmitOpts is the resolved command line for a single spark-submit run.
//
// Binary must be an absolute path (validated at cluster create/update time).
// Args are the post-binary CLI args, already including any cluster-default
// args prepended by the caller.
type SparkSubmitOpts struct {
	Binary string
	Args   []string
	// Timeout is the wall-clock budget for the entire run. Zero means no
	// timeout. We use context.WithTimeout to enforce it.
	Timeout time.Duration
}

// Result is what spark_submit returns to the LLM.
type Result struct {
	AppID      string `json:"app_id,omitempty"`
	ExitCode   int    `json:"exit_code"`
	StdoutTail string `json:"stdout_tail"`
	StderrTail string `json:"stderr_tail"`
	DurationMS int64  `json:"duration_ms"`
}

// ErrBinaryEmpty is returned when the cluster has no spark-submit binary
// configured. The agent should call list_clusters to surface this.
var ErrBinaryEmpty = errors.New("executor: spark_submit_execute_bin is empty")

// Run executes the configured spark-submit and returns the parsed result.
//
// Command construction: exec.Command(name, args...) — the slice form means
// no shell, so a malicious arg like "foo; rm -rf /" is passed as a single
// argument to spark-submit, never executed.
func Run(ctx context.Context, opts SparkSubmitOpts) (*Result, error) {
	if strings.TrimSpace(opts.Binary) == "" {
		return nil, ErrBinaryEmpty
	}

	if opts.Timeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, opts.Timeout)
		defer cancel()
	}

	cmd := exec.CommandContext(ctx, opts.Binary, opts.Args...)
	// Critical: do NOT call cmd.Run() with /bin/sh. The slice form above is
	// the entire injection defense.

	stdout := &capBuf{cap: stdoutCap}
	stderr := &capBuf{cap: stdoutCap}
	cmd.Stdout = stdout
	cmd.Stderr = stderr

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

	res := &Result{
		ExitCode:   exitCodeFromErr(err),
		StdoutTail: stdout.String(),
		StderrTail: stderr.String(),
		DurationMS: dur.Milliseconds(),
	}

	// app_id: prefer stdout, fall back to stderr. Both are searched only on
	// the tail (cap-bounded) so we don't pay for the full output.
	if id := matchAppID(res.StdoutTail); id != "" {
		res.AppID = id
	} else if id := matchAppID(res.StderrTail); id != "" {
		res.AppID = id
	}

	if err != nil {
		// Wrap the underlying error but keep the parsed result so the LLM
		// can see exit_code and stderr even on failure.
		return res, fmt.Errorf("executor: spark-submit failed: %w", err)
	}
	return res, nil
}

func matchAppID(s string) string {
	if s == "" {
		return ""
	}
	return appIDPattern.FindString(s)
}

// exitCodeFromErr returns 0 on nil, the process's exit code on *exec.ExitError,
// and 1 on other errors (e.g. context deadline). exec.ExitError is the typed
// way to ask the OS-level exit code without inspecting strings.
func exitCodeFromErr(err error) int {
	if err == nil {
		return 0
	}
	var exitErr *exec.ExitError
	if errors.As(err, &exitErr) {
		return exitErr.ExitCode()
	}
	// context.DeadlineExceeded or failed to start — surface as 1.
	return 1
}

// capBuf is an io.Writer that keeps at most `cap` bytes from the end.
//
// The point isn't to deliver partial stdout to the LLM (LLM rarely needs the
// full stream) — it's to bound memory and the per-Tool log file. We keep
// the tail because that's where "Submitted application application_xxx" and
// final error messages live.
type capBuf struct {
	cap int
	buf bytes.Buffer
}

func (b *capBuf) Write(p []byte) (int, error) {
	n, err := b.buf.Write(p)
	if err != nil {
		return n, err
	}
	if b.buf.Len() > b.cap {
		// Drop the front. Cheaper than a ring buffer; works fine at 10 MiB.
		overflow := b.buf.Len() - b.cap
		_ = b.buf.Next(overflow)
	}
	return len(p), nil
}

func (b *capBuf) String() string { return b.buf.String() }

// Ensure io.Discard stays referenced; we keep this in case future variants
// want to write to a logger instead of capturing.
var _ = io.Discard