diff --git a/internal/tools/bash.go b/internal/tools/bash.go index 090594fd..2f7b3696 100644 --- a/internal/tools/bash.go +++ b/internal/tools/bash.go @@ -18,6 +18,40 @@ import ( const defaultBashTimeoutMS = 120000 const maxBashTimeoutMS = 600000 +const maxBashOutputBytes = 16 * 1024 * 1024 // 16 MiB - prevent OOM on large command output + +// cappedWriter wraps a bytes.Buffer and stops accepting bytes after max, +// setting a flag instead. Used to cap stdout/stderr capture during +// command.Run() instead of only truncating afterward. +// +// The buffer is a named field rather than embedded on purpose: embedding +// *bytes.Buffer would promote its ReadFrom method, and exec.Cmd's Run() +// drives non-*os.File Stdout/Stderr writers through io.Copy, which prefers +// io.ReaderFrom over Write when the destination implements it. That would +// let the child process's output flow straight into the buffer via the +// promoted ReadFrom, bypassing the cap entirely and defeating the whole +// point of this type. +type cappedWriter struct { + buf *bytes.Buffer + max int + truncated *bool +} + +func (c *cappedWriter) Write(p []byte) (int, error) { + if c.buf.Len()+len(p) > c.max { + *c.truncated = true + remain := c.max - c.buf.Len() + if remain > 0 { + c.buf.Write(p[:remain]) + } + return len(p), nil + } + return c.buf.Write(p) +} + +func (c *cappedWriter) String() string { return c.buf.String() } +func (c *cappedWriter) Len() int { return c.buf.Len() } + type bashTool struct { baseTool workspaceRoot string @@ -124,10 +158,12 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS defer plan.Cleanup() addSandboxMeta(meta, plan) - var stdout bytes.Buffer - var stderr bytes.Buffer - command.Stdout = &stdout - command.Stderr = &stderr + stdoutTruncated := false + stderrTruncated := false + stdout := &cappedWriter{buf: &bytes.Buffer{}, max: maxBashOutputBytes, truncated: &stdoutTruncated} + stderr := &cappedWriter{buf: &bytes.Buffer{}, max: maxBashOutputBytes, truncated: &stderrTruncated} + command.Stdout = stdout + command.Stderr = stderr // Kill the shell as a process group on timeout and bound the post-kill I/O // wait, so a backgrounded child cannot outlive the command or hang Run(). @@ -159,7 +195,7 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS markLikelySandboxDenial(meta, plan, exitCode, stdout.String(), stderrText) return Result{ Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta), + Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta, stdoutTruncated, stderrTruncated), Meta: meta, } } @@ -168,13 +204,13 @@ func (tool bashTool) run(ctx context.Context, args map[string]any, engine *zeroS if meta[SandboxLikelyDeniedMeta] == "true" { return Result{ Status: StatusError, - Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta), + Output: formatBashOutputWithShellHint(commandText, stdout.String(), stderrText, exitCode, meta, stdoutTruncated, stderrTruncated), Meta: meta, } } return Result{ Status: StatusOK, - Output: formatBashOutput(stdout.String(), stderrText, exitCode), + Output: formatBashOutputWithTruncation(stdout.String(), stderrText, exitCode, stdoutTruncated, stderrTruncated), Meta: meta, } } @@ -319,6 +355,10 @@ func commandExitCode(err error) int { } func formatBashOutput(stdout string, stderr string, exitCode int) string { + return formatBashOutputWithTruncation(stdout, stderr, exitCode, false, false) +} + +func formatBashOutputWithTruncation(stdout string, stderr string, exitCode int, stdoutTruncated, stderrTruncated bool) string { parts := []string{} stdout = strings.TrimRight(stdout, "\r\n") stderr = strings.TrimRight(stderr, "\r\n") @@ -339,14 +379,20 @@ func formatBashOutput(stdout string, stderr string, exitCode int) string { if n := len(outFindings) + len(errFindings); n > 0 { parts = append(parts, fmt.Sprintf("[zero] redacted %d likely secret(s) from this output before showing it.", n)) } + if stdoutTruncated { + parts = append(parts, fmt.Sprintf("[zero] output truncated at %d MiB", maxBashOutputBytes/(1024*1024))) + } + if stderrTruncated { + parts = append(parts, fmt.Sprintf("[zero] stderr truncated at %d MiB", maxBashOutputBytes/(1024*1024))) + } if len(parts) == 0 { return "Command completed with no output." } return strings.Join(parts, "\n") } -func formatBashOutputWithShellHint(command string, stdout string, stderr string, exitCode int, meta map[string]string) string { - output := formatBashOutput(stdout, stderr, exitCode) +func formatBashOutputWithShellHint(command string, stdout string, stderr string, exitCode int, meta map[string]string, stdoutTruncated, stderrTruncated bool) string { + output := formatBashOutputWithTruncation(stdout, stderr, exitCode, stdoutTruncated, stderrTruncated) if issue := detectShellOutputIssue(command, stdout+"\n"+stderr, runtime.GOOS); issue != nil { meta["shell_issue"] = issue.Kind output = appendShellIssueHint(output, *issue) diff --git a/internal/tools/bash_tool_test.go b/internal/tools/bash_tool_test.go index 11073d5d..b6ce848a 100644 --- a/internal/tools/bash_tool_test.go +++ b/internal/tools/bash_tool_test.go @@ -1,8 +1,10 @@ package tools import ( + "bytes" "context" "fmt" + "io" "net" "net/http" "os" @@ -47,6 +49,13 @@ func runBashToolHelper(command string) { case "long-sleep": time.Sleep(5 * time.Second) fmt.Println("long sleep finished") + case "large-output": + // Writes well beyond maxBashOutputBytes so tests can verify the + // capped writer stops growing memory instead of buffering it all. + chunk := bytes.Repeat([]byte("x"), 1024*1024) + for i := 0; i < 20; i++ { + os.Stdout.Write(chunk) + } case "http-server": listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { @@ -298,6 +307,95 @@ func TestBashToolReturnsNonzeroExitAsError(t *testing.T) { } } +func TestBashToolCapsLargeStdoutAndReportsTruncation(t *testing.T) { + result := NewBashTool(t.TempDir()).Run(context.Background(), map[string]any{ + "command": helperCommand("large-output"), + }) + + if result.Status != StatusOK { + t.Fatalf("expected ok status, got %s: %s", result.Status, result.Output) + } + wantNotice := fmt.Sprintf("output truncated at %d MiB", maxBashOutputBytes/(1024*1024)) + if !strings.Contains(result.Output, wantNotice) { + t.Fatalf("expected truncation notice %q, got output of length %d", wantNotice, len(result.Output)) + } + if len(result.Output) > maxBashOutputBytes+4096 { + t.Fatalf("expected output to stay bounded near the cap, got length %d", len(result.Output)) + } +} + +func TestCappedWriterStopsGrowingAfterMax(t *testing.T) { + truncated := false + writer := &cappedWriter{buf: &bytes.Buffer{}, max: 10, truncated: &truncated} + + n, err := writer.Write([]byte("0123456789")) + if err != nil || n != 10 { + t.Fatalf("first write = (%d, %v), want (10, nil)", n, err) + } + if truncated { + t.Fatalf("truncated flag set early, buffer exactly at max") + } + if writer.Len() != 10 { + t.Fatalf("buffer length = %d, want 10", writer.Len()) + } + + n, err = writer.Write([]byte("overflow")) + if err != nil || n != len("overflow") { + t.Fatalf("second write = (%d, %v), want (%d, nil)", n, err, len("overflow")) + } + if !truncated { + t.Fatalf("expected truncated flag to be set once max exceeded") + } + if writer.Len() != 10 { + t.Fatalf("buffer should not grow past max, got length %d", writer.Len()) + } +} + +func TestCappedWriterCapturesPartialWriteThatCrossesMax(t *testing.T) { + truncated := false + writer := &cappedWriter{buf: &bytes.Buffer{}, max: 10, truncated: &truncated} + + n, err := writer.Write([]byte("012345")) + if err != nil || n != 6 { + t.Fatalf("first write = (%d, %v), want (6, nil)", n, err) + } + + n, err = writer.Write([]byte("6789ABCDEF")) + if err != nil || n != 10 { + t.Fatalf("second write = (%d, %v), want (10, nil)", n, err) + } + if !truncated { + t.Fatalf("expected truncated flag to be set") + } + if got, want := writer.String(), "0123456789"; got != want { + t.Fatalf("buffer contents = %q, want %q", got, want) + } +} + +func TestCappedWriterEnforcesCapThroughIOCopy(t *testing.T) { + // exec.Cmd drives a non-*os.File Stdout/Stderr through io.Copy, which + // prefers io.ReaderFrom over Write when the destination implements it. + // This reproduces that path directly (regression test for the bug where + // embedding *bytes.Buffer promoted ReadFrom and bypassed the cap). + truncated := false + writer := &cappedWriter{buf: &bytes.Buffer{}, max: 10, truncated: &truncated} + + if _, ok := any(writer).(io.ReaderFrom); ok { + t.Fatalf("cappedWriter must not implement io.ReaderFrom, or io.Copy bypasses Write and the cap") + } + + source := bytes.NewReader([]byte("0123456789overflow")) + if _, err := io.Copy(writer, source); err != nil { + t.Fatalf("io.Copy returned error: %v", err) + } + if !truncated { + t.Fatalf("expected truncated flag to be set") + } + if writer.Len() != 10 { + t.Fatalf("expected io.Copy through cappedWriter to stay capped at 10 bytes, got %d", writer.Len()) + } +} + func TestBashToolTimesOut(t *testing.T) { result := NewBashTool(t.TempDir()).Run(context.Background(), map[string]any{ "command": helperCommand("sleep"),