pexconn_test.go

   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         cl.initLogger()
  17         torrent := cl.newTorrentForTesting()
  18         addr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4747}
  19         c := cl.newConnection(nil, newConnectionOpts{
  20                 remoteAddr: addr,
  21                 network:    addr.Network(),
  22         })
  23         c.PeerExtensionIDs = make(map[pp.ExtensionName]pp.ExtensionNumber)
  24         c.PeerExtensionIDs[pp.ExtensionNamePex] = 1
  25         c.messageWriter.mu.Lock()
  26         c.setTorrent(torrent)
  27         if err := torrent.addPeerConn(c); err != nil {
  28                 t.Log(err)
  29         }
  30
  31         connWriteCond := c.messageWriter.writeCond.Signaled()
  32         c.pex.Init(c)
  33         require.True(t, c.pex.IsEnabled(), "should get enabled")
  34         defer c.pex.Close()
  35
  36         var out pp.Message
  37         writerCalled := false
  38         testWriter := func(m pp.Message) bool {
  39                 writerCalled = true
  40                 out = m
  41                 return true
  42         }
  43         <-connWriteCond
  44         c.pex.Share(testWriter)
  45         require.True(t, writerCalled)
  46         require.EqualValues(t, pp.Extended, out.Type)
  47         require.NotEqualValues(t, out.ExtendedID, 0)
  48         require.EqualValues(t, c.PeerExtensionIDs[pp.ExtensionNamePex], out.ExtendedID)
  49
  50         x, err := pp.LoadPexMsg(out.ExtendedPayload)
  51         require.NoError(t, err)
  52         targx := pp.PexMsg{
  53                 Added:      krpc.CompactIPv4NodeAddrs(nil),
  54                 AddedFlags: []pp.PexPeerFlags{},
  55                 Added6: krpc.CompactIPv6NodeAddrs{
  56                         krpcNodeAddrFromNetAddr(addr),
  57                 },
  58                 Added6Flags: []pp.PexPeerFlags{0},
  59         }
  60         require.EqualValues(t, targx, x)
  61 }