pex_test.go

   1 package torrent
   2
   3 import (
   4         "net"
   5         "testing"
   6
   7         "github.com/stretchr/testify/assert"
   8         "github.com/stretchr/testify/require"
   9
  10         "github.com/anacrolix/dht/v2/krpc"
  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         f = pp.PexOutgoingConn
  34 )
  35
  36 func TestPexAdded(t *testing.T) {
  37         t.Run("noHold", func(t *testing.T) {
  38                 s := new(pexState)
  39                 s.Add(&PeerConn{peer: peer{RemoteAddr: addrs[0], outgoing: true}})
  40                 targ := &pexState{
  41                         ev: []pexEvent{
  42                                 pexEvent{pexAdd, addrs[0], pp.PexOutgoingConn},
  43                         },
  44                         nc: 1,
  45                 }
  46                 require.EqualValues(t, targ, s)
  47         })
  48         t.Run("belowTarg", func(t *testing.T) {
  49                 s := &pexState{
  50                         hold: []pexEvent{
  51                                 pexEvent{pexDrop, addrs[1], 0},
  52                         },
  53                         nc: 0,
  54                 }
  55                 s.Add(&PeerConn{peer: peer{RemoteAddr: addrs[0]}})
  56                 targ := &pexState{
  57                         hold: []pexEvent{
  58                                 pexEvent{pexDrop, addrs[1], 0},
  59                         },
  60                         ev: []pexEvent{
  61                                 pexEvent{pexAdd, addrs[0], 0},
  62                         },
  63                         nc: 1,
  64                 }
  65                 require.EqualValues(t, targ, s)
  66         })
  67         t.Run("aboveTarg", func(t *testing.T) {
  68                 holdAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4848}
  69                 s := &pexState{
  70                         hold: []pexEvent{
  71                                 pexEvent{pexDrop, holdAddr, 0},
  72                         },
  73                         nc: pexTargAdded,
  74                 }
  75                 s.Add(&PeerConn{peer: peer{RemoteAddr: addrs[0]}})
  76                 targ := &pexState{
  77                         hold: []pexEvent{},
  78                         ev: []pexEvent{
  79                                 pexEvent{pexDrop, holdAddr, 0},
  80                                 pexEvent{pexAdd, addrs[0], 0},
  81                         },
  82                         nc: pexTargAdded + 1,
  83                 }
  84                 require.EqualValues(t, targ, s)
  85         })
  86 }
  87
  88 func TestPexDropped(t *testing.T) {
  89         t.Run("belowTarg", func(t *testing.T) {
  90                 s := &pexState{nc: 1}
  91                 s.Drop(&PeerConn{peer: peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
  92                 targ := &pexState{
  93                         hold: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
  94                         nc:   0,
  95                 }
  96                 require.EqualValues(t, targ, s)
  97         })
  98         t.Run("aboveTarg", func(t *testing.T) {
  99                 s := &pexState{nc: pexTargAdded + 1}
 100                 s.Drop(&PeerConn{peer: peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
 101                 targ := &pexState{
 102                         ev: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
 103                         nc: pexTargAdded,
 104                 }
 105                 require.EqualValues(t, targ, s)
 106         })
 107         t.Run("aboveTargNotListed", func(t *testing.T) {
 108                 s := &pexState{nc: pexTargAdded + 1}
 109                 s.Drop(&PeerConn{peer: peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: false}})
 110                 targ := &pexState{nc: pexTargAdded + 1}
 111                 require.EqualValues(t, targ, s)
 112         })
 113 }
 114
 115 func TestPexReset(t *testing.T) {
 116         s := &pexState{
 117                 hold: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
 118                 ev:   []pexEvent{pexEvent{pexAdd, addrs[1], 0}},
 119                 nc:   1,
 120         }
 121         s.Reset()
 122         targ := new(pexState)
 123         require.EqualValues(t, targ, s)
 124 }
 125
 126 func mustNodeAddr(addr net.Addr) krpc.NodeAddr {
 127         ret, ok := nodeAddr(addr)
 128         if !ok {
 129                 panic(addr)
 130         }
 131         return ret
 132 }
 133
 134 var testcases = []struct {
 135         name  string
 136         in    *pexState
 137         arg   int
 138         targM pp.PexMsg
 139         targS int
 140 }{
 141         {
 142                 name:  "empty",
 143                 in:    &pexState{},
 144                 arg:   0,
 145                 targM: pp.PexMsg{},
 146                 targS: 0,
 147         },
 148         {
 149                 name: "add4",
 150                 in: &pexState{
 151                         ev: []pexEvent{
 152                                 pexEvent{pexAdd, addrs[0], f},
 153                                 pexEvent{pexAdd, addrs[1], f},
 154                                 pexEvent{pexAdd, addrs[2], f},
 155                                 pexEvent{pexAdd, addrs[3], f},
 156                         },
 157                 },
 158                 arg: 0,
 159                 targM: pp.PexMsg{
 160                         Added: krpc.CompactIPv4NodeAddrs{
 161                                 mustNodeAddr(addrs[2]),
 162                                 mustNodeAddr(addrs[3]),
 163                         },
 164                         AddedFlags: []pp.PexPeerFlags{f, f},
 165                         Added6: krpc.CompactIPv6NodeAddrs{
 166                                 mustNodeAddr(addrs[0]),
 167                                 mustNodeAddr(addrs[1]),
 168                         },
 169                         Added6Flags: []pp.PexPeerFlags{f, f},
 170                 },
 171                 targS: 4,
 172         },
 173         {
 174                 name: "drop2",
 175                 arg:  0,
 176                 in: &pexState{
 177                         ev: []pexEvent{
 178                                 pexEvent{pexDrop, addrs[0], f},
 179                                 pexEvent{pexDrop, addrs[2], f},
 180                         },
 181                 },
 182                 targM: pp.PexMsg{
 183                         Dropped: krpc.CompactIPv4NodeAddrs{
 184                                 mustNodeAddr(addrs[2]),
 185                         },
 186                         Dropped6: krpc.CompactIPv6NodeAddrs{
 187                                 mustNodeAddr(addrs[0]),
 188                         },
 189                 },
 190                 targS: 2,
 191         },
 192         {
 193                 name: "add2drop1",
 194                 arg:  0,
 195                 in: &pexState{
 196                         ev: []pexEvent{
 197                                 pexEvent{pexAdd, addrs[0], f},
 198                                 pexEvent{pexAdd, addrs[1], f},
 199                                 pexEvent{pexDrop, addrs[0], f},
 200                         },
 201                 },
 202                 targM: pp.PexMsg{
 203                         Added6: krpc.CompactIPv6NodeAddrs{
 204                                 mustNodeAddr(addrs[1]),
 205                         },
 206                         Added6Flags: []pp.PexPeerFlags{f},
 207                 },
 208                 targS: 3,
 209         },
 210         {
 211                 name: "delayed",
 212                 arg:  0,
 213                 in: &pexState{
 214                         ev: []pexEvent{
 215                                 pexEvent{pexAdd, addrs[0], f},
 216                                 pexEvent{pexAdd, addrs[1], f},
 217                                 pexEvent{pexAdd, addrs[2], f},
 218                         },
 219                         hold: []pexEvent{
 220                                 pexEvent{pexDrop, addrs[0], f},
 221                                 pexEvent{pexDrop, addrs[2], f},
 222                                 pexEvent{pexDrop, addrs[1], f},
 223                         },
 224                 },
 225                 targM: pp.PexMsg{
 226                         Added: krpc.CompactIPv4NodeAddrs{
 227                                 mustNodeAddr(addrs[2]),
 228                         },
 229                         AddedFlags: []pp.PexPeerFlags{f},
 230                         Added6: krpc.CompactIPv6NodeAddrs{
 231                                 mustNodeAddr(addrs[0]),
 232                                 mustNodeAddr(addrs[1]),
 233                         },
 234                         Added6Flags: []pp.PexPeerFlags{f, f},
 235                 },
 236                 targS: 3,
 237         },
 238         {
 239                 name: "followup",
 240                 arg:  1,
 241                 in: &pexState{
 242                         ev: []pexEvent{
 243                                 pexEvent{pexAdd, addrs[0], f},
 244                                 pexEvent{pexAdd, addrs[1], f},
 245                         },
 246                 },
 247                 targM: pp.PexMsg{
 248                         Added6: krpc.CompactIPv6NodeAddrs{
 249                                 mustNodeAddr(addrs[1]),
 250                         },
 251                         Added6Flags: []pp.PexPeerFlags{f},
 252                 },
 253                 targS: 2,
 254         },
 255 }
 256
 257 // Represents the contents of a PexMsg in a way that supports equivalence checking in tests. This is
 258 // necessary because pexMsgFactory uses maps and so ordering of the resultant PexMsg isn't
 259 // deterministic. Because the flags are in a different array, we can't just use testify's
 260 // ElementsMatch because the ordering *does* still matter between an added addr and its flags.
 261 type comparablePexMsg struct {
 262         added, added6           []krpc.NodeAddr
 263         addedFlags, added6Flags []pp.PexPeerFlags
 264         dropped, dropped6       []krpc.NodeAddr
 265 }
 266
 267 // Such Rust-inspired.
 268 func (me *comparablePexMsg) From(f pp.PexMsg) {
 269         me.added = f.Added
 270         me.addedFlags = f.AddedFlags
 271         me.added6 = f.Added6
 272         me.added6Flags = f.Added6Flags
 273         me.dropped = f.Dropped
 274         me.dropped6 = f.Dropped6
 275 }
 276
 277 // For PexMsg created by pexMsgFactory, this is as good as it can get without using data structures
 278 // in pexMsgFactory that preserve insert ordering.
 279 func (actual comparablePexMsg) AssertEqual(t *testing.T, expected comparablePexMsg) {
 280         assert.ElementsMatch(t, expected.added, actual.added)
 281         assert.ElementsMatch(t, expected.addedFlags, actual.addedFlags)
 282         assert.ElementsMatch(t, expected.added6, actual.added6)
 283         assert.ElementsMatch(t, expected.added6Flags, actual.added6Flags)
 284         assert.ElementsMatch(t, expected.dropped, actual.dropped)
 285         assert.ElementsMatch(t, expected.dropped6, actual.dropped6)
 286 }
 287
 288 func assertPexMsgsEqual(t *testing.T, expected, actual pp.PexMsg) {
 289         var ec, ac comparablePexMsg
 290         ec.From(expected)
 291         ac.From(actual)
 292         ac.AssertEqual(t, ec)
 293 }
 294
 295 func TestPexGenmsg(t *testing.T) {
 296         for _, tc := range testcases {
 297                 t.Run(tc.name, func(t *testing.T) {
 298                         s := tc.in
 299                         m, seen := s.Genmsg(tc.arg)
 300                         assertPexMsgsEqual(t, tc.targM, m)
 301                         require.EqualValues(t, tc.targS, seen)
 302                 })
 303         }
 304 }
 305
 306 // generate 𝑛 distinct values of net.Addr
 307 func addrgen(n int) chan net.Addr {
 308         c := make(chan net.Addr)
 309         go func() {
 310                 for i := 4747; i < 65535 && n > 0; i++ {
 311                         c <- &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: i}
 312                         n--
 313                 }
 314                 close(c)
 315         }()
 316         return c
 317 }
 318
 319 func TestPexInitialNoCutoff(t *testing.T) {
 320         const n = 2 * pexMaxDelta
 321         var s pexState
 322
 323         c := addrgen(n)
 324         for addr := range c {
 325                 s.Add(&PeerConn{peer: peer{RemoteAddr: addr}})
 326         }
 327         m, seq := s.Genmsg(0)
 328
 329         require.EqualValues(t, n, seq)
 330         require.EqualValues(t, n, len(m.Added))
 331         require.EqualValues(t, n, len(m.AddedFlags))
 332         require.EqualValues(t, 0, len(m.Added6))
 333         require.EqualValues(t, 0, len(m.Added6Flags))
 334         require.EqualValues(t, 0, len(m.Dropped))
 335         require.EqualValues(t, 0, len(m.Dropped6))
 336 }
 337
 338 func TestPexAdd(t *testing.T) {
 339         t.Run("ipv4", func(t *testing.T) {
 340                 addrs := addrs4
 341                 var m pexMsgFactory
 342                 m.addEvent(pexEvent{pexDrop, addrs[0], 0})
 343                 m.addEvent(pexEvent{pexAdd, addrs[1], f})
 344                 for _, addr := range addrs {
 345                         m.addEvent(pexEvent{pexAdd, addr, f})
 346                 }
 347                 targ := pp.PexMsg{
 348                         Added: krpc.CompactIPv4NodeAddrs{
 349                                 mustNodeAddr(addrs[1]),
 350                                 mustNodeAddr(addrs[2]),
 351                                 mustNodeAddr(addrs[3]),
 352                         },
 353                         AddedFlags: []pp.PexPeerFlags{f, f, f},
 354                 }
 355                 out := m.PexMsg()
 356                 assertPexMsgsEqual(t, targ, out)
 357         })
 358         t.Run("ipv6", func(t *testing.T) {
 359                 addrs := addrs6
 360                 var m pexMsgFactory
 361                 m.addEvent(pexEvent{pexDrop, addrs[0], 0})
 362                 m.addEvent(pexEvent{pexAdd, addrs[1], f})
 363                 for _, addr := range addrs {
 364                         m.addEvent(pexEvent{pexAdd, addr, f})
 365                 }
 366                 targ := pp.PexMsg{
 367                         Added6: krpc.CompactIPv6NodeAddrs{
 368                                 mustNodeAddr(addrs[1]),
 369                                 mustNodeAddr(addrs[2]),
 370                                 mustNodeAddr(addrs[3]),
 371                         },
 372                         Added6Flags: []pp.PexPeerFlags{f, f, f},
 373                 }
 374                 assertPexMsgsEqual(t, targ, m.PexMsg())
 375         })
 376         t.Run("empty", func(t *testing.T) {
 377                 nullAddr := &net.TCPAddr{}
 378                 var xm pexMsgFactory
 379                 xm.addEvent(pexEvent{pexAdd, nullAddr, f})
 380                 m := xm.PexMsg()
 381                 require.EqualValues(t, 0, len(m.Added))
 382                 require.EqualValues(t, 0, len(m.AddedFlags))
 383                 require.EqualValues(t, 0, len(m.Added6))
 384                 require.EqualValues(t, 0, len(m.Added6Flags))
 385         })
 386 }
 387
 388 func TestPexDrop(t *testing.T) {
 389         t.Run("ipv4", func(t *testing.T) {
 390                 addrs := addrs4
 391                 var m pexMsgFactory
 392                 m.addEvent(pexEvent{pexAdd, addrs[0], f})
 393                 m.addEvent(pexEvent{pexDrop, addrs[1], 0})
 394                 for _, addr := range addrs {
 395                         m.addEvent(pexEvent{pexDrop, addr, 0})
 396                 }
 397                 targ := pp.PexMsg{
 398                         Dropped: krpc.CompactIPv4NodeAddrs{
 399                                 mustNodeAddr(addrs[1]),
 400                                 mustNodeAddr(addrs[2]),
 401                                 mustNodeAddr(addrs[3]),
 402                         },
 403                 }
 404                 assertPexMsgsEqual(t, targ, m.PexMsg())
 405         })
 406         t.Run("ipv6", func(t *testing.T) {
 407                 addrs := addrs6
 408                 var m pexMsgFactory
 409                 m.addEvent(pexEvent{pexAdd, addrs[0], f})
 410                 m.addEvent(pexEvent{pexDrop, addrs[1], 0})
 411                 for _, addr := range addrs {
 412                         m.addEvent(pexEvent{pexDrop, addr, 0})
 413                 }
 414                 targ := pp.PexMsg{
 415                         Dropped6: krpc.CompactIPv6NodeAddrs{
 416                                 mustNodeAddr(addrs[1]),
 417                                 mustNodeAddr(addrs[2]),
 418                                 mustNodeAddr(addrs[3]),
 419                         },
 420                 }
 421                 assertPexMsgsEqual(t, targ, m.PexMsg())
 422         })
 423         t.Run("empty", func(t *testing.T) {
 424                 nullAddr := &net.TCPAddr{}
 425                 var xm pexMsgFactory
 426                 xm.addEvent(pexEvent{pexDrop, nullAddr, f})
 427                 m := xm.PexMsg()
 428                 require.EqualValues(t, 0, len(m.Dropped))
 429                 require.EqualValues(t, 0, len(m.Dropped6))
 430         })
 431 }