do a round of refactoring & cleanup

checks/cli.go

@@ -1,43 +1,128 @@
 package checks
 
 import (
+	"bytes"
+	"context"
+	"errors"
 	"fmt"
 	"maps"
 	"os"
 	"os/exec"
 	"regexp"
 	"runtime"
 	"strings"
+	"time"
 
 	api "github.com/bootdotdev/bootdev/client"
 )
 
+const (
+	cliCommandTimeout          = 5 * time.Minute
+	maxCLIOutputBytesPerStream = 1024 * 1024
+	commandWaitDelay           = 2 * time.Second
+)
+
+var errCLIOutputLimitExceeded = errors.New("CLI command output limit exceeded")
+
+type boundedBuffer struct {
+	buffer     bytes.Buffer
+	limit      int
+	truncated  bool
+	onTruncate func()
+}
+
+func newBoundedBuffer(limit int, onTruncate func()) *boundedBuffer {
+	return &boundedBuffer{
+		limit:      max(limit, 0),
+		onTruncate: onTruncate,
+	}
+}
+
+func (b *boundedBuffer) Write(p []byte) (int, error) {
+	originalLength := len(p)
+	remaining := max(b.limit-b.buffer.Len(), 0)
+	toWrite := min(len(p), remaining)
+	if toWrite > 0 {
+		_, _ = b.buffer.Write(p[:toWrite])
+	}
+	if toWrite < len(p) && !b.truncated {
+		b.truncated = true
+		if b.onTruncate != nil {
+			b.onTruncate()
+		}
+	}
+	return originalLength, nil
+}
+
+func (b *boundedBuffer) String() string {
+	return b.buffer.String()
+}
+
+func (b *boundedBuffer) Truncated() bool {
+	return b.truncated
+}
+
 func runCLICommand(command api.CLIStepCLICommand, variables map[string]string) (result api.CLICommandResult) {
+	return runCLICommandWithLimits(command, variables, cliCommandTimeout, maxCLIOutputBytesPerStream)
+}
+
+func runCLICommandWithLimits(
+	command api.CLIStepCLICommand,
+	variables map[string]string,
+	timeout time.Duration,
+	maxOutputBytesPerStream int,
+) (result api.CLICommandResult) {
 	finalCommand := InterpolateVariables(command.Command, variables)
 	result.FinalCommand = finalCommand
 	result.Command = command
 
+	timeoutCtx, cancelTimeout := context.WithTimeout(context.Background(), timeout)
+	defer cancelTimeout()
+	ctx, cancelCommand := context.WithCancelCause(timeoutCtx)
+	defer cancelCommand(nil)
+
 	var cmd *exec.Cmd
 
 	if runtime.GOOS == "windows" {
-		cmd = exec.Command("powershell", "-Command", finalCommand)
+		cmd = exec.CommandContext(ctx, "powershell", "-Command", finalCommand)
 	} else {
-		cmd = exec.Command("sh", "-c", finalCommand)
+		cmd = exec.CommandContext(ctx, "sh", "-c", finalCommand)
 	}
 
 	cmd.Env = append(os.Environ(), "LANG=en_US.UTF-8")
-	b, err := cmd.CombinedOutput()
+	cmd.WaitDelay = commandWaitDelay
+	cancelForOutputLimit := func() {
+		cancelCommand(errCLIOutputLimitExceeded)
+	}
+	stdout := newBoundedBuffer(maxOutputBytesPerStream, cancelForOutputLimit)
+	stderr := newBoundedBuffer(maxOutputBytesPerStream, cancelForOutputLimit)
+	cmd.Stdout = stdout
+	cmd.Stderr = stderr
+	err := cmd.Run()
 	if ee, ok := err.(*exec.ExitError); ok {
 		result.ExitCode = ee.ExitCode()
 	} else if err != nil {
 		result.ExitCode = -2
 	}
-	result.Stdout = strings.TrimRight(string(b), " \n\t\r")
+	result.Stdout = strings.TrimRight(stdout.String(), " \n\t\r")
+	result.Stderr = strings.TrimRight(stderr.String(), " \n\t\r")
 	if command.StdoutFilterTmdl != nil {
 		result.Stdout = ExtractTmdlBlock(result.Stdout, *command.StdoutFilterTmdl)
 	}
-	if err := parseStdoutVariables(result.Stdout, command.StdoutVariables, variables); err != nil {
-		result.Err = err.Error()
+
+	switch {
+	case errors.Is(context.Cause(ctx), errCLIOutputLimitExceeded):
+		result.Err = fmt.Sprintf("command output exceeded the %d-byte per-stream limit", maxOutputBytesPerStream)
+		result.ExitCode = -2
+	case errors.Is(context.Cause(ctx), context.DeadlineExceeded):
+		result.Err = fmt.Sprintf("command timed out after %s", timeout)
+		result.ExitCode = -2
+	}
+
+	if result.Err == "" {
+		if err := parseStdoutVariables(result.Stdout, command.StdoutVariables, variables); err != nil {
+			result.Err = err.Error()
+		}
 	}
 	result.Variables = maps.Clone(variables)
 	return result

checks/cli_test.go

@@ -2,11 +2,102 @@ package checks
 
 import (
 	"runtime"
+	"strings"
 	"testing"
+	"time"
 
 	api "github.com/bootdotdev/bootdev/client"
 )
 
+func TestRunCLICommandTimesOut(t *testing.T) {
+	command := `while :; do :; done`
+	if runtime.GOOS == "windows" {
+		command = `while ($true) {}`
+	}
+
+	start := time.Now()
+	result := runCLICommandWithLimits(
+		api.CLIStepCLICommand{Command: command},
+		map[string]string{},
+		20*time.Millisecond,
+		1024,
+	)
+	elapsed := time.Since(start)
+
+	if !strings.Contains(result.Err, "command timed out") {
+		t.Fatalf("command error = %q, want timeout error", result.Err)
+	}
+	if result.ExitCode >= 0 {
+		t.Fatalf("exit code = %d, want internal failure", result.ExitCode)
+	}
+	if elapsed > time.Second {
+		t.Fatalf("command took %v, want a prompt timeout", elapsed)
+	}
+}
+
+func TestRunCLICommandCapsOutput(t *testing.T) {
+	command := `printf 'abcdefgh'; while :; do :; done`
+	if runtime.GOOS == "windows" {
+		command = `[Console]::Out.Write('abcdefgh'); while ($true) {}`
+	}
+
+	variables := map[string]string{}
+	start := time.Now()
+	result := runCLICommandWithLimits(
+		api.CLIStepCLICommand{
+			Command: command,
+			StdoutVariables: []api.CLICommandStdoutVariable{{
+				Name:  "partial",
+				Regex: `(abcd)`,
+			}},
+		},
+		variables,
+		5*time.Second,
+		4,
+	)
+	elapsed := time.Since(start)
+
+	if !strings.Contains(result.Err, "per-stream limit") {
+		t.Fatalf("command error = %q, want per-stream output limit error", result.Err)
+	}
+	if result.ExitCode >= 0 {
+		t.Fatalf("exit code = %d, want internal failure", result.ExitCode)
+	}
+	if result.Stdout != "abcd" {
+		t.Fatalf("stdout = %q, want capped output %q", result.Stdout, "abcd")
+	}
+	if _, ok := variables["partial"]; ok {
+		t.Fatal("truncated output unexpectedly populated a stdout variable")
+	}
+	if elapsed > time.Second {
+		t.Fatalf("command took %v, want cancellation immediately after exceeding the output limit", elapsed)
+	}
+}
+
+func TestBoundedBufferDiscardsExcessBytes(t *testing.T) {
+	truncations := 0
+	buffer := newBoundedBuffer(5, func() {
+		truncations++
+	})
+	written, err := buffer.Write([]byte("abcdefgh"))
+	if err != nil {
+		t.Fatalf("Write() error = %v", err)
+	}
+	if written != 8 {
+		t.Fatalf("Write() = %d, want 8", written)
+	}
+	if got := buffer.String(); got != "abcde" {
+		t.Fatalf("buffer = %q, want abcde", got)
+	}
+	if !buffer.Truncated() {
+		t.Fatal("buffer did not record truncation")
+	}
+	_, _ = buffer.Write([]byte("more"))
+	if truncations != 1 {
+		t.Fatalf("truncation callback invoked %d times, want once", truncations)
+	}
+}
+
 func TestRunCLICommandCapturesStdoutVariables(t *testing.T) {
 	variables := map[string]string{}
 	result := runCLICommand(api.CLIStepCLICommand{
@@ -28,6 +119,43 @@ func TestRunCLICommandCapturesStdoutVariables(t *testing.T) {
 	}
 }
 
+func TestRunCLICommandKeepsStderrSeparateFromStdoutChecks(t *testing.T) {
+	command := `printf 'stdout-value\n'; printf 'stderr-value\n' >&2`
+	if runtime.GOOS == "windows" {
+		command = `Write-Output 'stdout-value'; [Console]::Error.WriteLine('stderr-value')`
+	}
+
+	variables := map[string]string{}
+	step := api.CLIStepCLICommand{
+		Command: command,
+		StdoutVariables: []api.CLICommandStdoutVariable{{
+			Name:  "stderr_value",
+			Regex: `(stderr-value)`,
+		}},
+		Tests: []api.CLICommandTest{{
+			StdoutContainsAll: []string{"stderr-value"},
+		}},
+	}
+
+	result := runCLICommand(step, variables)
+
+	if result.Stdout != "stdout-value" {
+		t.Fatalf("stdout = %q, want stdout-value", result.Stdout)
+	}
+	if result.Stderr != "stderr-value" {
+		t.Fatalf("stderr = %q, want stderr-value", result.Stderr)
+	}
+	if strings.Contains(result.Stdout, "stderr-value") {
+		t.Fatalf("stdout unexpectedly contains stderr: %q", result.Stdout)
+	}
+	if _, ok := variables["stderr_value"]; ok {
+		t.Fatalf("stderr unexpectedly populated a stdout variable")
+	}
+	if failure := evaluateCLICommandTests(0, step, result); failure == nil {
+		t.Fatal("stderr unexpectedly satisfied a stdout check")
+	}
+}
+
 func TestRunCLICommandInterpolatesCapturedStdoutVariables(t *testing.T) {
 	variables := map[string]string{}
 

checks/http.go

@@ -17,6 +17,11 @@ import (
 	"github.com/spf13/cobra"
 )
 
+const (
+	maxHTTPResponseBodyBytes = 1024 * 1024
+	maxBinaryBodyBytes       = 16 * 1024
+)
+
 func runHTTPRequest(
 	client *http.Client,
 	baseURL string,
@@ -31,12 +36,12 @@ func runHTTPRequest(
 
 	var req *http.Request
 	if requestStep.Request.BodyJSON != nil {
-		dat, err := json.Marshal(requestStep.Request.BodyJSON)
+		bodyJSON := interpolateJSONStrings(requestStep.Request.BodyJSON, variables)
+		dat, err := json.Marshal(bodyJSON)
 		cobra.CheckErr(err)
-		interpolatedBodyJSONStr := InterpolateVariables(string(dat), variables)
 		req, err = http.NewRequest(
 			requestStep.Request.Method, completeURL,
-			bytes.NewBuffer([]byte(interpolatedBodyJSONStr)),
+			bytes.NewReader(dat),
 		)
 		if err != nil {
 			cobra.CheckErr("Failed to create request")
@@ -93,7 +98,7 @@ func runHTTPRequest(
 	}
 	defer resp.Body.Close()
 
-	body, err := io.ReadAll(resp.Body)
+	body, err := io.ReadAll(io.LimitReader(resp.Body, maxHTTPResponseBodyBytes+1))
 	if err != nil {
 		result = api.HTTPRequestResult{Err: "Failed to read response body"}
 		return result
@@ -109,7 +114,8 @@ func runHTTPRequest(
 		trailers[k] = strings.Join(v, ",")
 	}
 
-	if err := parseVariables(body, requestStep.ResponseVariables, variables); err != nil {
+	bodyString := truncateAndStringifyBody(body)
+	if err := parseVariables([]byte(bodyString), requestStep.ResponseVariables, variables); err != nil {
 		return api.HTTPRequestResult{Err: fmt.Sprintf("Failed to parse response variable: %s", err)}
 	}
 	if err := parseHeaderVariables(headers, requestStep.ResponseHeaderVariables, variables); err != nil {
@@ -120,13 +126,34 @@ func runHTTPRequest(
 		StatusCode:       resp.StatusCode,
 		ResponseHeaders:  headers,
 		ResponseTrailers: trailers,
-		BodyString:       truncateAndStringifyBody(body),
+		BodyString:       bodyString,
 		Variables:        maps.Clone(variables),
 		Request:          requestStep,
 	}
 	return result
 }
 
+func interpolateJSONStrings(value any, variables map[string]string) any {
+	switch value := value.(type) {
+	case string:
+		return InterpolateVariables(value, variables)
+	case []any:
+		interpolated := make([]any, len(value))
+		for i, item := range value {
+			interpolated[i] = interpolateJSONStrings(item, variables)
+		}
+		return interpolated
+	case map[string]any:
+		interpolated := make(map[string]any, len(value))
+		for key, item := range value {
+			interpolated[key] = interpolateJSONStrings(item, variables)
+		}
+		return interpolated
+	default:
+		return value
+	}
+}
+
 func prettyPrintHTTPTest(test api.HTTPRequestTest, variables map[string]string) string {
 	if test.StatusCode != nil {
 		return fmt.Sprintf("Expecting status code: %d", *test.StatusCode)
@@ -187,9 +214,9 @@ func prettyPrintHTTPTest(test api.HTTPRequestTest, variables map[string]string)
 // embedded in binary (e.g. "moov" in MP4 files) remain searchable by downstream checks.
 // The result is not guaranteed to be valid UTF-8 or lossless.
 func truncateAndStringifyBody(body []byte) string {
-	maxBodyLength := 1024 * 1024 // 1 MiB
+	maxBodyLength := maxHTTPResponseBodyBytes
 	if likelyBinary(body) {
-		maxBodyLength = 16 * 1024 // 16 KiB
+		maxBodyLength = maxBinaryBodyBytes
 	}
 	if len(body) > maxBodyLength {
 		body = body[:maxBodyLength]