pex_test.go

   1 package torrent
   2
   3 import (
   4         "net"
   5         "testing"
   6
   7         "github.com/stretchr/testify/require"
   8
   9         "github.com/anacrolix/dht/v2/krpc"
  10         pp "github.com/anacrolix/torrent/peer_protocol"
  11 )
  12
  13 var (
  14         addrs = []net.Addr{
  15                 &net.TCPAddr{IP: net.IPv6loopback, Port: 4747},
  16                 &net.TCPAddr{IP: net.IPv6loopback, Port: 4748},
  17                 &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747},
  18                 &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748},
  19         }
  20         f = pp.PexOutgoingConn
  21 )
  22
  23 func TestPexAdded(t *testing.T) {
  24         t.Run("noHold", func(t *testing.T) {
  25                 s := new(pexState)
  26                 s.Add(&PeerConn{peer: peer{RemoteAddr: addrs[0], outgoing: true}})
  27                 targ := &pexState{
  28                         ev: []pexEvent{
  29                                 pexEvent{pexAdd, addrs[0], pp.PexOutgoingConn},
  30                         },
  31                         nc: 1,
  32                 }
  33                 require.EqualValues(t, targ, s)
  34         })
  35         t.Run("belowTarg", func(t *testing.T) {
  36                 s := &pexState{
  37                         hold: []pexEvent{
  38                                 pexEvent{pexDrop, addrs[1], 0},
  39                         },
  40                         nc: 0,
  41                 }
  42                 s.Add(&PeerConn{peer: peer{RemoteAddr: addrs[0]}})
  43                 targ := &pexState{
  44                         hold: []pexEvent{
  45                                 pexEvent{pexDrop, addrs[1], 0},
  46                         },
  47                         ev: []pexEvent{
  48                                 pexEvent{pexAdd, addrs[0], 0},
  49                         },
  50                         nc: 1,
  51                 }
  52                 require.EqualValues(t, targ, s)
  53         })
  54         t.Run("aboveTarg", func(t *testing.T) {
  55                 holdAddr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4848}
  56                 s := &pexState{
  57                         hold: []pexEvent{
  58                                 pexEvent{pexDrop, holdAddr, 0},
  59                         },
  60                         nc: pexTargAdded,
  61                 }
  62                 s.Add(&PeerConn{peer: peer{RemoteAddr: addrs[0]}})
  63                 targ := &pexState{
  64                         hold: []pexEvent{},
  65                         ev: []pexEvent{
  66                                 pexEvent{pexDrop, holdAddr, 0},
  67                                 pexEvent{pexAdd, addrs[0], 0},
  68                         },
  69                         nc: pexTargAdded + 1,
  70                 }
  71                 require.EqualValues(t, targ, s)
  72         })
  73 }
  74
  75 func TestPexDropped(t *testing.T) {
  76         t.Run("belowTarg", func(t *testing.T) {
  77                 s := &pexState{nc: 1}
  78                 s.Drop(&PeerConn{peer: peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
  79                 targ := &pexState{
  80                         hold: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
  81                         nc:   0,
  82                 }
  83                 require.EqualValues(t, targ, s)
  84         })
  85         t.Run("aboveTarg", func(t *testing.T) {
  86                 s := &pexState{nc: pexTargAdded + 1}
  87                 s.Drop(&PeerConn{peer: peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
  88                 targ := &pexState{
  89                         ev: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
  90                         nc: pexTargAdded,
  91                 }
  92                 require.EqualValues(t, targ, s)
  93         })
  94         t.Run("aboveTargNotListed", func(t *testing.T) {
  95                 s := &pexState{nc: pexTargAdded + 1}
  96                 s.Drop(&PeerConn{peer: peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: false}})
  97                 targ := &pexState{nc: pexTargAdded + 1}
  98                 require.EqualValues(t, targ, s)
  99         })
 100 }
 101
 102 func TestPexReset(t *testing.T) {
 103         s := &pexState{
 104                 hold: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
 105                 ev:   []pexEvent{pexEvent{pexAdd, addrs[1], 0}},
 106                 nc:   1,
 107         }
 108         s.Reset()
 109         targ := new(pexState)
 110         require.EqualValues(t, targ, s)
 111 }
 112
 113 var testcases = []struct {
 114         name  string
 115         in    *pexState
 116         arg   int
 117         targM *pp.PexMsg
 118         targS int
 119 }{
 120         {
 121                 name:  "empty",
 122                 in:    &pexState{},
 123                 arg:   0,
 124                 targM: &pp.PexMsg{},
 125                 targS: 0,
 126         },
 127         {
 128                 name: "add4",
 129                 in: &pexState{
 130                         ev: []pexEvent{
 131                                 pexEvent{pexAdd, addrs[0], f},
 132                                 pexEvent{pexAdd, addrs[1], f},
 133                                 pexEvent{pexAdd, addrs[2], f},
 134                                 pexEvent{pexAdd, addrs[3], f},
 135                         },
 136                 },
 137                 arg: 0,
 138                 targM: &pp.PexMsg{
 139                         Added: krpc.CompactIPv4NodeAddrs{
 140                                 nodeAddr(addrs[2]),
 141                                 nodeAddr(addrs[3]),
 142                         },
 143                         AddedFlags: []pp.PexPeerFlags{f, f},
 144                         Added6: krpc.CompactIPv6NodeAddrs{
 145                                 nodeAddr(addrs[0]),
 146                                 nodeAddr(addrs[1]),
 147                         },
 148                         Added6Flags: []pp.PexPeerFlags{f, f},
 149                 },
 150                 targS: 4,
 151         },
 152         {
 153                 name: "drop2",
 154                 arg:  0,
 155                 in: &pexState{
 156                         ev: []pexEvent{
 157                                 pexEvent{pexDrop, addrs[0], f},
 158                                 pexEvent{pexDrop, addrs[2], f},
 159                         },
 160                 },
 161                 targM: &pp.PexMsg{
 162                         Dropped: krpc.CompactIPv4NodeAddrs{
 163                                 nodeAddr(addrs[2]),
 164                         },
 165                         Dropped6: krpc.CompactIPv6NodeAddrs{
 166                                 nodeAddr(addrs[0]),
 167                         },
 168                 },
 169                 targS: 2,
 170         },
 171         {
 172                 name: "add2drop1",
 173                 arg:  0,
 174                 in: &pexState{
 175                         ev: []pexEvent{
 176                                 pexEvent{pexAdd, addrs[0], f},
 177                                 pexEvent{pexAdd, addrs[1], f},
 178                                 pexEvent{pexDrop, addrs[0], f},
 179                         },
 180                 },
 181                 targM: &pp.PexMsg{
 182                         Added6: krpc.CompactIPv6NodeAddrs{
 183                                 nodeAddr(addrs[1]),
 184                         },
 185                         Added6Flags: []pp.PexPeerFlags{f},
 186                 },
 187                 targS: 3,
 188         },
 189         {
 190                 name: "delayed",
 191                 arg:  0,
 192                 in: &pexState{
 193                         ev: []pexEvent{
 194                                 pexEvent{pexAdd, addrs[0], f},
 195                                 pexEvent{pexAdd, addrs[1], f},
 196                                 pexEvent{pexAdd, addrs[2], f},
 197                         },
 198                         hold: []pexEvent{
 199                                 pexEvent{pexDrop, addrs[0], f},
 200                                 pexEvent{pexDrop, addrs[2], f},
 201                                 pexEvent{pexDrop, addrs[1], f},
 202                         },
 203                 },
 204                 targM: &pp.PexMsg{
 205                         Added: krpc.CompactIPv4NodeAddrs{
 206                                 nodeAddr(addrs[2]),
 207                         },
 208                         AddedFlags: []pp.PexPeerFlags{f},
 209                         Added6: krpc.CompactIPv6NodeAddrs{
 210                                 nodeAddr(addrs[0]),
 211                                 nodeAddr(addrs[1]),
 212                         },
 213                         Added6Flags: []pp.PexPeerFlags{f, f},
 214                 },
 215                 targS: 3,
 216         },
 217         {
 218                 name: "followup",
 219                 arg:  1,
 220                 in: &pexState{
 221                         ev: []pexEvent{
 222                                 pexEvent{pexAdd, addrs[0], f},
 223                                 pexEvent{pexAdd, addrs[1], f},
 224                         },
 225                 },
 226                 targM: &pp.PexMsg{
 227                         Added6: krpc.CompactIPv6NodeAddrs{
 228                                 nodeAddr(addrs[1]),
 229                         },
 230                         Added6Flags: []pp.PexPeerFlags{f},
 231                 },
 232                 targS: 2,
 233         },
 234 }
 235
 236 func TestPexGenmsg(t *testing.T) {
 237         for _, tc := range testcases {
 238                 t.Run(tc.name, func(t *testing.T) {
 239                         s := tc.in
 240                         m, seen := s.Genmsg(tc.arg)
 241                         require.EqualValues(t, tc.targM, m)
 242                         require.EqualValues(t, tc.targS, seen)
 243                 })
 244         }
 245 }