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

     1  package peer_protocol
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"net"
     7  	"testing"
     8  
     9  	"github.com/anacrolix/dht/v2/krpc"
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/anacrolix/torrent/bencode"
    13  )
    14  
    15  func TestUnmarshalPex(t *testing.T) {
    16  	var pem PexMsg
    17  	err := bencode.Unmarshal([]byte("d5:added12:\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0ce"), &pem)
    18  	require.NoError(t, err)
    19  	require.EqualValues(t, 2, len(pem.Added))
    20  	require.EqualValues(t, 1286, pem.Added[0].Port)
    21  	require.EqualValues(t, 0x100*0xb+0xc, pem.Added[1].Port)
    22  }
    23  
    24  func TestEmptyPexMsg(t *testing.T) {
    25  	pm := PexMsg{}
    26  	b, err := bencode.Marshal(pm)
    27  	t.Logf("%q", b)
    28  	require.NoError(t, err)
    29  	require.NoError(t, bencode.Unmarshal(b, &pm))
    30  }
    31  
    32  func TestMarshalPexMessage(t *testing.T) {
    33  	addr := krpc.NodeAddr{IP: net.IP{127, 0, 0, 1}, Port: 0x55aa}
    34  	f := PexPrefersEncryption | PexOutgoingConn
    35  	pm := new(PexMsg)
    36  	pm.Added = append(pm.Added, addr)
    37  	pm.AddedFlags = append(pm.AddedFlags, f)
    38  
    39  	_, err := bencode.Marshal(pm)
    40  	require.NoError(t, err)
    41  
    42  	pexExtendedId := ExtensionNumber(7)
    43  	msg := pm.Message(pexExtendedId)
    44  	expected := []byte("\x00\x00\x00\x4c\x14\x07d5:added6:\x7f\x00\x00\x01\x55\xaa7:added.f1:\x116:added60:8:added6.f0:7:dropped0:8:dropped60:e")
    45  	b, err := msg.MarshalBinary()
    46  	require.NoError(t, err)
    47  	require.EqualValues(t, b, expected)
    48  
    49  	msg = Message{}
    50  	dec := Decoder{
    51  		R:         bufio.NewReader(bytes.NewReader(b)),
    52  		MaxLength: 128,
    53  	}
    54  	pmOut := PexMsg{}
    55  	err = dec.Decode(&msg)
    56  	require.NoError(t, err)
    57  	require.EqualValues(t, Extended, msg.Type)
    58  	require.EqualValues(t, pexExtendedId, msg.ExtendedID)
    59  	err = bencode.Unmarshal(msg.ExtendedPayload, &pmOut)
    60  	require.NoError(t, err)
    61  	require.EqualValues(t, len(pm.Added), len(pmOut.Added))
    62  	require.EqualValues(t, pm.Added[0].IP, pmOut.Added[0].IP)
    63  	require.EqualValues(t, pm.Added[0].Port, pmOut.Added[0].Port)
    64  }