bufio: check buffer availability before reading in ReadFrom

Fixes #5947.

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

diff --git a/src/pkg/bufio/bufio.go b/src/pkg/bufio/bufio.go
index 4df5e32d9ae0a79f6eeec79481785475a9f45de0..993c22d61cd827e909617858661b9232e69e1cfc 100644 (file)
--- a/src/pkg/bufio/bufio.go
+++ b/src/pkg/bufio/bufio.go
@@ -678,23 +678,28 @@ 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 {
-               err = nil
+               // If we filled the buffer exactly, flush pre-emptively.
+               if b.Available() == 0 {
+                       err = b.flush()
+               } else {
+                       err = nil
+               }
        }
        return n, err
 }

diff --git a/src/pkg/bufio/bufio_test.go b/src/pkg/bufio/bufio_test.go
index 93f1b3fe086fabbdff4573cccc65851eb9e1a3fd..68a138e5c1e296ecff67d2c0cc5aed93e8b101b9 100644 (file)
--- a/src/pkg/bufio/bufio_test.go
+++ b/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