1
0
mirror of https://github.com/golang/go synced 2024-11-26 06:57:56 -07:00

bufio: check buffer availability before reading in ReadFrom

Fixes #5947.

R=golang-dev, bradfitz
CC=golang-dev
https://golang.org/cl/11801043
This commit is contained in:
Andrew Gerrand 2013-07-25 11:29:13 +10:00
parent f7dfeea90f
commit 93c6d0ef8f
2 changed files with 33 additions and 6 deletions

View File

@ -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
}

View File

@ -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