From: Matt Joiner <anacrolix@gmail.com>
Date: Sat, 14 Jul 2018 01:50:43 +0000 (+1000)
Subject: Rewrite piece data decoding and relax test
X-Git-Tag: v1.0.0~79
X-Git-Url: http://www.git.stargrave.org/?a=commitdiff_plain;h=91730696cfe5dd9f088e21a4f3b3a1e915348758;p=btrtrc.git

Rewrite piece data decoding and relax test
---

diff --git a/peer_protocol/decoder.go b/peer_protocol/decoder.go
index 47a8b65e..b7ccab14 100644
--- a/peer_protocol/decoder.go
+++ b/peer_protocol/decoder.go
@@ -3,11 +3,12 @@ package peer_protocol
 import (
 	"bufio"
 	"encoding/binary"
-	"errors"
 	"fmt"
 	"io"
 	"io/ioutil"
 	"sync"
+
+	"github.com/pkg/errors"
 )
 
 type Decoder struct {
@@ -69,24 +70,21 @@ func (d *Decoder) Decode(msg *Message) (err error) {
 		msg.Bitfield = unmarshalBitfield(b)
 	case Piece:
 		for _, pi := range []*Integer{&msg.Index, &msg.Begin} {
-			err = pi.Read(r)
+			err := pi.Read(r)
 			if err != nil {
-				break
+				return err
 			}
 		}
-		if err != nil {
-			break
+		dataLen := r.N
+		msg.Piece = (*d.Pool.Get().(*[]byte))
+		if int64(cap(msg.Piece)) < dataLen {
+			return errors.New("piece data longer than expected")
 		}
-		//msg.Piece, err = ioutil.ReadAll(r)
-		b := *d.Pool.Get().(*[]byte)
-		n, err := io.ReadFull(r, b)
+		msg.Piece = msg.Piece[:dataLen]
+		_, err := io.ReadFull(r, msg.Piece)
 		if err != nil {
-			if err != io.ErrUnexpectedEOF || n != int(length-9) {
-				return err
-			}
-			b = b[0:n]
+			return errors.Wrap(err, "reading piece data")
 		}
-		msg.Piece = b
 	case Extended:
 		b, err := readByte(r)
 		if err != nil {
diff --git a/peer_protocol/decoder_test.go b/peer_protocol/decoder_test.go
index bbd8194d..33909cdc 100644
--- a/peer_protocol/decoder_test.go
+++ b/peer_protocol/decoder_test.go
@@ -89,5 +89,5 @@ func TestDecodeOverlongPiece(t *testing.T) {
 		}},
 	}
 	var m Message
-	require.EqualError(t, d.Decode(&m), "piece data longer than expected")
+	require.Error(t, d.Decode(&m))
 }