fix panic when serializing big bitfields

internal/peerconn/peerwriter/piece_test.go

@@ -0,0 +1,22 @@
+package peerwriter
+
+import (
+	"bytes"
+	"testing"
+)
+
+func BenchmarkRead(b *testing.B) {
+	buf := make([]byte, 10)
+	buf2 := make([]byte, 25)
+	r := bytes.NewReader(buf)
+	p := Piece{
+		Piece:  r,
+		Begin:  2,
+		Length: 5,
+	}
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		p.Read(buf2)
+	}
+}

internal/peerprotocol/extension.go

@@ -35,22 +35,26 @@ type ExtensionMessage struct {
 func (m ExtensionMessage) ID() MessageID { return Extension }
 
 func (m ExtensionMessage) Read([]byte) (int, error) {
-	panic("read must not be called")
+	panic("Read must not be called, use WriteTo")
 }
 
-func (m ExtensionMessage) WriteTo(w io.Writer) (int64, error) {
-	_, err := w.Write([]byte{m.ExtendedMessageID})
+func (m ExtensionMessage) WriteTo(w io.Writer) (n int64, err error) {
+	nn, err := w.Write([]byte{m.ExtendedMessageID})
+	n += int64(nn)
 	if err != nil {
-		return 0, err
+		return
 	}
-	err = bencode.NewEncoder(w).Encode(m.Payload)
+	wc := newWriterCounter(w)
+	err = bencode.NewEncoder(wc).Encode(m.Payload)
+	n += wc.Count()
 	if err != nil {
-		return 0, err
+		return
 	}
 	if mm, ok := m.Payload.(ExtensionMetadataMessage); ok {
-		_, err = w.Write(mm.Data)
+		nn, err = w.Write(mm.Data)
+		n += int64(nn)
 	}
-	return 0, err
+	return
 }
 
 func (m *ExtensionMessage) UnmarshalBinary(data []byte) error {

internal/peerprotocol/messages.go

@@ -6,7 +6,7 @@ import (
 )
 
 type Message interface {
-	Read([]byte) (int, error)
+	io.Reader
 	ID() MessageID
 }
 
@@ -48,13 +48,18 @@ func (m PieceMessage) Read(b []byte) (int, error) {
 
 type BitfieldMessage struct {
 	Data []byte
+	pos  int
 }
 
 func (m BitfieldMessage) ID() MessageID { return Bitfield }
 
-func (m BitfieldMessage) Read(b []byte) (int, error) {
-	copy(b[0:len(m.Data)], m.Data)
-	return len(m.Data), io.EOF
+func (m BitfieldMessage) Read(b []byte) (n int, err error) {
+	n = copy(b, m.Data[m.pos:])
+	m.pos += n
+	if m.pos == len(m.Data) {
+		err = io.EOF
+	}
+	return
 }
 
 type emptyMessage struct{}

internal/peerprotocol/writecounter.go

@@ -0,0 +1,25 @@
+package peerprotocol
+
+import "io"
+
+type writerCounter struct {
+	io.Writer
+	count  int64
+	writer io.Writer
+}
+
+func newWriterCounter(w io.Writer) *writerCounter {
+	return &writerCounter{
+		writer: w,
+	}
+}
+
+func (w *writerCounter) Write(buf []byte) (int, error) {
+	n, err := w.writer.Write(buf)
+	w.count += int64(n)
+	return n, err
+}
+
+func (w *writerCounter) Count() int64 {
+	return w.count
+}