bytes: avoid overflow in (*Buffer).Grow and ReadFrom

fixes #21481

Change-Id: I26717876a1c0ee25a86c81159c6b3c59563dfec6
Reviewed-on: https://go-review.googlesource.com/56230
Reviewed-by: Ian Lance Taylor <iant@golang.org>

2 files changed, 25 insertions, 4 deletions

diff --git a/src/bytes/buffer.go b/src/bytes/buffer.go
index 20e42bbbbc..cf4f31d7d2 100644
--- a/src/bytes/buffer.go
+++ b/src/bytes/buffer.go
@@ -44,6 +44,8 @@ const (
 // ErrTooLarge is passed to panic if memory cannot be allocated to store data in a buffer.
 var ErrTooLarge = errors.New("bytes.Buffer: too large")
 
+const maxInt = int(^uint(0) >> 1)
+
 // Bytes returns a slice of length b.Len() holding the unread portion of the buffer.
 // The slice is valid for use only until the next buffer modification (that is,
 // only until the next call to a method like Read, Write, Reset, or Truncate).
@@ -97,7 +99,7 @@ func (b *Buffer) Reset() {
 // internal buffer only needs to be resliced.
 // It returns the index where bytes should be written and whether it succeeded.
 func (b *Buffer) tryGrowByReslice(n int) (int, bool) {
-	if l := len(b.buf); l+n <= cap(b.buf) {
+	if l := len(b.buf); n <= cap(b.buf)-l {
 		b.buf = b.buf[:l+n]
 		return l, true
 	}
@@ -122,15 +124,18 @@ func (b *Buffer) grow(n int) int {
 		b.buf = b.bootstrap[:n]
 		return 0
 	}
-	if m+n <= cap(b.buf)/2 {
+	c := cap(b.buf)
+	if n <= c/2-m {
 		// We can slide things down instead of allocating a new
 		// slice. We only need m+n <= cap(b.buf) to slide, but
 		// we instead let capacity get twice as large so we
 		// don't spend all our time copying.
 		copy(b.buf[:], b.buf[b.off:])
+	} else if c > maxInt-c-n {
+		panic(ErrTooLarge)
 	} else {
 		// Not enough space anywhere, we need to allocate.
-		buf := makeSlice(2*cap(b.buf) + n)
+		buf := makeSlice(2*c + n)
 		copy(buf, b.buf[b.off:])
 		b.buf = buf
 	}
@@ -200,7 +205,11 @@ func (b *Buffer) ReadFrom(r io.Reader) (n int64, err error) {
 			if b.off+free < MinRead {
 				// not enough space using beginning of buffer;
 				// double buffer capacity
-				newBuf = makeSlice(2*cap(b.buf) + MinRead)
+				c := cap(b.buf)
+				if c > maxInt-c-MinRead {
+					panic(ErrTooLarge)
+				}
+				newBuf = makeSlice(2*c + MinRead)
 			}
 			copy(newBuf, b.buf[b.off:])
 			b.buf = newBuf[:len(b.buf)-b.off]
diff --git a/src/bytes/buffer_test.go b/src/bytes/buffer_test.go
index ce2f01a0ad..dcfbfced92 100644
--- a/src/bytes/buffer_test.go
+++ b/src/bytes/buffer_test.go
@@ -473,6 +473,18 @@ func TestGrow(t *testing.T) {
 	}
 }
 
+func TestGrowOverflow(t *testing.T) {
+	defer func() {
+		if err := recover(); err != ErrTooLarge {
+			t.Errorf("after too-large Grow, recover() = %v; want %v", err, ErrTooLarge)
+		}
+	}()
+
+	buf := NewBuffer(make([]byte, 1))
+	const maxInt = int(^uint(0) >> 1)
+	buf.Grow(maxInt)
+}
+
 // Was a bug: used to give EOF reading empty slice at EOF.
 func TestReadEmptyAtEOF(t *testing.T) {
 	b := new(Buffer)