ltep_test.go

   1 package torrent_test
   2
   3 import (
   4         "strconv"
   5         "testing"
   6
   7         "github.com/anacrolix/sync"
   8         qt "github.com/frankban/quicktest"
   9
  10         . "github.com/anacrolix/torrent"
  11         "github.com/anacrolix/torrent/internal/testutil"
  12         pp "github.com/anacrolix/torrent/peer_protocol"
  13 )
  14
  15 const (
  16         testRepliesToOddsExtensionName  = "pm_me_odds"
  17         testRepliesToEvensExtensionName = "pm_me_evens"
  18 )
  19
  20 func countHandler(
  21         c *qt.C,
  22         wg *sync.WaitGroup,
  23         // Name of the endpoint that this handler is for, for logging.
  24         handlerName string,
  25         // Whether we expect evens or odds
  26         expectedMod2 uint,
  27         // Extension name of messages we expect to handle.
  28         answerToName pp.ExtensionName,
  29         // Extension name of messages we expect to send.
  30         replyToName pp.ExtensionName,
  31         // Signal done when this value is seen.
  32         doneValue uint,
  33 ) func(event PeerConnReadExtensionMessageEvent) {
  34         return func(event PeerConnReadExtensionMessageEvent) {
  35                 // Read handshake, don't look it up.
  36                 if event.ExtensionNumber == 0 {
  37                         return
  38                 }
  39                 name, builtin, err := event.PeerConn.LocalLtepProtocolMap.LookupId(event.ExtensionNumber)
  40                 c.Assert(err, qt.IsNil)
  41                 // Not a user protocol.
  42                 if builtin {
  43                         return
  44                 }
  45                 switch name {
  46                 case answerToName:
  47                         u64, err := strconv.ParseUint(string(event.Payload), 10, 0)
  48                         c.Assert(err, qt.IsNil)
  49                         i := uint(u64)
  50                         c.Logf("%v got %d", handlerName, i)
  51                         if i == doneValue {
  52                                 wg.Done()
  53                                 return
  54                         }
  55                         c.Assert(i%2, qt.Equals, expectedMod2)
  56                         go func() {
  57                                 c.Assert(
  58                                         event.PeerConn.WriteExtendedMessage(
  59                                                 replyToName,
  60                                                 []byte(strconv.FormatUint(uint64(i+1), 10))),
  61                                         qt.IsNil)
  62                         }()
  63                 default:
  64                         c.Fatalf("got unexpected extension name %q", name)
  65                 }
  66         }
  67 }
  68
  69 func TestUserLtep(t *testing.T) {
  70         c := qt.New(t)
  71         var wg sync.WaitGroup
  72
  73         makeCfg := func() *ClientConfig {
  74                 cfg := TestingConfig(t)
  75                 // Only want a single connection to between the clients.
  76                 cfg.DisableUTP = true
  77                 cfg.DisableIPv6 = true
  78                 return cfg
  79         }
  80
  81         evensCfg := makeCfg()
  82         evensCfg.Callbacks.ReadExtendedHandshake = func(pc *PeerConn, msg *pp.ExtendedHandshakeMessage) {
  83                 // The client lock is held while handling this event, so we have to do synchronous work in a
  84                 // separate goroutine.
  85                 go func() {
  86                         // Check sending an extended message for a protocol the peer doesn't support is an error.
  87                         c.Check(pc.WriteExtendedMessage("pm_me_floats", []byte("3.142")), qt.IsNotNil)
  88                         // Kick things off by sending a 1.
  89                         c.Check(pc.WriteExtendedMessage(testRepliesToOddsExtensionName, []byte("1")), qt.IsNil)
  90                 }()
  91         }
  92         evensCfg.Callbacks.PeerConnReadExtensionMessage = append(
  93                 evensCfg.Callbacks.PeerConnReadExtensionMessage,
  94                 countHandler(c, &wg, "evens", 0, testRepliesToEvensExtensionName, testRepliesToOddsExtensionName, 100))
  95         evensCfg.Callbacks.PeerConnAdded = append(evensCfg.Callbacks.PeerConnAdded, func(conn *PeerConn) {
  96                 conn.LocalLtepProtocolMap.AddUserProtocol(testRepliesToEvensExtensionName)
  97                 c.Assert(conn.LocalLtepProtocolMap.Index[conn.LocalLtepProtocolMap.NumBuiltin:], qt.HasLen, 1)
  98         })
  99
 100         oddsCfg := makeCfg()
 101         oddsCfg.Callbacks.PeerConnAdded = append(oddsCfg.Callbacks.PeerConnAdded, func(conn *PeerConn) {
 102                 conn.LocalLtepProtocolMap.AddUserProtocol(testRepliesToOddsExtensionName)
 103                 c.Assert(conn.LocalLtepProtocolMap.Index[conn.LocalLtepProtocolMap.NumBuiltin:], qt.HasLen, 1)
 104         })
 105         oddsCfg.Callbacks.PeerConnReadExtensionMessage = append(
 106                 oddsCfg.Callbacks.PeerConnReadExtensionMessage,
 107                 countHandler(c, &wg, "odds", 1, testRepliesToOddsExtensionName, testRepliesToEvensExtensionName, 100))
 108
 109         cl1, err := NewClient(oddsCfg)
 110         c.Assert(err, qt.IsNil)
 111         defer cl1.Close()
 112         cl2, err := NewClient(evensCfg)
 113         c.Assert(err, qt.IsNil)
 114         defer cl2.Close()
 115         addOpts := AddTorrentOpts{}
 116         t1, _ := cl1.AddTorrentOpt(addOpts)
 117         t2, _ := cl2.AddTorrentOpt(addOpts)
 118         defer testutil.ExportStatusWriter(cl1, "cl1", t)()
 119         defer testutil.ExportStatusWriter(cl2, "cl2", t)()
 120         // Expect one PeerConn to see the value.
 121         wg.Add(1)
 122         added := t1.AddClientPeer(cl2)
 123         // Ensure some addresses for the other client were added.
 124         c.Assert(added, qt.Not(qt.Equals), 0)
 125         wg.Wait()
 126         _ = t2
 127 }