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