From 91730696cfe5dd9f088e21a4f3b3a1e915348758 Mon Sep 17 00:00:00 2001 From: Matt Joiner Date: Sat, 14 Jul 2018 11:50:43 +1000 Subject: [PATCH] Rewrite piece data decoding and relax test --- peer_protocol/decoder.go | 24 +++++++++++------------- peer_protocol/decoder_test.go | 2 +- 2 files changed, 12 insertions(+), 14 deletions(-) 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)) } -- 2.44.0