-
-
Notifications
You must be signed in to change notification settings - Fork 631
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Add low level support for BEP 10 user protocols
- Loading branch information
Showing
9 changed files
with
302 additions
and
31 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,69 @@ | ||
| package torrent | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "slices" | ||
|
|
||
| g "github.com/anacrolix/generics" | ||
| pp "github.com/anacrolix/torrent/peer_protocol" | ||
| ) | ||
|
|
||
| type LocalLtepProtocolMap struct { | ||
| // 1-based mapping from extension number to extension name (subtract one from the extension ID | ||
| // to find the corresponding protocol name). The first LocalLtepProtocolBuiltinCount of these | ||
| // are use builtin handlers. If you want to handle builtin protocols yourself, you would move | ||
| // them above the threshold. You can disable them by removing them entirely, and add your own. | ||
| // These changes should be done in the PeerConnAdded callback. | ||
| Index []pp.ExtensionName | ||
| // How many of the protocols are using the builtin handlers. | ||
| NumBuiltin int | ||
| } | ||
|
|
||
| func (me *LocalLtepProtocolMap) toSupportedExtensionDict() (m map[pp.ExtensionName]pp.ExtensionNumber) { | ||
| g.MakeMapWithCap(&m, len(me.Index)) | ||
| for i, name := range me.Index { | ||
| old := g.MapInsert(m, name, pp.ExtensionNumber(i+1)) | ||
| if old.Ok { | ||
| panic(fmt.Sprintf("extension %q already defined with id %v", name, old.Value)) | ||
| } | ||
| } | ||
| return | ||
| } | ||
|
|
||
| // Returns the local extension name for the given ID. If builtin is true, the implementation intends | ||
| // to handle it itself. For incoming messages with extension ID 0, the message is a handshake, and | ||
| // should be treated specially. | ||
| func (me *LocalLtepProtocolMap) LookupId(id pp.ExtensionNumber) (name pp.ExtensionName, builtin bool, err error) { | ||
| if id == 0 { | ||
| err = fmt.Errorf("extension ID 0 is handshake") | ||
| builtin = true | ||
| return | ||
| } | ||
| protocolIndex := int(id - 1) | ||
| if protocolIndex >= len(me.Index) { | ||
| err = fmt.Errorf("unexpected extended message ID: %v", id) | ||
| return | ||
| } | ||
| builtin = protocolIndex < me.NumBuiltin | ||
| name = me.Index[protocolIndex] | ||
| return | ||
| } | ||
|
|
||
| func (me *LocalLtepProtocolMap) builtin() []pp.ExtensionName { | ||
| return me.Index[:me.NumBuiltin] | ||
| } | ||
|
|
||
| func (me *LocalLtepProtocolMap) user() []pp.ExtensionName { | ||
| return me.Index[me.NumBuiltin:] | ||
| } | ||
|
|
||
| func (me *LocalLtepProtocolMap) AddUserProtocol(name pp.ExtensionName) { | ||
| builtin := slices.DeleteFunc(me.builtin(), func(delName pp.ExtensionName) bool { | ||
| return delName == name | ||
| }) | ||
| user := slices.DeleteFunc(me.user(), func(delName pp.ExtensionName) bool { | ||
| return delName == name | ||
| }) | ||
| me.Index = append(append(builtin, user...), name) | ||
| me.NumBuiltin = len(builtin) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,130 @@ | ||
| package torrent_test | ||
|
|
||
| import ( | ||
| "strconv" | ||
| "testing" | ||
|
|
||
| pp "github.com/anacrolix/torrent/peer_protocol" | ||
|
|
||
| qt "github.com/frankban/quicktest" | ||
|
|
||
| "github.com/anacrolix/torrent/internal/testutil" | ||
|
|
||
| "github.com/anacrolix/sync" | ||
|
|
||
| . "github.com/anacrolix/torrent" | ||
| ) | ||
|
|
||
| const ( | ||
| testRepliesToOddsExtensionName = "pm_me_odds" | ||
| testRepliesToEvensExtensionName = "pm_me_evens" | ||
| ) | ||
|
|
||
| func countHandler( | ||
| c *qt.C, | ||
| wg *sync.WaitGroup, | ||
| // Name of the endpoint that this handler is for, for logging. | ||
| handlerName string, | ||
| // Whether we expect evens or odds | ||
| expectedMod2 uint, | ||
| // Extension name of messages we expect to handle. | ||
| answerToName pp.ExtensionName, | ||
| // Extension name of messages we expect to send. | ||
| replyToName pp.ExtensionName, | ||
| // Signal done when this value is seen. | ||
| doneValue uint, | ||
| ) func(event PeerConnReadExtensionMessageEvent) { | ||
| return func(event PeerConnReadExtensionMessageEvent) { | ||
| // Read handshake, don't look it up. | ||
| if event.ExtensionNumber == 0 { | ||
| return | ||
| } | ||
| name, builtin, err := event.PeerConn.LocalLtepProtocolMap.LookupId(event.ExtensionNumber) | ||
| c.Assert(err, qt.IsNil) | ||
| // Not a user protocol. | ||
| if builtin { | ||
| return | ||
| } | ||
| switch name { | ||
| case answerToName: | ||
| u64, err := strconv.ParseUint(string(event.Payload), 10, 0) | ||
| c.Assert(err, qt.IsNil) | ||
| i := uint(u64) | ||
| c.Logf("%v got %d", handlerName, i) | ||
| if i == doneValue { | ||
| wg.Done() | ||
| return | ||
| } | ||
| c.Assert(i%2, qt.Equals, expectedMod2) | ||
| go func() { | ||
| c.Assert( | ||
| event.PeerConn.WriteExtendedMessage( | ||
| replyToName, | ||
| []byte(strconv.FormatUint(uint64(i+1), 10))), | ||
| qt.IsNil) | ||
| }() | ||
| default: | ||
| c.Fatalf("got unexpected extension name %q", name) | ||
| } | ||
| } | ||
| } | ||
|
|
||
| func TestUserLtep(t *testing.T) { | ||
| c := qt.New(t) | ||
| var wg sync.WaitGroup | ||
|
|
||
| makeCfg := func() *ClientConfig { | ||
| cfg := TestingConfig(t) | ||
| // Only want a single connection to between the clients. | ||
| cfg.DisableUTP = true | ||
| cfg.DisableIPv6 = true | ||
| return cfg | ||
| } | ||
|
|
||
| evensCfg := makeCfg() | ||
| evensCfg.Callbacks.ReadExtendedHandshake = func(pc *PeerConn, msg *pp.ExtendedHandshakeMessage) { | ||
| // The client lock is held while handling this event, so we have to do synchronous work in a | ||
| // separate goroutine. | ||
| go func() { | ||
| // Check sending an extended message for a protocol the peer doesn't support is an error. | ||
| c.Check(pc.WriteExtendedMessage("pm_me_floats", []byte("3.142")), qt.IsNotNil) | ||
| // Kick things off by sending a 1. | ||
| c.Check(pc.WriteExtendedMessage(testRepliesToOddsExtensionName, []byte("1")), qt.IsNil) | ||
| }() | ||
| } | ||
| evensCfg.Callbacks.PeerConnReadExtensionMessage = append( | ||
| evensCfg.Callbacks.PeerConnReadExtensionMessage, | ||
| countHandler(c, &wg, "evens", 0, testRepliesToEvensExtensionName, testRepliesToOddsExtensionName, 100)) | ||
| evensCfg.Callbacks.PeerConnAdded = append(evensCfg.Callbacks.PeerConnAdded, func(conn *PeerConn) { | ||
| conn.LocalLtepProtocolMap.AddUserProtocol(testRepliesToEvensExtensionName) | ||
| c.Assert(conn.LocalLtepProtocolMap.Index[conn.LocalLtepProtocolMap.NumBuiltin:], qt.HasLen, 1) | ||
| }) | ||
|
|
||
| oddsCfg := makeCfg() | ||
| oddsCfg.Callbacks.PeerConnAdded = append(oddsCfg.Callbacks.PeerConnAdded, func(conn *PeerConn) { | ||
| conn.LocalLtepProtocolMap.AddUserProtocol(testRepliesToOddsExtensionName) | ||
| c.Assert(conn.LocalLtepProtocolMap.Index[conn.LocalLtepProtocolMap.NumBuiltin:], qt.HasLen, 1) | ||
| }) | ||
| oddsCfg.Callbacks.PeerConnReadExtensionMessage = append( | ||
| oddsCfg.Callbacks.PeerConnReadExtensionMessage, | ||
| countHandler(c, &wg, "odds", 1, testRepliesToOddsExtensionName, testRepliesToEvensExtensionName, 100)) | ||
|
|
||
| cl1, err := NewClient(oddsCfg) | ||
| c.Assert(err, qt.IsNil) | ||
| defer cl1.Close() | ||
| cl2, err := NewClient(evensCfg) | ||
| c.Assert(err, qt.IsNil) | ||
| defer cl2.Close() | ||
| addOpts := AddTorrentOpts{} | ||
| t1, _ := cl1.AddTorrentOpt(addOpts) | ||
| t2, _ := cl2.AddTorrentOpt(addOpts) | ||
| defer testutil.ExportStatusWriter(cl1, "cl1", t)() | ||
| defer testutil.ExportStatusWriter(cl2, "cl2", t)() | ||
| // Expect one PeerConn to see the value. | ||
| wg.Add(1) | ||
| added := t1.AddClientPeer(cl2) | ||
| // Ensure some addresses for the other client were added. | ||
| c.Assert(added, qt.Not(qt.Equals), 0) | ||
| wg.Wait() | ||
| _ = t2 | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -24,7 +24,7 @@ type ( | |
| } | ||
|
|
||
| ExtensionName string | ||
| ExtensionNumber int | ||
| ExtensionNumber uint8 | ||
| ) | ||
|
|
||
| const ( | ||
|
|
||
Oops, something went wrong.