os/exec: make Cmd.Output include stderr in ExitError

Change-Id: I3c6649d2f2521ab0843b13308569867d2e5f02da
Reviewed-on: https://go-review.googlesource.com/11415
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
Reviewed-by: Andrew Gerrand <adg@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

3 files changed, 186 insertions, 4 deletions

diff --git a/src/os/exec/exec.go b/src/os/exec/exec.go
index 8a84e263dc..a3ca98ce86 100644
--- a/src/os/exec/exec.go
+++ b/src/os/exec/exec.go
@@ -347,6 +347,18 @@ func (c *Cmd) Start() error {
 // An ExitError reports an unsuccessful exit by a command.
 type ExitError struct {
 	*os.ProcessState
+
+	// Stderr holds a subset of the standard error output from the
+	// Cmd.Output method if standard error was not otherwise being
+	// collected.
+	//
+	// If the error output is long, Stderr may contain only a prefix
+	// and suffix of the output, with the middle replaced with
+	// text about the number of omitted bytes.
+	//
+	// Stderr is provided for debugging, for inclusion in error messages.
+	// Users with other needs should redirect Cmd.Stderr as needed.
+	Stderr []byte
 }
 
 func (e *ExitError) Error() string {
@@ -392,21 +404,34 @@ func (c *Cmd) Wait() error {
 	if err != nil {
 		return err
 	} else if !state.Success() {
-		return &ExitError{state}
+		return &ExitError{ProcessState: state}
 	}
 
 	return copyError
 }
 
 // Output runs the command and returns its standard output.
+// Any returned error will usually be of type *ExitError.
+// If c.Stderr was nil, Output populates ExitError.Stderr.
 func (c *Cmd) Output() ([]byte, error) {
 	if c.Stdout != nil {
 		return nil, errors.New("exec: Stdout already set")
 	}
-	var b bytes.Buffer
-	c.Stdout = &b
+	var stdout bytes.Buffer
+	c.Stdout = &stdout
+
+	captureErr := c.Stderr == nil
+	if captureErr {
+		c.Stderr = &prefixSuffixSaver{N: 32 << 10}
+	}
+
 	err := c.Run()
-	return b.Bytes(), err
+	if err != nil && captureErr {
+		if ee, ok := err.(*ExitError); ok {
+			ee.Stderr = c.Stderr.(*prefixSuffixSaver).Bytes()
+		}
+	}
+	return stdout.Bytes(), err
 }
 
 // CombinedOutput runs the command and returns its combined standard
@@ -514,3 +539,80 @@ func (c *Cmd) StderrPipe() (io.ReadCloser, error) {
 	c.closeAfterWait = append(c.closeAfterWait, pr)
 	return pr, nil
 }
+
+// prefixSuffixSaver is an io.Writer which retains the first N bytes
+// and the last N bytes written to it. The Bytes() methods reconstructs
+// it with a pretty error message.
+type prefixSuffixSaver struct {
+	N         int // max size of prefix or suffix
+	prefix    []byte
+	suffix    []byte // ring buffer once len(suffix) == N
+	suffixOff int    // offset to write into suffix
+	skipped   int64
+
+	// TODO(bradfitz): we could keep one large []byte and use part of it for
+	// the prefix, reserve space for the '... Omitting N bytes ...' message,
+	// then the ring buffer suffix, and just rearrange the ring buffer
+	// suffix when Bytes() is called, but it doesn't seem worth it for
+	// now just for error messages. It's only ~64KB anyway.
+}
+
+func (w *prefixSuffixSaver) Write(p []byte) (n int, err error) {
+	lenp := len(p)
+	p = w.fill(&w.prefix, p)
+
+	// Only keep the last w.N bytes of suffix data.
+	if overage := len(p) - w.N; overage > 0 {
+		p = p[overage:]
+		w.skipped += int64(overage)
+	}
+	p = w.fill(&w.suffix, p)
+
+	// w.suffix is full now if p is non-empty. Overwrite it in a circle.
+	for len(p) > 0 { // 0, 1, or 2 iterations.
+		n := copy(w.suffix[w.suffixOff:], p)
+		p = p[n:]
+		w.skipped += int64(n)
+		w.suffixOff += n
+		if w.suffixOff == w.N {
+			w.suffixOff = 0
+		}
+	}
+	return lenp, nil
+}
+
+// fill appends up to len(p) bytes of p to *dst, such that *dst does not
+// grow larger than w.N. It returns the un-appended suffix of p.
+func (w *prefixSuffixSaver) fill(dst *[]byte, p []byte) (pRemain []byte) {
+	if remain := w.N - len(*dst); remain > 0 {
+		add := minInt(len(p), remain)
+		*dst = append(*dst, p[:add]...)
+		p = p[add:]
+	}
+	return p
+}
+
+func (w *prefixSuffixSaver) Bytes() []byte {
+	if w.suffix == nil {
+		return w.prefix
+	}
+	if w.skipped == 0 {
+		return append(w.prefix, w.suffix...)
+	}
+	var buf bytes.Buffer
+	buf.Grow(len(w.prefix) + len(w.suffix) + 50)
+	buf.Write(w.prefix)
+	buf.WriteString("\n... omitting ")
+	buf.WriteString(strconv.FormatInt(w.skipped, 10))
+	buf.WriteString(" bytes ...\n")
+	buf.Write(w.suffix[w.suffixOff:])
+	buf.Write(w.suffix[:w.suffixOff])
+	return buf.Bytes()
+}
+
+func minInt(a, b int) int {
+	if a < b {
+		return a
+	}
+	return b
+}
diff --git a/src/os/exec/exec_test.go b/src/os/exec/exec_test.go
index 28be21ce63..52b4724ab0 100644
--- a/src/os/exec/exec_test.go
+++ b/src/os/exec/exec_test.go
@@ -760,6 +760,9 @@ func TestHelperProcess(*testing.T) {
 		}
 		fmt.Print(p)
 		os.Exit(0)
+	case "stderrfail":
+		fmt.Fprintf(os.Stderr, "some stderr text\n")
+		os.Exit(1)
 	default:
 		fmt.Fprintf(os.Stderr, "Unknown command %q\n", cmd)
 		os.Exit(2)
@@ -816,3 +819,19 @@ func TestClosePipeOnCopyError(t *testing.T) {
 		t.Fatalf("yes got stuck writing to bad writer")
 	}
 }
+
+func TestOutputStderrCapture(t *testing.T) {
+	testenv.MustHaveExec(t)
+
+	cmd := helperCommand(t, "stderrfail")
+	_, err := cmd.Output()
+	ee, ok := err.(*exec.ExitError)
+	if !ok {
+		t.Fatalf("Output error type = %T; want ExitError", err)
+	}
+	got := string(ee.Stderr)
+	want := "some stderr text\n"
+	if got != want {
+		t.Errorf("ExitError.Stderr = %q; want %q", got, want)
+	}
+}
diff --git a/src/os/exec/internal_test.go b/src/os/exec/internal_test.go
new file mode 100644
index 0000000000..68d517ffb9
--- /dev/null
+++ b/src/os/exec/internal_test.go
@@ -0,0 +1,61 @@
+// Copyright 2015 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package exec
+
+import (
+	"io"
+	"testing"
+)
+
+func TestPrefixSuffixSaver(t *testing.T) {
+	tests := []struct {
+		N      int
+		writes []string
+		want   string
+	}{
+		{
+			N:      2,
+			writes: nil,
+			want:   "",
+		},
+		{
+			N:      2,
+			writes: []string{"a"},
+			want:   "a",
+		},
+		{
+			N:      2,
+			writes: []string{"abc", "d"},
+			want:   "abcd",
+		},
+		{
+			N:      2,
+			writes: []string{"abc", "d", "e"},
+			want:   "ab\n... omitting 1 bytes ...\nde",
+		},
+		{
+			N:      2,
+			writes: []string{"ab______________________yz"},
+			want:   "ab\n... omitting 22 bytes ...\nyz",
+		},
+		{
+			N:      2,
+			writes: []string{"ab_______________________y", "z"},
+			want:   "ab\n... omitting 23 bytes ...\nyz",
+		},
+	}
+	for i, tt := range tests {
+		w := &prefixSuffixSaver{N: tt.N}
+		for _, s := range tt.writes {
+			n, err := io.WriteString(w, s)
+			if err != nil || n != len(s) {
+				t.Errorf("%d. WriteString(%q) = %v, %v; want %v, %v", i, s, n, err, len(s), nil)
+			}
+		}
+		if got := string(w.Bytes()); got != tt.want {
+			t.Errorf("%d. Bytes = %q; want %q", i, got, tt.want)
+		}
+	}
+}