Add a test that short pieces are decoded correctly

diff --git a/peer_protocol/decoder_test.go b/peer_protocol/decoder_test.go
index 348f8247e67848eb92d6c7b93ebff6c3a022c974..5dfd09c5ec80f5e2c407d42cdb1ade402f310412 100644 (file)
--- a/peer_protocol/decoder_test.go
+++ b/peer_protocol/decoder_test.go
@@ -7,6 +7,7 @@ import (
        "testing"
 
        "github.com/bradfitz/iter"
+       "github.com/stretchr/testify/assert"
        "github.com/stretchr/testify/require"
 )
 
@@ -52,3 +53,23 @@ func BenchmarkDecodePieces(t *testing.B) {
                d.Pool.Put(&msg.Piece)
        }
 }
+
+func TestDecodeShortPieceEOF(t *testing.T) {
+       r, w := io.Pipe()
+       go func() {
+               w.Write(Message{Type: Piece, Piece: make([]byte, 1)}.MustMarshalBinary())
+               w.Close()
+       }()
+       d := Decoder{
+               R:         bufio.NewReader(r),
+               MaxLength: 1 << 15,
+               Pool: &sync.Pool{New: func() interface{} {
+                       b := make([]byte, 2)
+                       return &b
+               }},
+       }
+       var m Message
+       require.NoError(t, d.Decode(&m))
+       assert.Len(t, m.Piece, 1)
+       assert.Equal(t, io.EOF, d.Decode(&m))
+}