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         "github.com/anacrolix/torrent/metainfo"
  11         pp "github.com/anacrolix/torrent/peer_protocol"
  12 )
  13
  14 func TestPexConnState(t *testing.T) {
  15         cl := Client{
  16                 config: TestingConfig(t),
  17         }
  18         cl.initLogger()
  19         torrent := cl.newTorrent(metainfo.Hash{}, nil)
  20         addr := &net.TCPAddr{IP: net.IPv6loopback, Port: 4747}
  21         c := cl.newConnection(nil, false, addr, addr.Network(), "")
  22         c.PeerExtensionIDs = make(map[pp.ExtensionName]pp.ExtensionNumber)
  23         c.PeerExtensionIDs[pp.ExtensionNamePex] = pexExtendedId
  24         c.messageWriter.mu.Lock()
  25         c.setTorrent(torrent)
  26         torrent.addPeerConn(c)
  27
  28         c.pex.Init(c)
  29         require.True(t, c.pex.IsEnabled(), "should get enabled")
  30         defer c.pex.Close()
  31
  32         var out pp.Message
  33         writerCalled := false
  34         testWriter := func(m pp.Message) bool {
  35                 writerCalled = true
  36                 out = m
  37                 return true
  38         }
  39         <-c.messageWriter.writeCond.Signaled()
  40         c.pex.Share(testWriter)
  41         require.True(t, writerCalled)
  42         require.EqualValues(t, pp.Extended, out.Type)
  43         require.EqualValues(t, pexExtendedId, out.ExtendedID)
  44
  45         x, err := pp.LoadPexMsg(out.ExtendedPayload)
  46         require.NoError(t, err)
  47         targx := pp.PexMsg{
  48                 Added:      krpc.CompactIPv4NodeAddrs(nil),
  49                 AddedFlags: []pp.PexPeerFlags{},
  50                 Added6: krpc.CompactIPv6NodeAddrs{
  51                         mustNodeAddr(addr),
  52                 },
  53                 Added6Flags: []pp.PexPeerFlags{0},
  54         }
  55         require.EqualValues(t, targx, x)
  56 }