pex.go

   1 package torrent
   2
   3 import (
   4         "net"
   5
   6         "github.com/anacrolix/dht/v2/krpc"
   7         pp "github.com/anacrolix/torrent/peer_protocol"
   8 )
   9
  10 type pexEventType int
  11
  12 const (
  13         pexAdd pexEventType = iota
  14         pexDrop
  15 )
  16
  17 // internal, based on BEP11
  18 const (
  19         pexTargAdded = 25 // put drops on hold when the number of alive connections is lower than this
  20         pexMaxHold   = 25 // length of the drop hold-back buffer
  21         pexMaxDelta  = 50 // upper bound on added+added6 and dropped+dropped6 in a single PEX message
  22 )
  23
  24 // represents a single connection (t=pexAdd) or disconnection (t=pexDrop) event
  25 type pexEvent struct {
  26         t    pexEventType
  27         addr net.Addr
  28         f    pp.PexPeerFlags
  29 }
  30
  31 // Combines the node addr, as required for pp.PexMsg.
  32 type pexMsgAdded struct {
  33         krpc.NodeAddr
  34         pp.PexPeerFlags
  35 }
  36
  37 // Makes generating a PexMsg more efficient.
  38 type pexMsgFactory struct {
  39         added   map[string]pexMsgAdded
  40         dropped map[string]krpc.NodeAddr
  41 }
  42
  43 func (me *pexMsgFactory) DeltaLen() int {
  44         return int(max(
  45                 int64(len(me.added)),
  46                 int64(len(me.dropped))))
  47 }
  48
  49 // Returns the key to use to identify a given addr in the factory. Panics if we can't support the
  50 // addr later in generating a PexMsg (since adding an unusable addr will cause DeltaLen to be out.)
  51 func (me *pexMsgFactory) addrKey(addr krpc.NodeAddr) string {
  52         if addr.IP.To4() != nil {
  53                 addr.IP = addr.IP.To4()
  54         }
  55         keyBytes, err := addr.MarshalBinary()
  56         if err != nil {
  57                 panic(err)
  58         }
  59         switch len(keyBytes) {
  60         case compactIpv4NodeAddrElemSize:
  61         case compactIpv6NodeAddrElemSize:
  62         default:
  63                 panic(len(keyBytes))
  64         }
  65         return string(keyBytes)
  66 }
  67
  68 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
  69 // won't hit the limit consuming this event).
  70 func (me *pexMsgFactory) Add(addr krpc.NodeAddr, flags pp.PexPeerFlags) bool {
  71         key := me.addrKey(addr)
  72         if _, ok := me.dropped[key]; ok {
  73                 delete(me.dropped, key)
  74                 return true
  75         }
  76         if me.DeltaLen() >= pexMaxDelta {
  77                 return false
  78         }
  79         if me.added == nil {
  80                 me.added = make(map[string]pexMsgAdded, pexMaxDelta)
  81         }
  82         me.added[key] = pexMsgAdded{addr, flags}
  83         return true
  84
  85 }
  86
  87 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
  88 // won't hit the limit consuming this event).
  89 func (me *pexMsgFactory) Drop(addr krpc.NodeAddr) bool {
  90         key := me.addrKey(addr)
  91         if _, ok := me.added[key]; ok {
  92                 delete(me.added, key)
  93                 return true
  94         }
  95         if me.DeltaLen() >= pexMaxDelta {
  96                 return false
  97         }
  98         if me.dropped == nil {
  99                 me.dropped = make(map[string]krpc.NodeAddr, pexMaxDelta)
 100         }
 101         me.dropped[key] = addr
 102         return true
 103 }
 104
 105 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
 106 // won't hit the limit consuming this event).
 107 func (me *pexMsgFactory) addEvent(event pexEvent) bool {
 108         addr, ok := nodeAddr(event.addr)
 109         if !ok {
 110                 return true
 111         }
 112         switch event.t {
 113         case pexAdd:
 114                 return me.Add(addr, event.f)
 115         case pexDrop:
 116                 return me.Drop(addr)
 117         default:
 118                 panic(event.t)
 119         }
 120 }
 121
 122 var compactIpv4NodeAddrElemSize = krpc.CompactIPv4NodeAddrs{}.ElemSize()
 123 var compactIpv6NodeAddrElemSize = krpc.CompactIPv6NodeAddrs{}.ElemSize()
 124
 125 func (me *pexMsgFactory) PexMsg() (ret pp.PexMsg) {
 126         for key, added := range me.added {
 127                 switch len(key) {
 128                 case compactIpv4NodeAddrElemSize:
 129                         ret.Added = append(ret.Added, added.NodeAddr)
 130                         ret.AddedFlags = append(ret.AddedFlags, added.PexPeerFlags)
 131                 case compactIpv6NodeAddrElemSize:
 132                         ret.Added6 = append(ret.Added6, added.NodeAddr)
 133                         ret.Added6Flags = append(ret.Added6Flags, added.PexPeerFlags)
 134                 default:
 135                         panic(key)
 136                 }
 137         }
 138         for key, addr := range me.dropped {
 139                 switch len(key) {
 140                 case compactIpv4NodeAddrElemSize:
 141                         ret.Dropped = append(ret.Dropped, addr)
 142                 case compactIpv6NodeAddrElemSize:
 143                         ret.Dropped6 = append(ret.Dropped6, addr)
 144                 default:
 145                         panic(key)
 146                 }
 147         }
 148         return
 149 }
 150
 151 func mustNodeAddr(addr net.Addr) krpc.NodeAddr {
 152         ret, ok := nodeAddr(addr)
 153         if !ok {
 154                 panic(addr)
 155         }
 156         return ret
 157 }
 158
 159 // Convert an arbitrary torrent peer Addr into one that can be represented by the compact addr
 160 // format.
 161 func nodeAddr(addr net.Addr) (_ krpc.NodeAddr, ok bool) {
 162         ipport, ok := tryIpPortFromNetAddr(addr)
 163         if !ok {
 164                 return
 165         }
 166         return krpc.NodeAddr{IP: shortestIP(ipport.IP), Port: ipport.Port}, true
 167 }
 168
 169 // mainly for the krpc marshallers
 170 func shortestIP(ip net.IP) net.IP {
 171         if ip4 := ip.To4(); ip4 != nil {
 172                 return ip4
 173         }
 174         return ip
 175 }
 176
 177 // Per-torrent PEX state
 178 type pexState struct {
 179         ev   []pexEvent // event feed, append-only
 180         hold []pexEvent // delayed drops
 181         nc   int        // net number of alive conns
 182 }
 183
 184 func (s *pexState) Reset() {
 185         s.ev = nil
 186         s.hold = nil
 187         s.nc = 0
 188 }
 189
 190 func (s *pexState) Add(c *PeerConn) {
 191         s.nc++
 192         if s.nc >= pexTargAdded {
 193                 s.ev = append(s.ev, s.hold...)
 194                 s.hold = s.hold[:0]
 195         }
 196         e := c.pexEvent(pexAdd)
 197         s.ev = append(s.ev, e)
 198         c.pex.Listed = true
 199 }
 200
 201 func (s *pexState) Drop(c *PeerConn) {
 202         if !c.pex.Listed {
 203                 // skip connections which were not previously Added
 204                 return
 205         }
 206         e := c.pexEvent(pexDrop)
 207         s.nc--
 208         if s.nc < pexTargAdded && len(s.hold) < pexMaxHold {
 209                 s.hold = append(s.hold, e)
 210         } else {
 211                 s.ev = append(s.ev, e)
 212         }
 213 }
 214
 215 // Generate a PEX message based on the event feed. Also returns an index to pass to the subsequent
 216 // calls, producing incremental deltas.
 217 func (s *pexState) Genmsg(start int) (pp.PexMsg, int) {
 218         var factory pexMsgFactory
 219         n := start
 220         for _, e := range s.ev[start:] {
 221                 if !factory.addEvent(e) {
 222                         break
 223                 }
 224                 n++
 225         }
 226         return factory.PexMsg(), n
 227 }