WIP support for ut_holepunch

[btrtrc.git] / torrent.go

diff --git a/torrent.go b/torrent.go
index b0820e5b3b0e97f142ba71d1b7d7ac1fde10356c..285f5cc804cb07e159aed181c42aee2eff97c4e6 100644 (file)
--- a/torrent.go
+++ b/torrent.go
@@ -17,6 +17,8 @@ import (
        "time"
        "unsafe"
 
+       utHolepunch "github.com/anacrolix/torrent/peer_protocol/ut-holepunch"
+
        "github.com/RoaringBitmap/roaring"
        "github.com/anacrolix/chansync"
        "github.com/anacrolix/chansync/events"
@@ -2357,7 +2359,8 @@ func (t *Torrent) VerifyData() {
        }
 }
 
-// Start the process of connecting to the given peer for the given torrent if appropriate.
+// Start the process of connecting to the given peer for the given torrent if appropriate. I'm not
+// sure all the PeerInfo fields are being used.
 func (t *Torrent) initiateConn(peer PeerInfo) {
        if peer.Id == t.cl.peerID {
                return
@@ -2664,3 +2667,72 @@ func (t *Torrent) checkValidReceiveChunk(r Request) error {
        // catch most of the overflow manipulation stuff by checking index and begin above.
        return nil
 }
+
+func (t *Torrent) peerConnsWithRemoteAddrPort(addrPort netip.AddrPort) (ret []*PeerConn) {
+       for pc := range t.conns {
+               addr := pc.remoteAddrPort()
+               if !(addr.Ok && addr.Value == addrPort) {
+                       continue
+               }
+               ret = append(ret, pc)
+       }
+       return
+}
+
+func makeUtHolepunchMsgForPeerConn(
+       recipient *PeerConn,
+       msgType utHolepunch.MsgType,
+       addrPort netip.AddrPort,
+       errCode utHolepunch.ErrCode,
+) pp.Message {
+       utHolepunchMsg := utHolepunch.Msg{
+               MsgType:  msgType,
+               AddrPort: addrPort,
+               ErrCode:  errCode,
+       }
+       extendedPayload, err := utHolepunchMsg.MarshalBinary()
+       if err != nil {
+               panic(err)
+       }
+       return pp.Message{
+               Type:            pp.Extended,
+               ExtendedID:      MapMustGet(recipient.PeerExtensionIDs, utHolepunch.ExtensionName),
+               ExtendedPayload: extendedPayload,
+       }
+}
+
+func (t *Torrent) handleReceivedUtHolepunchMsg(msg utHolepunch.Msg, sender *PeerConn) error {
+       switch msg.MsgType {
+       case utHolepunch.Rendezvous:
+               sendMsg := func(
+                       pc *PeerConn,
+                       msgType utHolepunch.MsgType,
+                       addrPort netip.AddrPort,
+                       errCode utHolepunch.ErrCode,
+               ) {
+                       pc.write(makeUtHolepunchMsgForPeerConn(pc, msgType, addrPort, errCode))
+               }
+               targets := t.peerConnsWithRemoteAddrPort(msg.AddrPort)
+               if len(targets) == 0 {
+                       sendMsg(sender, utHolepunch.Error, msg.AddrPort, utHolepunch.NotConnected)
+                       break
+               }
+               for _, pc := range targets {
+                       if !pc.supportsExtension(utHolepunch.ExtensionName) {
+                               sendMsg(sender, utHolepunch.Error, msg.AddrPort, utHolepunch.NoSupport)
+                               continue
+                       }
+                       sendMsg(sender, utHolepunch.Connect, msg.AddrPort, 0)
+                       sendMsg(pc, utHolepunch.Connect, sender.remoteAddrPort().Unwrap(), 0)
+               }
+       case utHolepunch.Connect:
+               t.initiateConn(PeerInfo{
+                       Addr:   msg.AddrPort,
+                       Source: PeerSourceUtHolepunch,
+               })
+       case utHolepunch.Error:
+
+       default:
+               return fmt.Errorf("unhandled msg type %v", msg.MsgType)
+       }
+}