peer_protocol/pex_test.go

   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 }