ltep.go

   1 package torrent
   2
   3 import (
   4         "fmt"
   5         "slices"
   6
   7         g "github.com/anacrolix/generics"
   8
   9         pp "github.com/anacrolix/torrent/peer_protocol"
  10 )
  11
  12 type LocalLtepProtocolMap struct {
  13         // 1-based mapping from extension number to extension name (subtract one from the extension ID
  14         // to find the corresponding protocol name). The first LocalLtepProtocolBuiltinCount of these
  15         // are use builtin handlers. If you want to handle builtin protocols yourself, you would move
  16         // them above the threshold. You can disable them by removing them entirely, and add your own.
  17         // These changes should be done in the PeerConnAdded callback.
  18         Index []pp.ExtensionName
  19         // How many of the protocols are using the builtin handlers.
  20         NumBuiltin int
  21 }
  22
  23 func (me *LocalLtepProtocolMap) toSupportedExtensionDict() (m map[pp.ExtensionName]pp.ExtensionNumber) {
  24         g.MakeMapWithCap(&m, len(me.Index))
  25         for i, name := range me.Index {
  26                 old := g.MapInsert(m, name, pp.ExtensionNumber(i+1))
  27                 if old.Ok {
  28                         panic(fmt.Sprintf("extension %q already defined with id %v", name, old.Value))
  29                 }
  30         }
  31         return
  32 }
  33
  34 // Returns the local extension name for the given ID. If builtin is true, the implementation intends
  35 // to handle it itself. For incoming messages with extension ID 0, the message is a handshake, and
  36 // should be treated specially.
  37 func (me *LocalLtepProtocolMap) LookupId(id pp.ExtensionNumber) (name pp.ExtensionName, builtin bool, err error) {
  38         if id == 0 {
  39                 err = fmt.Errorf("extension ID 0 is handshake")
  40                 builtin = true
  41                 return
  42         }
  43         protocolIndex := int(id - 1)
  44         if protocolIndex >= len(me.Index) {
  45                 err = fmt.Errorf("unexpected extended message ID: %v", id)
  46                 return
  47         }
  48         builtin = protocolIndex < me.NumBuiltin
  49         name = me.Index[protocolIndex]
  50         return
  51 }
  52
  53 func (me *LocalLtepProtocolMap) builtin() []pp.ExtensionName {
  54         return me.Index[:me.NumBuiltin]
  55 }
  56
  57 func (me *LocalLtepProtocolMap) user() []pp.ExtensionName {
  58         return me.Index[me.NumBuiltin:]
  59 }
  60
  61 func (me *LocalLtepProtocolMap) AddUserProtocol(name pp.ExtensionName) {
  62         builtin := slices.DeleteFunc(me.builtin(), func(delName pp.ExtensionName) bool {
  63                 return delName == name
  64         })
  65         user := slices.DeleteFunc(me.user(), func(delName pp.ExtensionName) bool {
  66                 return delName == name
  67         })
  68         me.Index = append(append(builtin, user...), name)
  69         me.NumBuiltin = len(builtin)
  70 }