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, newConnectionOpts{
  21                 remoteAddr: addr,
  22                 network:    addr.Network(),
  23         })
  24         c.PeerExtensionIDs = make(map[pp.ExtensionName]pp.ExtensionNumber)
  25         c.PeerExtensionIDs[pp.ExtensionNamePex] = pexExtendedId
  26         c.messageWriter.mu.Lock()
  27         c.setTorrent(torrent)
  28         if err := torrent.addPeerConn(c); err != nil {
  29                 t.Log(err)
  30         }
  31
  32         connWriteCond := c.messageWriter.writeCond.Signaled()
  33         c.pex.Init(c)
  34         require.True(t, c.pex.IsEnabled(), "should get enabled")
  35         defer c.pex.Close()
  36
  37         var out pp.Message
  38         writerCalled := false
  39         testWriter := func(m pp.Message) bool {
  40                 writerCalled = true
  41                 out = m
  42                 return true
  43         }
  44         <-connWriteCond
  45         c.pex.Share(testWriter)
  46         require.True(t, writerCalled)
  47         require.EqualValues(t, pp.Extended, out.Type)
  48         require.EqualValues(t, pexExtendedId, 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 }