Merge pull request #891 from buildpacks/fix/890-index-out-of-range

Fix: Index out-of-bounds, CR processing

logging/prefix_writer.go

@@ -11,66 +11,88 @@ import (
 
 // PrefixWriter is a buffering writer that prefixes each new line. Close should be called to properly flush the buffer.
 type PrefixWriter struct {
-	out    io.Writer
-	buf    *bytes.Buffer
-	prefix string
+	out           io.Writer
+	buf           *bytes.Buffer
+	prefix        string
+	readerFactory func(data []byte) io.Reader
+}
+
+type PrefixWriterOption func(c *PrefixWriter)
+
+func WithReaderFactory(factory func(data []byte) io.Reader) PrefixWriterOption {
+	return func(writer *PrefixWriter) {
+		writer.readerFactory = factory
+	}
 }
 
 // NewPrefixWriter writes by w will be prefixed
-func NewPrefixWriter(w io.Writer, prefix string) *PrefixWriter {
-	return &PrefixWriter{
+func NewPrefixWriter(w io.Writer, prefix string, opts ...PrefixWriterOption) *PrefixWriter {
+	writer := &PrefixWriter{
 		out:    w,
 		prefix: fmt.Sprintf("[%s] ", style.Prefix(prefix)),
 		buf:    &bytes.Buffer{},
+		readerFactory: func(data []byte) io.Reader {
+			return bytes.NewReader(data)
+		},
+	}
+
+	for _, opt := range opts {
+		opt(writer)
 	}
+
+	return writer
 }
 
 // Write writes bytes to the embedded log function
 func (w *PrefixWriter) Write(data []byte) (int, error) {
-	scanner := bufio.NewScanner(bytes.NewReader(data))
+	scanner := bufio.NewScanner(w.readerFactory(data))
 	scanner.Split(ScanLinesKeepNewLine)
 	for scanner.Scan() {
 		newBits := scanner.Bytes()
-		if newBits[len(newBits)-1] != '\n' { // just append if we don't have a new line
+		if len(newBits) > 0 && newBits[len(newBits)-1] != '\n' { // just append if we don't have a new line
 			_, err := w.buf.Write(newBits)
 			if err != nil {
 				return 0, err
 			}
 		} else { // write our complete message
-			var allBits []byte
-			if w.buf.Len() > 0 {
-				allBits = append(w.buf.Bytes(), newBits...)
-				w.buf.Reset()
-			} else {
-				allBits = newBits
+			_, err := w.buf.Write(bytes.TrimRight(newBits, "\n"))
+			if err != nil {
+				return 0, err
 			}
 
-			err := w.writeWithPrefix(allBits)
+			err = w.flush()
 			if err != nil {
 				return 0, err
 			}
 		}
 	}
 
+	if err := scanner.Err(); err != nil {
+		return 0, err
+	}
+
 	return len(data), nil
 }
 
 // Close writes any pending data in the buffer
 func (w *PrefixWriter) Close() error {
 	if w.buf.Len() > 0 {
-		err := w.writeWithPrefix(w.buf.Bytes())
-		if err != nil {
-			return err
-		}
+		return w.flush()
 	}
 
-	w.buf.Reset()
-
 	return nil
 }
 
-func (w *PrefixWriter) writeWithPrefix(bits []byte) error {
-	_, err := fmt.Fprint(w.out, w.prefix+string(bits))
+func (w *PrefixWriter) flush() error {
+	bits := w.buf.Bytes()
+	w.buf.Reset()
+
+	// process any CR in message
+	if i := bytes.LastIndexByte(bits, '\r'); i >= 0 {
+		bits = bits[i+1:]
+	}
+
+	_, err := fmt.Fprint(w.out, w.prefix+string(bits)+"\n")
 	return err
 }
 
@@ -79,13 +101,17 @@ func ScanLinesKeepNewLine(data []byte, atEOF bool) (advance int, token []byte, e
 	if atEOF && len(data) == 0 {
 		return 0, nil, nil
 	}
+
+	// first we'll split by LF (\n)
+	// then remove any preceding CR (\r) [due to CR+LF]
 	if i := bytes.IndexByte(data, '\n'); i >= 0 {
 		// We have a full newline-terminated line.
 		return i + 1, append(dropCR(data[0:i]), '\n'), nil
 	}
+
 	// If we're at EOF, we have a final, non-terminated line. Return it.
 	if atEOF {
-		return len(data), dropCR(data), nil
+		return len(data), data, nil
 	}
 	// Request more data.
 	return 0, nil, nil

logging/prefix_writer_test.go

@@ -2,7 +2,8 @@ package logging_test
 
 import (
 	"bytes"
-	"fmt"
+	"errors"
+	"io"
 	"testing"
 
 	"github.com/heroku/color"
@@ -20,36 +21,129 @@ func TestPrefixWriter(t *testing.T) {
 }
 
 func testPrefixWriter(t *testing.T, when spec.G, it spec.S) {
+	var (
+		assert = h.NewAssertionManager(t)
+	)
+
 	when("#Write", func() {
 		it("prepends prefix to string", func() {
 			var w bytes.Buffer
-			prefix := "test prefix"
-			writer := logging.NewPrefixWriter(&w, prefix)
-			_, _ = writer.Write([]byte("test"))
-			_ = writer.Close()
 
-			h.AssertEq(t, w.String(), fmt.Sprintf("[%s] %s", prefix, "test"))
+			writer := logging.NewPrefixWriter(&w, "prefix")
+			_, err := writer.Write([]byte("test"))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
+
+			h.AssertEq(t, w.String(), "[prefix] test\n")
 		})
 
 		it("prepends prefix to multi-line string", func() {
 			var w bytes.Buffer
 
 			writer := logging.NewPrefixWriter(&w, "prefix")
-			_, _ = writer.Write([]byte("line 1\nline 2\nline 3"))
-			_ = writer.Close()
-			h.AssertEq(t, w.String(), "[prefix] line 1\n[prefix] line 2\n[prefix] line 3")
+			_, err := writer.Write([]byte("line 1\nline 2\nline 3"))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
+
+			h.AssertEq(t, w.String(), "[prefix] line 1\n[prefix] line 2\n[prefix] line 3\n")
 		})
 
 		it("buffers mid-line calls", func() {
 			var buf bytes.Buffer
 
 			writer := logging.NewPrefixWriter(&buf, "prefix")
-			_, _ = writer.Write([]byte("word 1, "))
-			_, _ = writer.Write([]byte("word 2, "))
-			_, _ = writer.Write([]byte("word 3."))
-			_ = writer.Close()
+			_, err := writer.Write([]byte("word 1, "))
+			assert.Nil(err)
+			_, err = writer.Write([]byte("word 2, "))
+			assert.Nil(err)
+			_, err = writer.Write([]byte("word 3."))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
 
-			h.AssertEq(t, buf.String(), "[prefix] word 1, word 2, word 3.")
+			h.AssertEq(t, buf.String(), "[prefix] word 1, word 2, word 3.\n")
+		})
+
+		it("handles empty lines", func() {
+			var buf bytes.Buffer
+
+			writer := logging.NewPrefixWriter(&buf, "prefix")
+			_, err := writer.Write([]byte("\n"))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
+
+			h.AssertEq(t, buf.String(), "[prefix] \n")
+		})
+
+		it("handles empty input", func() {
+			var buf bytes.Buffer
+
+			writer := logging.NewPrefixWriter(&buf, "prefix")
+			_, err := writer.Write([]byte(""))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
+
+			assert.Equal(buf.String(), "")
+		})
+
+		it("propagates reader errors", func() {
+			var buf bytes.Buffer
+
+			factory := &boobyTrapReaderFactory{failAtCallNumber: 2}
+			writer := logging.NewPrefixWriter(&buf, "prefix", logging.WithReaderFactory(factory.NewReader))
+			_, err := writer.Write([]byte("word 1,"))
+			assert.Nil(err)
+			_, err = writer.Write([]byte("word 2."))
+			assert.ErrorContains(err, "some error")
+		})
+
+		it("handles requests to clear line", func() {
+			var buf bytes.Buffer
+
+			writer := logging.NewPrefixWriter(&buf, "prefix")
+			_, err := writer.Write([]byte("progress 1\rprogress 2\rprogress 3\rcomplete!"))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
+
+			h.AssertEq(t, buf.String(), "[prefix] complete!\n")
+		})
+
+		it("handles requests clear line (amidst content)", func() {
+			var buf bytes.Buffer
+
+			writer := logging.NewPrefixWriter(&buf, "prefix")
+			_, err := writer.Write([]byte("downloading\rcompleted!      \r\nall done!\nnevermind\r"))
+			assert.Nil(err)
+			err = writer.Close()
+			assert.Nil(err)
+
+			h.AssertEq(t, buf.String(), "[prefix] completed!      \n[prefix] all done!\n[prefix] \n")
 		})
 	})
 }
+
+type boobyTrapReaderFactory struct {
+	numberOfCalls    int
+	failAtCallNumber int
+}
+
+func (b *boobyTrapReaderFactory) NewReader(data []byte) io.Reader {
+	b.numberOfCalls++
+	if b.numberOfCalls >= b.failAtCallNumber {
+		return &faultyReader{}
+	}
+
+	return bytes.NewReader(data)
+}
+
+type faultyReader struct {
+}
+
+func (f faultyReader) Read(b []byte) (n int, err error) {
+	return 0, errors.New("some error")
+}