bufio: check buffer availability before reading in ReadFrom

Fixes #5947.

R=golang-dev, bradfitz
CC=golang-dev
https://golang.org/cl/11801043

src/pkg/bufio/bufio.go

@@ -678,24 +678,29 @@ func (b *Writer) ReadFrom(r io.Reader) (n int64, err error) {
 	}
 	var m int
 	for {
+		if b.Available() == 0 {
+			if err1 := b.flush(); err1 != nil {
+				return n, err1
+			}
+		}
 		m, err = r.Read(b.buf[b.n:])
 		if m == 0 {
 			break
 		}
 		b.n += m
 		n += int64(m)
-		if b.Available() == 0 {
-			if err1 := b.flush(); err1 != nil {
-				return n, err1
-			}
-		}
 		if err != nil {
 			break
 		}
 	}
 	if err == io.EOF {
+		// If we filled the buffer exactly, flush pre-emptively.
+		if b.Available() == 0 {
+			err = b.flush()
+		} else {
 			err = nil
 		}
+	}
 	return n, err
 }
 

src/pkg/bufio/bufio_test.go

@@ -847,6 +847,10 @@ func TestWriterReadFrom(t *testing.T) {
 				t.Errorf("ws[%d],rs[%d]: w.ReadFrom(r) = %d, %v, want %d, nil", wi, ri, n, err, len(input))
 				continue
 			}
+			if err := w.Flush(); err != nil {
+				t.Errorf("Flush returned %v", err)
+				continue
+			}
 			if got, want := b.String(), string(input); got != want {
 				t.Errorf("ws[%d], rs[%d]:\ngot  %q\nwant %q\n", wi, ri, got, want)
 			}
@@ -1003,6 +1007,24 @@ func TestReaderClearError(t *testing.T) {
 	}
 }
 
+// Test for golang.org/issue/5947
+func TestWriterReadFromWhileFull(t *testing.T) {
+	buf := new(bytes.Buffer)
+	w := NewWriterSize(buf, 10)
+
+	// Fill buffer exactly.
+	n, err := w.Write([]byte("0123456789"))
+	if n != 10 || err != nil {
+		t.Fatalf("Write returned (%v, %v), want (10, nil)", n, err)
+	}
+
+	// Use ReadFrom to read in some data.
+	n2, err := w.ReadFrom(strings.NewReader("abcdef"))
+	if n2 != 6 || err != nil {
+		t.Fatalf("ReadFrom returned (%v, %v), want (6, nil)", n, err)
+	}
+}
+
 // An onlyReader only implements io.Reader, no matter what other methods the underlying implementation may have.
 type onlyReader struct {
 	r io.Reader