github.com/anacrolix/torrent@v1.61.0/pexconn_test.go (about)

     1  package torrent
     2  
     3  import (
     4  	"net"
     5  	"testing"
     6  
     7  	"github.com/anacrolix/dht/v2/krpc"
     8  	"github.com/stretchr/testify/require"
     9  
    10  	pp "github.com/anacrolix/torrent/peer_protocol"
    11  )
    12  
    13  func TestPexConnState(t *testing.T) {
    14  	var cl Client
    15  	cl.init(TestingConfig(t))
    16  	torrent := cl.newTorrentForTesting()
    17  	addr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4747}
    18  	c := cl.newConnection(nil, newConnectionOpts{
    19  		remoteAddr: addr,
    20  		network:    addr.Network(),
    21  	})
    22  	c.PeerExtensionIDs = make(map[pp.ExtensionName]pp.ExtensionNumber)
    23  	c.PeerExtensionIDs[pp.ExtensionNamePex] = 1
    24  	c.messageWriter.mu.Lock()
    25  	c.setTorrent(torrent)
    26  	if err := torrent.addPeerConn(c); err != nil {
    27  		t.Log(err)
    28  	}
    29  
    30  	connWriteCond := c.messageWriter.writeCond.Signaled()
    31  	c.pex.Init(c)
    32  	require.True(t, c.pex.IsEnabled(), "should get enabled")
    33  	defer c.pex.Close()
    34  
    35  	var out pp.Message
    36  	writerCalled := false
    37  	testWriter := func(m pp.Message) bool {
    38  		writerCalled = true
    39  		out = m
    40  		return true
    41  	}
    42  	<-connWriteCond
    43  	c.pex.Share(testWriter)
    44  	require.True(t, writerCalled)
    45  	require.EqualValues(t, pp.Extended, out.Type)
    46  	require.NotEqualValues(t, out.ExtendedID, 0)
    47  	require.EqualValues(t, c.PeerExtensionIDs[pp.ExtensionNamePex], out.ExtendedID)
    48  
    49  	x, err := pp.LoadPexMsg(out.ExtendedPayload)
    50  	require.NoError(t, err)
    51  	targx := pp.PexMsg{
    52  		Added:      krpc.CompactIPv4NodeAddrs(nil),
    53  		AddedFlags: []pp.PexPeerFlags{},
    54  		Added6: krpc.CompactIPv6NodeAddrs{
    55  			krpcNodeAddrFromNetAddr(addr),
    56  		},
    57  		Added6Flags: []pp.PexPeerFlags{0},
    58  	}
    59  	require.EqualValues(t, targx, x)
    60  }