pex_test.go

   1 package torrent
   2
   3 import (
   4         "net"
   5         "testing"
   6
   7         "github.com/anacrolix/dht/v2/krpc"
   8         "github.com/stretchr/testify/assert"
   9         "github.com/stretchr/testify/require"
  10
  11         pp "github.com/anacrolix/torrent/peer_protocol"
  12 )
  13
  14 var (
  15         addrs6 = []net.Addr{
  16                 &net.TCPAddr{IP: net.IPv6loopback, Port: 4747},
  17                 &net.TCPAddr{IP: net.IPv6loopback, Port: 4748},
  18                 &net.TCPAddr{IP: net.IPv6loopback, Port: 4749},
  19                 &net.TCPAddr{IP: net.IPv6loopback, Port: 4750},
  20         }
  21         addrs4 = []net.Addr{
  22                 &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747},
  23                 &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748},
  24                 &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4749},
  25                 &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4750},
  26         }
  27         addrs = []net.Addr{
  28                 addrs6[0],
  29                 addrs6[1],
  30                 addrs4[0],
  31                 addrs4[1],
  32         }
  33 )
  34
  35 func TestPexReset(t *testing.T) {
  36         s := &pexState{}
  37         conns := []PeerConn{
  38                 {Peer: Peer{RemoteAddr: addrs[0]}},
  39                 {Peer: Peer{RemoteAddr: addrs[1]}},
  40                 {Peer: Peer{RemoteAddr: addrs[2]}},
  41         }
  42         s.Add(&conns[0])
  43         s.Add(&conns[1])
  44         s.Drop(&conns[0])
  45         s.Reset()
  46         targ := new(pexState)
  47         require.EqualValues(t, targ, s)
  48 }
  49
  50 func krpcNodeAddrFromNetAddr(addr net.Addr) krpc.NodeAddr {
  51         addrPort, err := addrPortFromPeerRemoteAddr(addr)
  52         if err != nil {
  53                 panic(err)
  54         }
  55         return krpcNodeAddrFromAddrPort(addrPort)
  56 }
  57
  58 var testcases = []struct {
  59         name   string
  60         in     *pexState
  61         targ   pp.PexMsg
  62         update func(*pexState)
  63         targ1  pp.PexMsg
  64 }{
  65         {
  66                 name: "empty",
  67                 in:   &pexState{},
  68                 targ: pp.PexMsg{},
  69         },
  70         {
  71                 name: "add0",
  72                 in: func() *pexState {
  73                         s := new(pexState)
  74                         nullAddr := &net.TCPAddr{}
  75                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: nullAddr}})
  76                         return s
  77                 }(),
  78                 targ: pp.PexMsg{},
  79         },
  80         {
  81                 name: "drop0",
  82                 in: func() *pexState {
  83                         nullAddr := &net.TCPAddr{}
  84                         s := new(pexState)
  85                         s.Drop(&PeerConn{Peer: Peer{RemoteAddr: nullAddr}, pex: pexConnState{Listed: true}})
  86                         return s
  87                 }(),
  88                 targ: pp.PexMsg{},
  89         },
  90         {
  91                 name: "add4",
  92                 in: func() *pexState {
  93                         s := new(pexState)
  94                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}})
  95                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1], outgoing: true}})
  96                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[2], outgoing: true}})
  97                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[3]}})
  98                         return s
  99                 }(),
 100                 targ: pp.PexMsg{
 101                         Added: krpc.CompactIPv4NodeAddrs{
 102                                 krpcNodeAddrFromNetAddr(addrs[2]),
 103                                 krpcNodeAddrFromNetAddr(addrs[3]),
 104                         },
 105                         AddedFlags: []pp.PexPeerFlags{pp.PexOutgoingConn, 0},
 106                         Added6: krpc.CompactIPv6NodeAddrs{
 107                                 krpcNodeAddrFromNetAddr(addrs[0]),
 108                                 krpcNodeAddrFromNetAddr(addrs[1]),
 109                         },
 110                         Added6Flags: []pp.PexPeerFlags{0, pp.PexOutgoingConn},
 111                 },
 112         },
 113         {
 114                 name: "drop2",
 115                 in: func() *pexState {
 116                         s := &pexState{nc: pexTargAdded + 2}
 117                         s.Drop(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
 118                         s.Drop(&PeerConn{Peer: Peer{RemoteAddr: addrs[2]}, pex: pexConnState{Listed: true}})
 119                         return s
 120                 }(),
 121                 targ: pp.PexMsg{
 122                         Dropped: krpc.CompactIPv4NodeAddrs{
 123                                 krpcNodeAddrFromNetAddr(addrs[2]),
 124                         },
 125                         Dropped6: krpc.CompactIPv6NodeAddrs{
 126                                 krpcNodeAddrFromNetAddr(addrs[0]),
 127                         },
 128                 },
 129         },
 130         {
 131                 name: "add2drop1",
 132                 in: func() *pexState {
 133                         conns := []PeerConn{
 134                                 {Peer: Peer{RemoteAddr: addrs[0]}},
 135                                 {Peer: Peer{RemoteAddr: addrs[1]}},
 136                                 {Peer: Peer{RemoteAddr: addrs[2]}},
 137                         }
 138                         s := &pexState{nc: pexTargAdded}
 139                         s.Add(&conns[0])
 140                         s.Add(&conns[1])
 141                         s.Drop(&conns[0])
 142                         s.Drop(&conns[2]) // to be ignored: it wasn't added
 143                         return s
 144                 }(),
 145                 targ: pp.PexMsg{
 146                         Added6: krpc.CompactIPv6NodeAddrs{
 147                                 krpcNodeAddrFromNetAddr(addrs[1]),
 148                         },
 149                         Added6Flags: []pp.PexPeerFlags{0},
 150                 },
 151         },
 152         {
 153                 name: "delayed",
 154                 in: func() *pexState {
 155                         conns := []PeerConn{
 156                                 {Peer: Peer{RemoteAddr: addrs[0]}},
 157                                 {Peer: Peer{RemoteAddr: addrs[1]}},
 158                                 {Peer: Peer{RemoteAddr: addrs[2]}},
 159                         }
 160                         s := new(pexState)
 161                         s.Add(&conns[0])
 162                         s.Add(&conns[1])
 163                         s.Add(&conns[2])
 164                         s.Drop(&conns[0]) // on hold: s.nc < pexTargAdded
 165                         s.Drop(&conns[2])
 166                         s.Drop(&conns[1])
 167                         return s
 168                 }(),
 169                 targ: pp.PexMsg{
 170                         Added: krpc.CompactIPv4NodeAddrs{
 171                                 krpcNodeAddrFromNetAddr(addrs[2]),
 172                         },
 173                         AddedFlags: []pp.PexPeerFlags{0},
 174                         Added6: krpc.CompactIPv6NodeAddrs{
 175                                 krpcNodeAddrFromNetAddr(addrs[0]),
 176                                 krpcNodeAddrFromNetAddr(addrs[1]),
 177                         },
 178                         Added6Flags: []pp.PexPeerFlags{0, 0},
 179                 },
 180         },
 181         {
 182                 name: "unheld",
 183                 in: func() *pexState {
 184                         conns := []PeerConn{
 185                                 {Peer: Peer{RemoteAddr: addrs[0]}},
 186                                 {Peer: Peer{RemoteAddr: addrs[1]}},
 187                         }
 188                         s := &pexState{nc: pexTargAdded - 1}
 189                         s.Add(&conns[0])
 190                         s.Drop(&conns[0]) // on hold: s.nc < pexTargAdded
 191                         s.Add(&conns[1])  // unholds the above
 192                         return s
 193                 }(),
 194                 targ: pp.PexMsg{
 195                         Added6: krpc.CompactIPv6NodeAddrs{
 196                                 krpcNodeAddrFromNetAddr(addrs[1]),
 197                         },
 198                         Added6Flags: []pp.PexPeerFlags{0},
 199                 },
 200         },
 201         {
 202                 name: "followup",
 203                 in: func() *pexState {
 204                         s := new(pexState)
 205                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}})
 206                         return s
 207                 }(),
 208                 targ: pp.PexMsg{
 209                         Added6: krpc.CompactIPv6NodeAddrs{
 210                                 krpcNodeAddrFromNetAddr(addrs[0]),
 211                         },
 212                         Added6Flags: []pp.PexPeerFlags{0},
 213                 },
 214                 update: func(s *pexState) {
 215                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1]}})
 216                 },
 217                 targ1: pp.PexMsg{
 218                         Added6: krpc.CompactIPv6NodeAddrs{
 219                                 krpcNodeAddrFromNetAddr(addrs[1]),
 220                         },
 221                         Added6Flags: []pp.PexPeerFlags{0},
 222                 },
 223         },
 224 }
 225
 226 // Represents the contents of a PexMsg in a way that supports equivalence checking in tests. This is
 227 // necessary because pexMsgFactory uses maps and so ordering of the resultant PexMsg isn't
 228 // deterministic. Because the flags are in a different array, we can't just use testify's
 229 // ElementsMatch because the ordering *does* still matter between an added addr and its flags.
 230 type comparablePexMsg struct {
 231         added, added6           []krpc.NodeAddr
 232         addedFlags, added6Flags []pp.PexPeerFlags
 233         dropped, dropped6       []krpc.NodeAddr
 234 }
 235
 236 // Such Rust-inspired.
 237 func (me *comparablePexMsg) From(f pp.PexMsg) {
 238         me.added = f.Added
 239         me.addedFlags = f.AddedFlags
 240         me.added6 = f.Added6
 241         me.added6Flags = f.Added6Flags
 242         me.dropped = f.Dropped
 243         me.dropped6 = f.Dropped6
 244 }
 245
 246 // For PexMsg created by pexMsgFactory, this is as good as it can get without using data structures
 247 // in pexMsgFactory that preserve insert ordering.
 248 func (actual comparablePexMsg) AssertEqual(t *testing.T, expected comparablePexMsg) {
 249         assert.ElementsMatch(t, expected.added, actual.added)
 250         assert.ElementsMatch(t, expected.addedFlags, actual.addedFlags)
 251         assert.ElementsMatch(t, expected.added6, actual.added6)
 252         assert.ElementsMatch(t, expected.added6Flags, actual.added6Flags)
 253         assert.ElementsMatch(t, expected.dropped, actual.dropped)
 254         assert.ElementsMatch(t, expected.dropped6, actual.dropped6)
 255 }
 256
 257 func assertPexMsgsEqual(t *testing.T, expected, actual pp.PexMsg) {
 258         var ec, ac comparablePexMsg
 259         ec.From(expected)
 260         ac.From(actual)
 261         ac.AssertEqual(t, ec)
 262 }
 263
 264 func TestPexGenmsg0(t *testing.T) {
 265         for _, tc := range testcases {
 266                 t.Run(tc.name, func(t *testing.T) {
 267                         s := *tc.in
 268                         m, last := s.Genmsg(nil)
 269                         assertPexMsgsEqual(t, tc.targ, m)
 270                         if tc.update != nil {
 271                                 tc.update(&s)
 272                                 m1, last := s.Genmsg(last)
 273                                 assertPexMsgsEqual(t, tc.targ1, m1)
 274                                 assert.NotNil(t, last)
 275                         }
 276                 })
 277         }
 278 }
 279
 280 // generate 𝑛 distinct values of net.Addr
 281 func addrgen(n int) chan net.Addr {
 282         c := make(chan net.Addr)
 283         go func() {
 284                 defer close(c)
 285                 for i := 4747; i < 65535 && n > 0; i++ {
 286                         c <- &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: i}
 287                         n--
 288                 }
 289         }()
 290         return c
 291 }
 292
 293 func TestPexInitialNoCutoff(t *testing.T) {
 294         const n = 2 * pexMaxDelta
 295         var s pexState
 296
 297         c := addrgen(n)
 298         for addr := range c {
 299                 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addr}})
 300         }
 301         m, _ := s.Genmsg(nil)
 302
 303         require.EqualValues(t, n, len(m.Added))
 304         require.EqualValues(t, n, len(m.AddedFlags))
 305         require.EqualValues(t, 0, len(m.Added6))
 306         require.EqualValues(t, 0, len(m.Added6Flags))
 307         require.EqualValues(t, 0, len(m.Dropped))
 308         require.EqualValues(t, 0, len(m.Dropped6))
 309 }
 310
 311 func benchmarkPexInitialN(b *testing.B, npeers int) {
 312         for i := 0; i < b.N; i++ {
 313                 var s pexState
 314                 c := addrgen(npeers)
 315                 for addr := range c {
 316                         s.Add(&PeerConn{Peer: Peer{RemoteAddr: addr}})
 317                         s.Genmsg(nil)
 318                 }
 319         }
 320 }
 321
 322 // obtain at least 5 points, e.g. to plot a graph
 323 func BenchmarkPexInitial4(b *testing.B)   { benchmarkPexInitialN(b, 4) }
 324 func BenchmarkPexInitial50(b *testing.B)  { benchmarkPexInitialN(b, 50) }
 325 func BenchmarkPexInitial100(b *testing.B) { benchmarkPexInitialN(b, 100) }
 326 func BenchmarkPexInitial200(b *testing.B) { benchmarkPexInitialN(b, 200) }
 327 func BenchmarkPexInitial400(b *testing.B) { benchmarkPexInitialN(b, 400) }