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 TestPexAdd(t *testing.T) {
  33         addrs4 := []krpc.NodeAddr{
  34                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}, // 0
  35                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748}, // 1
  36                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4747}, // 2
  37                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4748}, // 3
  38         }
  39         addrs6 := []krpc.NodeAddr{
  40                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}, // 0
  41                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4748}, // 1
  42                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4749}, // 2
  43                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4750}, // 3
  44         }
  45         f := PexPrefersEncryption | PexOutgoingConn
  46
  47         t.Run("ipv4", func(t *testing.T) {
  48                 addrs := addrs4
  49                 m := new(PexMsg)
  50                 m.Drop(addrs[0])
  51                 m.Add(addrs[1], f)
  52                 for _, addr := range addrs {
  53                         m.Add(addr, f)
  54                 }
  55                 targ := &PexMsg{
  56                         Added: krpc.CompactIPv4NodeAddrs{
  57                                 addrs[1],
  58                                 addrs[2],
  59                                 addrs[3],
  60                         },
  61                         AddedFlags: []PexPeerFlags{f, f, f},
  62                         Dropped:    krpc.CompactIPv4NodeAddrs{},
  63                 }
  64                 require.EqualValues(t, targ, m)
  65         })
  66         t.Run("ipv6", func(t *testing.T) {
  67                 addrs := addrs6
  68                 m := new(PexMsg)
  69                 m.Drop(addrs[0])
  70                 m.Add(addrs[1], f)
  71                 for _, addr := range addrs {
  72                         m.Add(addr, f)
  73                 }
  74                 targ := &PexMsg{
  75                         Added6: krpc.CompactIPv6NodeAddrs{
  76                                 addrs[1],
  77                                 addrs[2],
  78                                 addrs[3],
  79                         },
  80                         Added6Flags: []PexPeerFlags{f, f, f},
  81                         Dropped6:    krpc.CompactIPv6NodeAddrs{},
  82                 }
  83                 require.EqualValues(t, targ, m)
  84         })
  85         t.Run("empty", func(t *testing.T) {
  86                 addr := krpc.NodeAddr{}
  87                 xm := new(PexMsg)
  88                 xm.Add(addr, f)
  89                 require.EqualValues(t, 0, len(xm.Added))
  90                 require.EqualValues(t, 0, len(xm.AddedFlags))
  91                 require.EqualValues(t, 0, len(xm.Added6))
  92                 require.EqualValues(t, 0, len(xm.Added6Flags))
  93         })
  94 }
  95
  96 func TestPexDrop(t *testing.T) {
  97         addrs4 := []krpc.NodeAddr{
  98                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}, // 0
  99                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748}, // 1
 100                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4747}, // 2
 101                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4748}, // 3
 102         }
 103         addrs6 := []krpc.NodeAddr{
 104                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}, // 0
 105                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4748}, // 1
 106                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4749}, // 2
 107                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4750}, // 3
 108         }
 109         f := PexPrefersEncryption | PexOutgoingConn
 110
 111         t.Run("ipv4", func(t *testing.T) {
 112                 addrs := addrs4
 113                 m := new(PexMsg)
 114                 m.Add(addrs[0], f)
 115                 m.Drop(addrs[1])
 116                 for _, addr := range addrs {
 117                         m.Drop(addr)
 118                 }
 119                 targ := &PexMsg{
 120                         AddedFlags: []PexPeerFlags{},
 121                         Added:    krpc.CompactIPv4NodeAddrs{},
 122                         Dropped: krpc.CompactIPv4NodeAddrs{
 123                                 addrs[1],
 124                                 addrs[2],
 125                                 addrs[3],
 126                         },
 127                 }
 128                 require.EqualValues(t, targ, m)
 129         })
 130         t.Run("ipv6", func(t *testing.T) {
 131                 addrs := addrs6
 132                 m := new(PexMsg)
 133                 m.Add(addrs[0], f)
 134                 m.Drop(addrs[1])
 135                 for _, addr := range addrs {
 136                         m.Drop(addr)
 137                 }
 138                 targ := &PexMsg{
 139                         Added6Flags: []PexPeerFlags{},
 140                         Added6:    krpc.CompactIPv6NodeAddrs{},
 141                         Dropped6: krpc.CompactIPv6NodeAddrs{
 142                                 addrs[1],
 143                                 addrs[2],
 144                                 addrs[3],
 145                         },
 146                 }
 147                 require.EqualValues(t, targ, m)
 148         })
 149         t.Run("empty", func(t *testing.T) {
 150                 addr := krpc.NodeAddr{}
 151                 xm := new(PexMsg)
 152                 xm.Drop(addr)
 153                 require.EqualValues(t, 0, len(xm.Dropped))
 154                 require.EqualValues(t, 0, len(xm.Dropped6))
 155         })
 156 }
 157
 158 func TestMarshalPexMessage(t *testing.T) {
 159         addr := krpc.NodeAddr{IP: net.IP{127, 0, 0, 1}, Port: 0x55aa}
 160         f := PexPrefersEncryption | PexOutgoingConn
 161         pm := new(PexMsg)
 162         pm.Added = append(pm.Added, addr)
 163         pm.AddedFlags = append(pm.AddedFlags, f)
 164
 165         b, err := bencode.Marshal(pm)
 166         require.NoError(t, err)
 167
 168         pexExtendedId := ExtensionNumber(7)
 169         msg := pm.Message(pexExtendedId)
 170         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")
 171         b, err = msg.MarshalBinary()
 172         require.NoError(t, err)
 173         require.EqualValues(t, b, expected)
 174
 175         msg = Message{}
 176         dec := Decoder{
 177                 R:         bufio.NewReader(bytes.NewBuffer(b)),
 178                 MaxLength: 128,
 179         }
 180         pmOut := PexMsg{}
 181         err = dec.Decode(&msg)
 182         require.NoError(t, err)
 183         require.EqualValues(t, Extended, msg.Type)
 184         require.EqualValues(t, pexExtendedId, msg.ExtendedID)
 185         err = bencode.Unmarshal(msg.ExtendedPayload, &pmOut)
 186         require.NoError(t, err)
 187         require.EqualValues(t, len(pm.Added), len(pmOut.Added))
 188         require.EqualValues(t, pm.Added[0].IP, pmOut.Added[0].IP)
 189         require.EqualValues(t, pm.Added[0].Port, pmOut.Added[0].Port)
 190 }