github.com/anacrolix/torrent@v1.61.0/peer_protocol/decoder_test.go (about)

     1  package peer_protocol
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"io"
     7  	"sync"
     8  	"testing"
     9  
    10  	qt "github.com/go-quicktest/qt"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  func BenchmarkDecodePieces(t *testing.B) {
    16  	const pieceLen = 1 << 14
    17  	inputMsg := Message{
    18  		Type:  Piece,
    19  		Index: 0,
    20  		Begin: 1,
    21  		Piece: make([]byte, pieceLen),
    22  	}
    23  	b := inputMsg.MustMarshalBinary()
    24  	t.SetBytes(int64(len(b)))
    25  	var r bytes.Reader
    26  	// Try to somewhat emulate what torrent.Client would do. But the goal is to get decoding as fast
    27  	// as possible and let consumers apply their own adjustments.
    28  	d := Decoder{
    29  		R:         bufio.NewReaderSize(&r, 1<<10),
    30  		MaxLength: 1 << 18,
    31  		Pool: &sync.Pool{
    32  			New: func() interface{} {
    33  				b := make([]byte, pieceLen)
    34  				return &b
    35  			},
    36  		},
    37  	}
    38  	t.ReportAllocs()
    39  	t.ResetTimer()
    40  	for i := 0; i < t.N; i += 1 {
    41  		r.Reset(b)
    42  		var msg Message
    43  		err := d.Decode(&msg)
    44  		if err != nil {
    45  			t.Fatal(err)
    46  		}
    47  		// This is very expensive, and should be discovered in tests rather than a benchmark.
    48  		if false {
    49  			qt.Assert(t, qt.DeepEquals(msg, inputMsg))
    50  		}
    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.ErrorIs(t, d.Decode(&m), io.EOF)
    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  }