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