github.com/anacrolix/torrent@v1.61.0/ltep.go (about)

     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  }