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{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{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{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{remoteAddr: addrs[0]})
  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{remoteAddr: addrs[0]})
  88                 targ := &pexState{
  89                         ev: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
  90                         nc: pexTargAdded,
  91                 }
  92                 require.EqualValues(t, targ, s)
  93         })
  94 }
  95
  96 func TestPexReset(t *testing.T) {
  97         s := &pexState{
  98                 hold: []pexEvent{pexEvent{pexDrop, addrs[0], 0}},
  99                 ev:   []pexEvent{pexEvent{pexAdd, addrs[1], 0}},
 100                 nc:   1,
 101         }
 102         s.Reset()
 103         targ := new(pexState)
 104         require.EqualValues(t, targ, s)
 105 }
 106
 107 var testcases = []struct {
 108         name  string
 109         in    *pexState
 110         arg   int
 111         targM *pp.PexMsg
 112         targS int
 113 }{
 114         {
 115                 name:  "empty",
 116                 in:    &pexState{},
 117                 arg:   0,
 118                 targM: &pp.PexMsg{},
 119                 targS: 0,
 120         },
 121         {
 122                 name: "add4",
 123                 in: &pexState{
 124                         ev: []pexEvent{
 125                                 pexEvent{pexAdd, addrs[0], f},
 126                                 pexEvent{pexAdd, addrs[1], f},
 127                                 pexEvent{pexAdd, addrs[2], f},
 128                                 pexEvent{pexAdd, addrs[3], f},
 129                         },
 130                 },
 131                 arg: 0,
 132                 targM: &pp.PexMsg{
 133                         Added: krpc.CompactIPv4NodeAddrs{
 134                                 nodeAddr(addrs[2]),
 135                                 nodeAddr(addrs[3]),
 136                         },
 137                         AddedFlags: []pp.PexPeerFlags{f, f},
 138                         Added6: krpc.CompactIPv6NodeAddrs{
 139                                 nodeAddr(addrs[0]),
 140                                 nodeAddr(addrs[1]),
 141                         },
 142                         Added6Flags: []pp.PexPeerFlags{f, f},
 143                 },
 144                 targS: 4,
 145         },
 146         {
 147                 name: "drop2",
 148                 arg:  0,
 149                 in: &pexState{
 150                         ev: []pexEvent{
 151                                 pexEvent{pexDrop, addrs[0], f},
 152                                 pexEvent{pexDrop, addrs[2], f},
 153                         },
 154                 },
 155                 targM: &pp.PexMsg{
 156                         Dropped: krpc.CompactIPv4NodeAddrs{
 157                                 nodeAddr(addrs[2]),
 158                         },
 159                         Dropped6: krpc.CompactIPv6NodeAddrs{
 160                                 nodeAddr(addrs[0]),
 161                         },
 162                 },
 163                 targS: 2,
 164         },
 165         {
 166                 name: "add2drop1",
 167                 arg:  0,
 168                 in: &pexState{
 169                         ev: []pexEvent{
 170                                 pexEvent{pexAdd, addrs[0], f},
 171                                 pexEvent{pexAdd, addrs[1], f},
 172                                 pexEvent{pexDrop, addrs[0], f},
 173                         },
 174                 },
 175                 targM: &pp.PexMsg{
 176                         Added6: krpc.CompactIPv6NodeAddrs{
 177                                 nodeAddr(addrs[1]),
 178                         },
 179                         Added6Flags: []pp.PexPeerFlags{f},
 180                 },
 181                 targS: 3,
 182         },
 183         {
 184                 name: "delayed",
 185                 arg:  0,
 186                 in: &pexState{
 187                         ev: []pexEvent{
 188                                 pexEvent{pexAdd, addrs[0], f},
 189                                 pexEvent{pexAdd, addrs[1], f},
 190                                 pexEvent{pexAdd, addrs[2], f},
 191                         },
 192                         hold: []pexEvent{
 193                                 pexEvent{pexDrop, addrs[0], f},
 194                                 pexEvent{pexDrop, addrs[2], f},
 195                                 pexEvent{pexDrop, addrs[1], f},
 196                         },
 197                 },
 198                 targM: &pp.PexMsg{
 199                         Added: krpc.CompactIPv4NodeAddrs{
 200                                 nodeAddr(addrs[2]),
 201                         },
 202                         AddedFlags: []pp.PexPeerFlags{f},
 203                         Added6: krpc.CompactIPv6NodeAddrs{
 204                                 nodeAddr(addrs[0]),
 205                                 nodeAddr(addrs[1]),
 206                         },
 207                         Added6Flags: []pp.PexPeerFlags{f, f},
 208                 },
 209                 targS: 3,
 210         },
 211         {
 212                 name: "followup",
 213                 arg:  1,
 214                 in: &pexState{
 215                         ev: []pexEvent{
 216                                 pexEvent{pexAdd, addrs[0], f},
 217                                 pexEvent{pexAdd, addrs[1], f},
 218                         },
 219                 },
 220                 targM: &pp.PexMsg{
 221                         Added6: krpc.CompactIPv6NodeAddrs{
 222                                 nodeAddr(addrs[1]),
 223                         },
 224                         Added6Flags: []pp.PexPeerFlags{f},
 225                 },
 226                 targS: 2,
 227         },
 228 }
 229
 230 func TestPexGenmsg(t *testing.T) {
 231         for _, tc := range testcases {
 232                 t.Run(tc.name, func(t *testing.T) {
 233                         s := tc.in
 234                         m, seen := s.Genmsg(tc.arg)
 235                         require.EqualValues(t, tc.targM, m)
 236                         require.EqualValues(t, tc.targS, seen)
 237                 })
 238         }
 239 }