pex.go

   1 package torrent
   2
   3 import (
   4         "net"
   5         "sync"
   6         "time"
   7
   8         "github.com/anacrolix/dht/v2/krpc"
   9
  10         pp "github.com/anacrolix/torrent/peer_protocol"
  11 )
  12
  13 type pexEventType int
  14
  15 const (
  16         pexAdd pexEventType = iota
  17         pexDrop
  18 )
  19
  20 // internal, based on BEP11
  21 const (
  22         pexTargAdded = 25 // put drops on hold when the number of alive connections is lower than this
  23         pexMaxHold   = 25 // length of the drop hold-back buffer
  24         pexMaxDelta  = 50 // upper bound on added+added6 and dropped+dropped6 in a single PEX message
  25 )
  26
  27 // represents a single connection (t=pexAdd) or disconnection (t=pexDrop) event
  28 type pexEvent struct {
  29         t    pexEventType
  30         addr PeerRemoteAddr
  31         f    pp.PexPeerFlags
  32         next *pexEvent // event feed list
  33 }
  34
  35 // facilitates efficient de-duplication while generating PEX messages
  36 type pexMsgFactory struct {
  37         msg     pp.PexMsg
  38         added   map[addrKey]struct{}
  39         dropped map[addrKey]struct{}
  40 }
  41
  42 func (me *pexMsgFactory) DeltaLen() int {
  43         return int(max(
  44                 int64(len(me.added)),
  45                 int64(len(me.dropped))))
  46 }
  47
  48 type addrKey string
  49
  50 // Returns the key to use to identify a given addr in the factory.
  51 func (me *pexMsgFactory) addrKey(addr PeerRemoteAddr) addrKey {
  52         return addrKey(addr.String())
  53 }
  54
  55 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
  56 // won't hit the limit consuming this event).
  57 func (me *pexMsgFactory) add(e pexEvent) {
  58         key := me.addrKey(e.addr)
  59         if _, ok := me.added[key]; ok {
  60                 return
  61         }
  62         if me.added == nil {
  63                 me.added = make(map[addrKey]struct{}, pexMaxDelta)
  64         }
  65         addr, ok := nodeAddr(e.addr)
  66         if !ok {
  67                 return
  68         }
  69         m := &me.msg
  70         switch {
  71         case addr.IP.To4() != nil:
  72                 if _, ok := me.dropped[key]; ok {
  73                         if i := m.Dropped.Index(addr); i >= 0 {
  74                                 m.Dropped = append(m.Dropped[:i], m.Dropped[i+1:]...)
  75                         }
  76                         delete(me.dropped, key)
  77                         return
  78                 }
  79                 m.Added = append(m.Added, addr)
  80                 m.AddedFlags = append(m.AddedFlags, e.f)
  81         case len(addr.IP) == net.IPv6len:
  82                 if _, ok := me.dropped[key]; ok {
  83                         if i := m.Dropped6.Index(addr); i >= 0 {
  84                                 m.Dropped6 = append(m.Dropped6[:i], m.Dropped6[i+1:]...)
  85                         }
  86                         delete(me.dropped, key)
  87                         return
  88                 }
  89                 m.Added6 = append(m.Added6, addr)
  90                 m.Added6Flags = append(m.Added6Flags, e.f)
  91         default:
  92                 panic(addr)
  93         }
  94         me.added[key] = struct{}{}
  95 }
  96
  97 // Returns whether the entry was added (we can check if we're cancelling out another entry and so
  98 // won't hit the limit consuming this event).
  99 func (me *pexMsgFactory) drop(e pexEvent) {
 100         addr, ok := nodeAddr(e.addr)
 101         if !ok {
 102                 return
 103         }
 104         key := me.addrKey(e.addr)
 105         if me.dropped == nil {
 106                 me.dropped = make(map[addrKey]struct{}, pexMaxDelta)
 107         }
 108         if _, ok := me.dropped[key]; ok {
 109                 return
 110         }
 111         m := &me.msg
 112         switch {
 113         case addr.IP.To4() != nil:
 114                 if _, ok := me.added[key]; ok {
 115                         if i := m.Added.Index(addr); i >= 0 {
 116                                 m.Added = append(m.Added[:i], m.Added[i+1:]...)
 117                                 m.AddedFlags = append(m.AddedFlags[:i], m.AddedFlags[i+1:]...)
 118                         }
 119                         delete(me.added, key)
 120                         return
 121                 }
 122                 m.Dropped = append(m.Dropped, addr)
 123         case len(addr.IP) == net.IPv6len:
 124                 if _, ok := me.added[key]; ok {
 125                         if i := m.Added6.Index(addr); i >= 0 {
 126                                 m.Added6 = append(m.Added6[:i], m.Added6[i+1:]...)
 127                                 m.Added6Flags = append(m.Added6Flags[:i], m.Added6Flags[i+1:]...)
 128                         }
 129                         delete(me.added, key)
 130                         return
 131                 }
 132                 m.Dropped6 = append(m.Dropped6, addr)
 133         }
 134         me.dropped[key] = struct{}{}
 135 }
 136
 137 func (me *pexMsgFactory) append(event pexEvent) {
 138         switch event.t {
 139         case pexAdd:
 140                 me.add(event)
 141         case pexDrop:
 142                 me.drop(event)
 143         default:
 144                 panic(event.t)
 145         }
 146 }
 147
 148 func (me *pexMsgFactory) PexMsg() *pp.PexMsg {
 149         return &me.msg
 150 }
 151
 152 // Convert an arbitrary torrent peer Addr into one that can be represented by the compact addr
 153 // format.
 154 func nodeAddr(addr PeerRemoteAddr) (krpc.NodeAddr, bool) {
 155         ipport, _ := tryIpPortFromNetAddr(addr)
 156         ok := ipport.IP != nil
 157         return krpc.NodeAddr{IP: ipport.IP, Port: ipport.Port}, ok
 158 }
 159
 160 // Per-torrent PEX state
 161 type pexState struct {
 162         sync.RWMutex
 163         tail *pexEvent     // event feed list
 164         hold []pexEvent    // delayed drops
 165         rest time.Time     // cooldown deadline on inbound
 166         nc   int           // net number of alive conns
 167         msg0 pexMsgFactory // initial message
 168 }
 169
 170 // Reset wipes the state clean, releasing resources. Called from Torrent.Close().
 171 func (s *pexState) Reset() {
 172         s.Lock()
 173         defer s.Unlock()
 174         s.tail = nil
 175         s.hold = nil
 176         s.nc = 0
 177         s.rest = time.Time{}
 178         s.msg0 = pexMsgFactory{}
 179 }
 180
 181 func (s *pexState) append(e *pexEvent) {
 182         if s.tail != nil {
 183                 s.tail.next = e
 184         }
 185         s.tail = e
 186         s.msg0.append(*e)
 187 }
 188
 189 func (s *pexState) Add(c *PeerConn) {
 190         s.Lock()
 191         defer s.Unlock()
 192         s.nc++
 193         if s.nc >= pexTargAdded {
 194                 for _, e := range s.hold {
 195                         ne := e
 196                         s.append(&ne)
 197                 }
 198                 s.hold = s.hold[:0]
 199         }
 200         e := c.pexEvent(pexAdd)
 201         c.pex.Listed = true
 202         s.append(&e)
 203 }
 204
 205 func (s *pexState) Drop(c *PeerConn) {
 206         if !c.pex.Listed {
 207                 // skip connections which were not previously Added
 208                 return
 209         }
 210         s.Lock()
 211         defer s.Unlock()
 212         e := c.pexEvent(pexDrop)
 213         s.nc--
 214         if s.nc < pexTargAdded && len(s.hold) < pexMaxHold {
 215                 s.hold = append(s.hold, e)
 216         } else {
 217                 s.append(&e)
 218         }
 219 }
 220
 221 // Generate a PEX message based on the event feed.
 222 // Also returns a pointer to pass to the subsequent calls
 223 // to produce incremental deltas.
 224 func (s *pexState) Genmsg(start *pexEvent) (pp.PexMsg, *pexEvent) {
 225         s.RLock()
 226         defer s.RUnlock()
 227         if start == nil {
 228                 return *s.msg0.PexMsg(), s.tail
 229         }
 230         var msg pexMsgFactory
 231         last := start
 232         for e := start.next; e != nil; e = e.next {
 233                 if msg.DeltaLen() >= pexMaxDelta {
 234                         break
 235                 }
 236                 msg.append(*e)
 237                 last = e
 238         }
 239         return *msg.PexMsg(), last
 240 }
 241
 242 // The same as Genmsg but just counts up the distinct events that haven't been sent.
 243 func (s *pexState) numPending(start *pexEvent) (num int) {
 244         s.RLock()
 245         defer s.RUnlock()
 246         if start == nil {
 247                 return s.msg0.PexMsg().Len()
 248         }
 249         for e := start.next; e != nil; e = e.next {
 250                 num++
 251         }
 252         return
 253 }