peer_protocol/decoder_test.go

   1 package peer_protocol
   2
   3 import (
   4         "bufio"
   5         "io"
   6         "sync"
   7         "testing"
   8
   9         "github.com/bradfitz/iter"
  10         "github.com/stretchr/testify/assert"
  11         "github.com/stretchr/testify/require"
  12 )
  13
  14 func BenchmarkDecodePieces(t *testing.B) {
  15         r, w := io.Pipe()
  16         const pieceLen = 1 << 14
  17         msg := Message{
  18                 Type:  Piece,
  19                 Index: 0,
  20                 Begin: 1,
  21                 Piece: make([]byte, pieceLen),
  22         }
  23         b, err := msg.MarshalBinary()
  24         require.NoError(t, err)
  25         t.SetBytes(int64(len(b)))
  26         defer r.Close()
  27         go func() {
  28                 defer w.Close()
  29                 for {
  30                         n, err := w.Write(b)
  31                         if err == io.ErrClosedPipe {
  32                                 return
  33                         }
  34                         require.NoError(t, err)
  35                         require.Equal(t, len(b), n)
  36                 }
  37         }()
  38         d := Decoder{
  39                 // Emulate what package torrent's client would do.
  40                 R:         bufio.NewReader(r),
  41                 MaxLength: 1 << 18,
  42                 Pool: &sync.Pool{
  43                         New: func() interface{} {
  44                                 b := make([]byte, pieceLen)
  45                                 return &b
  46                         },
  47                 },
  48         }
  49         for range iter.N(t.N) {
  50                 var msg Message
  51                 require.NoError(t, d.Decode(&msg))
  52                 // WWJD
  53                 d.Pool.Put(&msg.Piece)
  54         }
  55 }
  56
  57 func TestDecodeShortPieceEOF(t *testing.T) {
  58         r, w := io.Pipe()
  59         go func() {
  60                 w.Write(Message{Type: Piece, Piece: make([]byte, 1)}.MustMarshalBinary())
  61                 w.Close()
  62         }()
  63         d := Decoder{
  64                 R:         bufio.NewReader(r),
  65                 MaxLength: 1 << 15,
  66                 Pool: &sync.Pool{New: func() interface{} {
  67                         b := make([]byte, 2)
  68                         return &b
  69                 }},
  70         }
  71         var m Message
  72         require.NoError(t, d.Decode(&m))
  73         assert.Len(t, m.Piece, 1)
  74         assert.Equal(t, io.EOF, d.Decode(&m))
  75 }
  76
  77 func TestDecodeOverlongPiece(t *testing.T) {
  78         r, w := io.Pipe()
  79         go func() {
  80                 w.Write(Message{Type: Piece, Piece: make([]byte, 3)}.MustMarshalBinary())
  81                 w.Close()
  82         }()
  83         d := Decoder{
  84                 R:         bufio.NewReader(r),
  85                 MaxLength: 1 << 15,
  86                 Pool: &sync.Pool{New: func() interface{} {
  87                         b := make([]byte, 2)
  88                         return &b
  89                 }},
  90         }
  91         var m Message
  92         require.EqualError(t, d.Decode(&m), "piece data longer than expected")
  93 }