github.com/anacrolix/torrent@v1.61.0/pex_test.go

package torrent

import (
	"net"
	"testing"

	"github.com/anacrolix/dht/v2/krpc"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	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 krpcNodeAddrFromNetAddr(addr net.Addr) krpc.NodeAddr {
	addrPort, err := addrPortFromPeerRemoteAddr(addr)
	if err != nil {
		panic(err)
	}
	return krpcNodeAddrFromAddrPort(addrPort)
}

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{
				krpcNodeAddrFromNetAddr(addrs[2]),
				krpcNodeAddrFromNetAddr(addrs[3]),
			},
			AddedFlags: []pp.PexPeerFlags{pp.PexOutgoingConn, 0},
			Added6: krpc.CompactIPv6NodeAddrs{
				krpcNodeAddrFromNetAddr(addrs[0]),
				krpcNodeAddrFromNetAddr(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{
				krpcNodeAddrFromNetAddr(addrs[2]),
			},
			Dropped6: krpc.CompactIPv6NodeAddrs{
				krpcNodeAddrFromNetAddr(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{
				krpcNodeAddrFromNetAddr(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{
				krpcNodeAddrFromNetAddr(addrs[2]),
			},
			AddedFlags: []pp.PexPeerFlags{0},
			Added6: krpc.CompactIPv6NodeAddrs{
				krpcNodeAddrFromNetAddr(addrs[0]),
				krpcNodeAddrFromNetAddr(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{
				krpcNodeAddrFromNetAddr(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{
				krpcNodeAddrFromNetAddr(addrs[0]),
			},
			Added6Flags: []pp.PexPeerFlags{0},
		},
		update: func(s *pexState) {
			s.Add(&PeerConn{Peer: Peer{RemoteAddr: addrs[1]}})
		},
		targ1: pp.PexMsg{
			Added6: krpc.CompactIPv6NodeAddrs{
				krpcNodeAddrFromNetAddr(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) }