peer_protocol/pex_test.go

   1 package peer_protocol
   2
   3 import (
   4         "bufio"
   5         "bytes"
   6         "net"
   7         "testing"
   8
   9         "github.com/stretchr/testify/require"
  10
  11         "github.com/anacrolix/dht/v2/krpc"
  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 TestPexAppendAdded(t *testing.T) {
  33         t.Run("ipv4", func(t *testing.T) {
  34                 addr := krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}
  35                 f := PexPrefersEncryption | PexOutgoingConn
  36                 xm := PexMsg{}
  37                 xm.AppendAdded(addr, f)
  38                 require.EqualValues(t, len(xm.Added), 1)
  39                 require.EqualValues(t, len(xm.AddedFlags), 1)
  40                 require.EqualValues(t, len(xm.Added6), 0)
  41                 require.EqualValues(t, len(xm.Added6Flags), 0)
  42                 require.True(t, xm.Added[0].IP.Equal(addr.IP), "IPs should match")
  43                 require.EqualValues(t, xm.Added[0].Port, addr.Port)
  44                 require.EqualValues(t, xm.AddedFlags[0], f)
  45         })
  46         t.Run("ipv6", func(t *testing.T) {
  47                 addr := krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}
  48                 f := PexPrefersEncryption | PexOutgoingConn
  49                 xm := PexMsg{}
  50                 xm.AppendAdded(addr, f)
  51                 require.EqualValues(t, len(xm.Added), 0)
  52                 require.EqualValues(t, len(xm.AddedFlags), 0)
  53                 require.EqualValues(t, len(xm.Added6), 1)
  54                 require.EqualValues(t, len(xm.Added6Flags), 1)
  55                 require.True(t, xm.Added6[0].IP.Equal(addr.IP), "IPs should match")
  56                 require.EqualValues(t, xm.Added6[0].Port, addr.Port)
  57                 require.EqualValues(t, xm.Added6Flags[0], f)
  58         })
  59         t.Run("unspecified", func(t *testing.T) {
  60                 addr := krpc.NodeAddr{}
  61                 xm := PexMsg{}
  62                 xm.AppendAdded(addr, 0)
  63                 require.EqualValues(t, len(xm.Added), 0)
  64                 require.EqualValues(t, len(xm.AddedFlags), 0)
  65                 require.EqualValues(t, len(xm.Added6), 0)
  66                 require.EqualValues(t, len(xm.Added6Flags), 0)
  67         })
  68 }
  69
  70 func TestPexAppendDropped(t *testing.T) {
  71         t.Run("ipv4", func(t *testing.T) {
  72                 addr := krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}
  73                 xm := PexMsg{}
  74                 xm.AppendDropped(addr)
  75                 require.EqualValues(t, len(xm.Dropped), 1)
  76                 require.EqualValues(t, len(xm.Dropped6), 0)
  77                 require.True(t, xm.Dropped[0].IP.Equal(addr.IP), "IPs should match")
  78                 require.EqualValues(t, xm.Dropped[0].Port, addr.Port)
  79         })
  80         t.Run("ipv6", func(t *testing.T) {
  81                 addr := krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}
  82                 xm := PexMsg{}
  83                 xm.AppendDropped(addr)
  84                 require.EqualValues(t, len(xm.Dropped), 0)
  85                 require.EqualValues(t, len(xm.Dropped6), 1)
  86                 require.True(t, xm.Dropped6[0].IP.Equal(addr.IP), "IPs should match")
  87                 require.EqualValues(t, xm.Dropped6[0].Port, addr.Port)
  88         })
  89         t.Run("unspecified", func(t *testing.T) {
  90                 addr := krpc.NodeAddr{}
  91                 xm := PexMsg{}
  92                 xm.AppendDropped(addr)
  93                 require.EqualValues(t, len(xm.Dropped), 0)
  94                 require.EqualValues(t, len(xm.Dropped6), 0)
  95         })
  96 }
  97
  98 func TestMarshalPexMessage(t *testing.T) {
  99         addr := krpc.NodeAddr{IP: net.IP{127, 0, 0, 1}, Port: 0x55aa}
 100         f := PexPrefersEncryption | PexOutgoingConn
 101         pm := PexMsg{}
 102         pm.AppendAdded(addr, f)
 103
 104         b, err := bencode.Marshal(pm)
 105         require.NoError(t, err)
 106
 107         pexExtendedId := ExtensionNumber(7)
 108         msg := pm.Message(pexExtendedId)
 109         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")
 110         b, err = msg.MarshalBinary()
 111         require.NoError(t, err)
 112         require.EqualValues(t, b, expected)
 113
 114         msg = Message{}
 115         dec := Decoder{
 116                 R:         bufio.NewReader(bytes.NewBuffer(b)),
 117                 MaxLength: 128,
 118         }
 119         pmOut := PexMsg{}
 120         err = dec.Decode(&msg)
 121         require.NoError(t, err)
 122         require.EqualValues(t, Extended, msg.Type)
 123         require.EqualValues(t, pexExtendedId, msg.ExtendedID)
 124         err = bencode.Unmarshal(msg.ExtendedPayload, &pmOut)
 125         require.NoError(t, err)
 126         require.EqualValues(t, len(pm.Added), len(pmOut.Added))
 127         require.EqualValues(t, pm.Added[0].IP, pmOut.Added[0].IP)
 128         require.EqualValues(t, pm.Added[0].Port, pmOut.Added[0].Port)
 129 }