pexconn_test.go

package torrent

import (
	"net"
	"testing"

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

	pp "github.com/anacrolix/torrent/peer_protocol"
)

func TestPexConnState(t *testing.T) {
	var cl Client
	cl.init(TestingConfig(t))
	cl.initLogger()
	torrent := cl.newTorrentForTesting()
	addr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4747}
	c := cl.newConnection(nil, newConnectionOpts{
		remoteAddr: addr,
		network:    addr.Network(),
	})
	c.PeerExtensionIDs = make(map[pp.ExtensionName]pp.ExtensionNumber)
	c.PeerExtensionIDs[pp.ExtensionNamePex] = 1
	c.messageWriter.mu.Lock()
	c.setTorrent(torrent)
	if err := torrent.addPeerConn(c); err != nil {
		t.Log(err)
	}

	connWriteCond := c.messageWriter.writeCond.Signaled()
	c.pex.Init(c)
	require.True(t, c.pex.IsEnabled(), "should get enabled")
	defer c.pex.Close()

	var out pp.Message
	writerCalled := false
	testWriter := func(m pp.Message) bool {
		writerCalled = true
		out = m
		return true
	}
	<-connWriteCond
	c.pex.Share(testWriter)
	require.True(t, writerCalled)
	require.EqualValues(t, pp.Extended, out.Type)
	require.NotEqualValues(t, out.ExtendedID, 0)
	require.EqualValues(t, c.PeerExtensionIDs[pp.ExtensionNamePex], out.ExtendedID)

	x, err := pp.LoadPexMsg(out.ExtendedPayload)
	require.NoError(t, err)
	targx := pp.PexMsg{
		Added:      krpc.CompactIPv4NodeAddrs(nil),
		AddedFlags: []pp.PexPeerFlags{},
		Added6: krpc.CompactIPv6NodeAddrs{
			krpcNodeAddrFromNetAddr(addr),
		},
		Added6Flags: []pp.PexPeerFlags{0},
	}
	require.EqualValues(t, targx, x)
}