pex.go

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