328 lines
8.2 KiB
Go
328 lines
8.2 KiB
Go
package torrent
|
||
|
||
import (
|
||
"net"
|
||
"testing"
|
||
|
||
"github.com/stretchr/testify/assert"
|
||
"github.com/stretchr/testify/require"
|
||
|
||
"github.com/anacrolix/dht/v2/krpc"
|
||
pp "github.com/anacrolix/torrent/peer_protocol"
|
||
)
|
||
|
||
var (
|
||
addrs6 = []net.Addr{
|
||
&net.TCPAddr{IP: net.IPv6loopback, Port: 4747},
|
||
&net.TCPAddr{IP: net.IPv6loopback, Port: 4748},
|
||
&net.TCPAddr{IP: net.IPv6loopback, Port: 4749},
|
||
&net.TCPAddr{IP: net.IPv6loopback, Port: 4750},
|
||
}
|
||
addrs4 = []net.Addr{
|
||
&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4747},
|
||
&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4748},
|
||
&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4749},
|
||
&net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 4750},
|
||
}
|
||
addrs = []net.Addr{
|
||
addrs6[0],
|
||
addrs6[1],
|
||
addrs4[0],
|
||
addrs4[1],
|
||
}
|
||
)
|
||
|
||
func TestPexReset(t *testing.T) {
|
||
s := &pexState{}
|
||
conns := []PeerConn{
|
||
{Peer: Peer{RemoteAddr: addrs[0]}},
|
||
{Peer: Peer{RemoteAddr: addrs[1]}},
|
||
{Peer: Peer{RemoteAddr: addrs[2]}},
|
||
}
|
||
s.Add(&conns[0])
|
||
s.Add(&conns[1])
|
||
s.Drop(&conns[0])
|
||
s.Reset()
|
||
targ := new(pexState)
|
||
require.EqualValues(t, targ, s)
|
||
}
|
||
|
||
func mustNodeAddr(addr net.Addr) krpc.NodeAddr {
|
||
ret, ok := nodeAddr(addr)
|
||
if !ok {
|
||
panic(addr)
|
||
}
|
||
return ret
|
||
}
|
||
|
||
var testcases = []struct {
|
||
name string
|
||
in *pexState
|
||
targ pp.PexMsg
|
||
update func(*pexState)
|
||
targ1 pp.PexMsg
|
||
}{
|
||
{
|
||
name: "empty",
|
||
in: &pexState{},
|
||
targ: pp.PexMsg{},
|
||
},
|
||
{
|
||
name: "add0",
|
||
in: func() *pexState {
|
||
s := new(pexState)
|
||
nullAddr := &net.TCPAddr{}
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: nullAddr}})
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{},
|
||
},
|
||
{
|
||
name: "drop0",
|
||
in: func() *pexState {
|
||
nullAddr := &net.TCPAddr{}
|
||
s := new(pexState)
|
||
s.Drop(&PeerConn{Peer: Peer{RemoteAddr: nullAddr}, pex: pexConnState{Listed: true}})
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{},
|
||
},
|
||
{
|
||
name: "add4",
|
||
in: func() *pexState {
|
||
s := new(pexState)
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}})
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1], outgoing: true}})
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[2], outgoing: true}})
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[3]}})
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{
|
||
Added: krpc.CompactIPv4NodeAddrs{
|
||
mustNodeAddr(addrs[2]),
|
||
mustNodeAddr(addrs[3]),
|
||
},
|
||
AddedFlags: []pp.PexPeerFlags{pp.PexOutgoingConn, 0},
|
||
Added6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[0]),
|
||
mustNodeAddr(addrs[1]),
|
||
},
|
||
Added6Flags: []pp.PexPeerFlags{0, pp.PexOutgoingConn},
|
||
},
|
||
},
|
||
{
|
||
name: "drop2",
|
||
in: func() *pexState {
|
||
s := &pexState{nc: pexTargAdded + 2}
|
||
s.Drop(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}, pex: pexConnState{Listed: true}})
|
||
s.Drop(&PeerConn{Peer: Peer{RemoteAddr: addrs[2]}, pex: pexConnState{Listed: true}})
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{
|
||
Dropped: krpc.CompactIPv4NodeAddrs{
|
||
mustNodeAddr(addrs[2]),
|
||
},
|
||
Dropped6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[0]),
|
||
},
|
||
},
|
||
},
|
||
{
|
||
name: "add2drop1",
|
||
in: func() *pexState {
|
||
conns := []PeerConn{
|
||
{Peer: Peer{RemoteAddr: addrs[0]}},
|
||
{Peer: Peer{RemoteAddr: addrs[1]}},
|
||
{Peer: Peer{RemoteAddr: addrs[2]}},
|
||
}
|
||
s := &pexState{nc: pexTargAdded}
|
||
s.Add(&conns[0])
|
||
s.Add(&conns[1])
|
||
s.Drop(&conns[0])
|
||
s.Drop(&conns[2]) // to be ignored: it wasn't added
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{
|
||
Added6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[1]),
|
||
},
|
||
Added6Flags: []pp.PexPeerFlags{0},
|
||
},
|
||
},
|
||
{
|
||
name: "delayed",
|
||
in: func() *pexState {
|
||
conns := []PeerConn{
|
||
{Peer: Peer{RemoteAddr: addrs[0]}},
|
||
{Peer: Peer{RemoteAddr: addrs[1]}},
|
||
{Peer: Peer{RemoteAddr: addrs[2]}},
|
||
}
|
||
s := new(pexState)
|
||
s.Add(&conns[0])
|
||
s.Add(&conns[1])
|
||
s.Add(&conns[2])
|
||
s.Drop(&conns[0]) // on hold: s.nc < pexTargAdded
|
||
s.Drop(&conns[2])
|
||
s.Drop(&conns[1])
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{
|
||
Added: krpc.CompactIPv4NodeAddrs{
|
||
mustNodeAddr(addrs[2]),
|
||
},
|
||
AddedFlags: []pp.PexPeerFlags{0},
|
||
Added6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[0]),
|
||
mustNodeAddr(addrs[1]),
|
||
},
|
||
Added6Flags: []pp.PexPeerFlags{0, 0},
|
||
},
|
||
},
|
||
{
|
||
name: "unheld",
|
||
in: func() *pexState {
|
||
conns := []PeerConn{
|
||
{Peer: Peer{RemoteAddr: addrs[0]}},
|
||
{Peer: Peer{RemoteAddr: addrs[1]}},
|
||
}
|
||
s := &pexState{nc: pexTargAdded - 1}
|
||
s.Add(&conns[0])
|
||
s.Drop(&conns[0]) // on hold: s.nc < pexTargAdded
|
||
s.Add(&conns[1]) // unholds the above
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{
|
||
Added6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[1]),
|
||
},
|
||
Added6Flags: []pp.PexPeerFlags{0},
|
||
},
|
||
},
|
||
{
|
||
name: "followup",
|
||
in: func() *pexState {
|
||
s := new(pexState)
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[0]}})
|
||
return s
|
||
}(),
|
||
targ: pp.PexMsg{
|
||
Added6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[0]),
|
||
},
|
||
Added6Flags: []pp.PexPeerFlags{0},
|
||
},
|
||
update: func(s *pexState) {
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1]}})
|
||
},
|
||
targ1: pp.PexMsg{
|
||
Added6: krpc.CompactIPv6NodeAddrs{
|
||
mustNodeAddr(addrs[1]),
|
||
},
|
||
Added6Flags: []pp.PexPeerFlags{0},
|
||
},
|
||
},
|
||
}
|
||
|
||
// Represents the contents of a PexMsg in a way that supports equivalence checking in tests. This is
|
||
// necessary because pexMsgFactory uses maps and so ordering of the resultant PexMsg isn't
|
||
// deterministic. Because the flags are in a different array, we can't just use testify's
|
||
// ElementsMatch because the ordering *does* still matter between an added addr and its flags.
|
||
type comparablePexMsg struct {
|
||
added, added6 []krpc.NodeAddr
|
||
addedFlags, added6Flags []pp.PexPeerFlags
|
||
dropped, dropped6 []krpc.NodeAddr
|
||
}
|
||
|
||
// Such Rust-inspired.
|
||
func (me *comparablePexMsg) From(f pp.PexMsg) {
|
||
me.added = f.Added
|
||
me.addedFlags = f.AddedFlags
|
||
me.added6 = f.Added6
|
||
me.added6Flags = f.Added6Flags
|
||
me.dropped = f.Dropped
|
||
me.dropped6 = f.Dropped6
|
||
}
|
||
|
||
// For PexMsg created by pexMsgFactory, this is as good as it can get without using data structures
|
||
// in pexMsgFactory that preserve insert ordering.
|
||
func (actual comparablePexMsg) AssertEqual(t *testing.T, expected comparablePexMsg) {
|
||
assert.ElementsMatch(t, expected.added, actual.added)
|
||
assert.ElementsMatch(t, expected.addedFlags, actual.addedFlags)
|
||
assert.ElementsMatch(t, expected.added6, actual.added6)
|
||
assert.ElementsMatch(t, expected.added6Flags, actual.added6Flags)
|
||
assert.ElementsMatch(t, expected.dropped, actual.dropped)
|
||
assert.ElementsMatch(t, expected.dropped6, actual.dropped6)
|
||
}
|
||
|
||
func assertPexMsgsEqual(t *testing.T, expected, actual pp.PexMsg) {
|
||
var ec, ac comparablePexMsg
|
||
ec.From(expected)
|
||
ac.From(actual)
|
||
ac.AssertEqual(t, ec)
|
||
}
|
||
|
||
func TestPexGenmsg0(t *testing.T) {
|
||
for _, tc := range testcases {
|
||
t.Run(tc.name, func(t *testing.T) {
|
||
s := *tc.in
|
||
m, last := s.Genmsg(nil)
|
||
assertPexMsgsEqual(t, tc.targ, m)
|
||
if tc.update != nil {
|
||
tc.update(&s)
|
||
m1, last := s.Genmsg(last)
|
||
assertPexMsgsEqual(t, tc.targ1, m1)
|
||
assert.NotNil(t, last)
|
||
}
|
||
})
|
||
}
|
||
}
|
||
|
||
// generate 𝑛 distinct values of net.Addr
|
||
func addrgen(n int) chan net.Addr {
|
||
c := make(chan net.Addr)
|
||
go func() {
|
||
defer close(c)
|
||
for i := 4747; i < 65535 && n > 0; i++ {
|
||
c <- &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: i}
|
||
n--
|
||
}
|
||
}()
|
||
return c
|
||
}
|
||
|
||
func TestPexInitialNoCutoff(t *testing.T) {
|
||
const n = 2 * pexMaxDelta
|
||
var s pexState
|
||
|
||
c := addrgen(n)
|
||
for addr := range c {
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addr}})
|
||
}
|
||
m, _ := s.Genmsg(nil)
|
||
|
||
require.EqualValues(t, n, len(m.Added))
|
||
require.EqualValues(t, n, len(m.AddedFlags))
|
||
require.EqualValues(t, 0, len(m.Added6))
|
||
require.EqualValues(t, 0, len(m.Added6Flags))
|
||
require.EqualValues(t, 0, len(m.Dropped))
|
||
require.EqualValues(t, 0, len(m.Dropped6))
|
||
}
|
||
|
||
func benchmarkPexInitialN(b *testing.B, npeers int) {
|
||
for i := 0; i < b.N; i++ {
|
||
var s pexState
|
||
c := addrgen(npeers)
|
||
for addr := range c {
|
||
s.Add(&PeerConn{Peer: Peer{RemoteAddr: addr}})
|
||
s.Genmsg(nil)
|
||
}
|
||
}
|
||
}
|
||
|
||
// obtain at least 5 points, e.g. to plot a graph
|
||
func BenchmarkPexInitial4(b *testing.B) { benchmarkPexInitialN(b, 4) }
|
||
func BenchmarkPexInitial50(b *testing.B) { benchmarkPexInitialN(b, 50) }
|
||
func BenchmarkPexInitial100(b *testing.B) { benchmarkPexInitialN(b, 100) }
|
||
func BenchmarkPexInitial200(b *testing.B) { benchmarkPexInitialN(b, 200) }
|
||
func BenchmarkPexInitial400(b *testing.B) { benchmarkPexInitialN(b, 400) }
|