Implement io.ReaderFrom

bytebuffer.go

@@ -16,7 +16,46 @@ type ByteBuffer struct {
 	B []byte
 }
 
-// WriteTo implements io.WriterTo
+// Len returns the size of the byte buffer.
+func (b *ByteBuffer) Len() int {
+	return len(b.B)
+}
+
+// ReadFrom implements io.ReaderFrom.
+//
+// The function appends all the data read from r to b.
+func (b *ByteBuffer) ReadFrom(r io.Reader) (int64, error) {
+	p := b.B
+	nStart := int64(len(p))
+	nMax := int64(cap(p))
+	n := nStart
+	if nMax == 0 {
+		nMax = 64
+		p = make([]byte, nMax)
+	} else {
+		p = p[:nMax]
+	}
+	for {
+		if n == nMax {
+			nMax *= 2
+			bNew := make([]byte, nMax)
+			copy(bNew, p)
+			p = bNew
+		}
+		nn, err := r.Read(p[n:])
+		n += int64(nn)
+		if err != nil {
+			b.B = p[:n]
+			n -= nStart
+			if err == io.EOF {
+				return n, nil
+			}
+			return n, err
+		}
+	}
+}
+
+// WriteTo implements io.WriterTo.
 func (b *ByteBuffer) WriteTo(w io.Writer) (int64, error) {
 	n, err := w.Write(b.B)
 	return int64(n), err

bytebuffer_test.go

@@ -8,6 +8,40 @@ import (
 	"time"
 )
 
+func TestByteBufferReadFrom(t *testing.T) {
+	prefix := "foobar"
+	expectedS := "asadfsdafsadfasdfisdsdfa"
+	prefixLen := int64(len(prefix))
+	expectedN := int64(len(expectedS))
+
+	var bb ByteBuffer
+	bb.WriteString(prefix)
+
+	rf := (io.ReaderFrom)(&bb)
+	for i := 0; i < 20; i++ {
+		r := bytes.NewBufferString(expectedS)
+		n, err := rf.ReadFrom(r)
+		if n != expectedN {
+			t.Fatalf("unexpected n=%d. Expecting %d. iteration %d", n, expectedN, i)
+		}
+		if err != nil {
+			t.Fatalf("unexpected error: %s", err)
+		}
+		bbLen := int64(bb.Len())
+		expectedLen := prefixLen + int64(i+1)*expectedN
+		if bbLen != expectedLen {
+			t.Fatalf("unexpected byteBuffer length: %d. Expecting %d", bbLen, expectedLen)
+		}
+		for j := 0; j < i; j++ {
+			start := prefixLen + int64(j)*expectedN
+			b := bb.B[start : start+expectedN]
+			if string(b) != expectedS {
+				t.Fatalf("unexpected byteBuffer contents: %q. Expecting %q", b, expectedS)
+			}
+		}
+	}
+}
+
 func TestByteBufferWriteTo(t *testing.T) {
 	expectedS := "foobarbaz"
 	var bb ByteBuffer