peer_protocol/decoder_test.go

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