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 func TestPexAdd(t *testing.T) {
 292         addrs4 := []krpc.NodeAddr{
 293                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}, // 0
 294                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748}, // 1
 295                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4747}, // 2
 296                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4748}, // 3
 297         }
 298         addrs6 := []krpc.NodeAddr{
 299                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}, // 0
 300                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4748}, // 1
 301                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4749}, // 2
 302                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4750}, // 3
 303         }
 304         f := pp.PexPrefersEncryption | pp.PexOutgoingConn
 305
 306         t.Run("ipv4", func(t *testing.T) {
 307                 addrs := addrs4
 308                 var m pexMsgFactory
 309                 m.Drop(addrs[0])
 310                 m.Add(addrs[1], f)
 311                 for _, addr := range addrs {
 312                         m.Add(addr, f)
 313                 }
 314                 targ := pp.PexMsg{
 315                         Added: krpc.CompactIPv4NodeAddrs{
 316                                 addrs[1],
 317                                 addrs[2],
 318                                 addrs[3],
 319                         },
 320                         AddedFlags: []pp.PexPeerFlags{f, f, f},
 321                 }
 322                 assertPexMsgsEqual(t, targ, m.PexMsg())
 323         })
 324         t.Run("ipv6", func(t *testing.T) {
 325                 addrs := addrs6
 326                 var m pexMsgFactory
 327                 m.Drop(addrs[0])
 328                 m.Add(addrs[1], f)
 329                 for _, addr := range addrs {
 330                         m.Add(addr, f)
 331                 }
 332                 targ := pp.PexMsg{
 333                         Added6: krpc.CompactIPv6NodeAddrs{
 334                                 addrs[1],
 335                                 addrs[2],
 336                                 addrs[3],
 337                         },
 338                         Added6Flags: []pp.PexPeerFlags{f, f, f},
 339                 }
 340                 assertPexMsgsEqual(t, targ, m.PexMsg())
 341         })
 342         t.Run("empty", func(t *testing.T) {
 343                 addr := krpc.NodeAddr{}
 344                 var xm pexMsgFactory
 345                 assert.Panics(t, func() { xm.Add(addr, f) })
 346                 m := xm.PexMsg()
 347                 require.EqualValues(t, 0, len(m.Added))
 348                 require.EqualValues(t, 0, len(m.AddedFlags))
 349                 require.EqualValues(t, 0, len(m.Added6))
 350                 require.EqualValues(t, 0, len(m.Added6Flags))
 351         })
 352 }
 353
 354 func TestPexDrop(t *testing.T) {
 355         addrs4 := []krpc.NodeAddr{
 356                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747}, // 0
 357                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748}, // 1
 358                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4747}, // 2
 359                 krpc.NodeAddr{IP: net.IPv4(127, 0, 0, 2), Port: 4748}, // 3
 360         }
 361         addrs6 := []krpc.NodeAddr{
 362                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4747}, // 0
 363                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4748}, // 1
 364                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4749}, // 2
 365                 krpc.NodeAddr{IP: net.IPv6loopback, Port: 4750}, // 3
 366         }
 367         f := pp.PexPrefersEncryption | pp.PexOutgoingConn
 368
 369         t.Run("ipv4", func(t *testing.T) {
 370                 addrs := addrs4
 371                 var m pexMsgFactory
 372                 m.Add(addrs[0], f)
 373                 m.Drop(addrs[1])
 374                 for _, addr := range addrs {
 375                         m.Drop(addr)
 376                 }
 377                 targ := pp.PexMsg{
 378                         Dropped: krpc.CompactIPv4NodeAddrs{
 379                                 addrs[1],
 380                                 addrs[2],
 381                                 addrs[3],
 382                         },
 383                 }
 384                 assertPexMsgsEqual(t, targ, m.PexMsg())
 385         })
 386         t.Run("ipv6", func(t *testing.T) {
 387                 addrs := addrs6
 388                 var m pexMsgFactory
 389                 m.Add(addrs[0], f)
 390                 m.Drop(addrs[1])
 391                 for _, addr := range addrs {
 392                         m.Drop(addr)
 393                 }
 394                 targ := pp.PexMsg{
 395                         Dropped6: krpc.CompactIPv6NodeAddrs{
 396                                 addrs[1],
 397                                 addrs[2],
 398                                 addrs[3],
 399                         },
 400                 }
 401                 assertPexMsgsEqual(t, targ, m.PexMsg())
 402         })
 403         t.Run("empty", func(t *testing.T) {
 404                 addr := krpc.NodeAddr{}
 405                 var xm pexMsgFactory
 406                 require.Panics(t, func() { xm.Drop(addr) })
 407                 m := xm.PexMsg()
 408                 require.EqualValues(t, 0, len(m.Dropped))
 409                 require.EqualValues(t, 0, len(m.Dropped6))
 410         })
 411 }