7 "github.com/anacrolix/dht/v2/krpc"
8 "github.com/stretchr/testify/assert"
9 "github.com/stretchr/testify/require"
11 pp "github.com/anacrolix/torrent/peer_protocol"
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},
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},
35 func TestPexReset(t *testing.T) {
38 {Peer: Peer{RemoteAddr: addrs[0]}},
39 {Peer: Peer{RemoteAddr: addrs[1]}},
40 {Peer: Peer{RemoteAddr: addrs[2]}},
47 require.EqualValues(t, targ, s)
50 func krpcNodeAddrFromNetAddr(addr net.Addr) krpc.NodeAddr {
51 addrPort, err := addrPortFromPeerRemoteAddr(addr)
55 return krpcNodeAddrFromAddrPort(addrPort)
58 var testcases = []struct {
62 update func(*pexState)
72 in: func() *pexState {
74 nullAddr := &net.TCPAddr{}
75 s.Add(&PeerConn{Peer: Peer{RemoteAddr: nullAddr}})
82 in: func() *pexState {
83 nullAddr := &net.TCPAddr{}
85 s.Drop(&PeerConn{Peer: Peer{RemoteAddr: nullAddr}, pex: pexConnState{Listed: true}})
92 in: func() *pexState {
94 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}})
95 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1], outgoing: true}})
96 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[2], outgoing: true}})
97 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[3]}})
101 Added: krpc.CompactIPv4NodeAddrs{
102 krpcNodeAddrFromNetAddr(addrs[2]),
103 krpcNodeAddrFromNetAddr(addrs[3]),
105 AddedFlags: []pp.PexPeerFlags{pp.PexOutgoingConn, 0},
106 Added6: krpc.CompactIPv6NodeAddrs{
107 krpcNodeAddrFromNetAddr(addrs[0]),
108 krpcNodeAddrFromNetAddr(addrs[1]),
110 Added6Flags: []pp.PexPeerFlags{0, pp.PexOutgoingConn},
115 in: func() *pexState {
116 s := &pexState{nc: pexTargAdded + 2}
117 s.Drop(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
118 s.Drop(&PeerConn{Peer: Peer{RemoteAddr: addrs[2]}, pex: pexConnState{Listed: true}})
122 Dropped: krpc.CompactIPv4NodeAddrs{
123 krpcNodeAddrFromNetAddr(addrs[2]),
125 Dropped6: krpc.CompactIPv6NodeAddrs{
126 krpcNodeAddrFromNetAddr(addrs[0]),
132 in: func() *pexState {
134 {Peer: Peer{RemoteAddr: addrs[0]}},
135 {Peer: Peer{RemoteAddr: addrs[1]}},
136 {Peer: Peer{RemoteAddr: addrs[2]}},
138 s := &pexState{nc: pexTargAdded}
142 s.Drop(&conns[2]) // to be ignored: it wasn't added
146 Added6: krpc.CompactIPv6NodeAddrs{
147 krpcNodeAddrFromNetAddr(addrs[1]),
149 Added6Flags: []pp.PexPeerFlags{0},
154 in: func() *pexState {
156 {Peer: Peer{RemoteAddr: addrs[0]}},
157 {Peer: Peer{RemoteAddr: addrs[1]}},
158 {Peer: Peer{RemoteAddr: addrs[2]}},
164 s.Drop(&conns[0]) // on hold: s.nc < pexTargAdded
170 Added: krpc.CompactIPv4NodeAddrs{
171 krpcNodeAddrFromNetAddr(addrs[2]),
173 AddedFlags: []pp.PexPeerFlags{0},
174 Added6: krpc.CompactIPv6NodeAddrs{
175 krpcNodeAddrFromNetAddr(addrs[0]),
176 krpcNodeAddrFromNetAddr(addrs[1]),
178 Added6Flags: []pp.PexPeerFlags{0, 0},
183 in: func() *pexState {
185 {Peer: Peer{RemoteAddr: addrs[0]}},
186 {Peer: Peer{RemoteAddr: addrs[1]}},
188 s := &pexState{nc: pexTargAdded - 1}
190 s.Drop(&conns[0]) // on hold: s.nc < pexTargAdded
191 s.Add(&conns[1]) // unholds the above
195 Added6: krpc.CompactIPv6NodeAddrs{
196 krpcNodeAddrFromNetAddr(addrs[1]),
198 Added6Flags: []pp.PexPeerFlags{0},
203 in: func() *pexState {
205 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}})
209 Added6: krpc.CompactIPv6NodeAddrs{
210 krpcNodeAddrFromNetAddr(addrs[0]),
212 Added6Flags: []pp.PexPeerFlags{0},
214 update: func(s *pexState) {
215 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1]}})
218 Added6: krpc.CompactIPv6NodeAddrs{
219 krpcNodeAddrFromNetAddr(addrs[1]),
221 Added6Flags: []pp.PexPeerFlags{0},
226 // Represents the contents of a PexMsg in a way that supports equivalence checking in tests. This is
227 // necessary because pexMsgFactory uses maps and so ordering of the resultant PexMsg isn't
228 // deterministic. Because the flags are in a different array, we can't just use testify's
229 // ElementsMatch because the ordering *does* still matter between an added addr and its flags.
230 type comparablePexMsg struct {
231 added, added6 []krpc.NodeAddr
232 addedFlags, added6Flags []pp.PexPeerFlags
233 dropped, dropped6 []krpc.NodeAddr
236 // Such Rust-inspired.
237 func (me *comparablePexMsg) From(f pp.PexMsg) {
239 me.addedFlags = f.AddedFlags
241 me.added6Flags = f.Added6Flags
242 me.dropped = f.Dropped
243 me.dropped6 = f.Dropped6
246 // For PexMsg created by pexMsgFactory, this is as good as it can get without using data structures
247 // in pexMsgFactory that preserve insert ordering.
248 func (actual comparablePexMsg) AssertEqual(t *testing.T, expected comparablePexMsg) {
249 assert.ElementsMatch(t, expected.added, actual.added)
250 assert.ElementsMatch(t, expected.addedFlags, actual.addedFlags)
251 assert.ElementsMatch(t, expected.added6, actual.added6)
252 assert.ElementsMatch(t, expected.added6Flags, actual.added6Flags)
253 assert.ElementsMatch(t, expected.dropped, actual.dropped)
254 assert.ElementsMatch(t, expected.dropped6, actual.dropped6)
257 func assertPexMsgsEqual(t *testing.T, expected, actual pp.PexMsg) {
258 var ec, ac comparablePexMsg
261 ac.AssertEqual(t, ec)
264 func TestPexGenmsg0(t *testing.T) {
265 for _, tc := range testcases {
266 t.Run(tc.name, func(t *testing.T) {
268 m, last := s.Genmsg(nil)
269 assertPexMsgsEqual(t, tc.targ, m)
270 if tc.update != nil {
272 m1, last := s.Genmsg(last)
273 assertPexMsgsEqual(t, tc.targ1, m1)
274 assert.NotNil(t, last)
280 // generate 𝑛 distinct values of net.Addr
281 func addrgen(n int) chan net.Addr {
282 c := make(chan net.Addr)
285 for i := 4747; i < 65535 && n > 0; i++ {
286 c <- &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: i}
293 func TestPexInitialNoCutoff(t *testing.T) {
294 const n = 2 * pexMaxDelta
298 for addr := range c {
299 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addr}})
301 m, _ := s.Genmsg(nil)
303 require.EqualValues(t, n, len(m.Added))
304 require.EqualValues(t, n, len(m.AddedFlags))
305 require.EqualValues(t, 0, len(m.Added6))
306 require.EqualValues(t, 0, len(m.Added6Flags))
307 require.EqualValues(t, 0, len(m.Dropped))
308 require.EqualValues(t, 0, len(m.Dropped6))
311 func benchmarkPexInitialN(b *testing.B, npeers int) {
312 for i := 0; i < b.N; i++ {
315 for addr := range c {
316 s.Add(&PeerConn{Peer: Peer{RemoteAddr: addr}})
322 // obtain at least 5 points, e.g. to plot a graph
323 func BenchmarkPexInitial4(b *testing.B) { benchmarkPexInitialN(b, 4) }
324 func BenchmarkPexInitial50(b *testing.B) { benchmarkPexInitialN(b, 50) }
325 func BenchmarkPexInitial100(b *testing.B) { benchmarkPexInitialN(b, 100) }
326 func BenchmarkPexInitial200(b *testing.B) { benchmarkPexInitialN(b, 200) }
327 func BenchmarkPexInitial400(b *testing.B) { benchmarkPexInitialN(b, 400) }