torrent/pex.go

106 lines
2.2 KiB
Go

package torrent
import (
"net"
"github.com/anacrolix/dht/v2/krpc"
pp "github.com/anacrolix/torrent/peer_protocol"
)
type pexEventType int
const (
pexAdd pexEventType = iota
pexDrop
)
// internal, based on BEP11
const (
pexTargAdded = 25 // put drops on hold when the number of alive connections is lower than this
pexMaxHold = 25 // length of the drop hold-back buffer
pexMaxDelta = 50 // upper bound on added+added6 and dropped+dropped6 in a single PEX message
)
// represents a single connection (t=pexAdd) or disconnection (t=pexDrop) event
type pexEvent struct {
t pexEventType
addr net.Addr
f pp.PexPeerFlags
}
// records the event into the peer protocol PEX message
func (e *pexEvent) put(m *pp.PexMsg) {
switch e.t {
case pexAdd:
m.Add(nodeAddr(e.addr), e.f)
case pexDrop:
m.Drop(nodeAddr(e.addr))
}
}
func nodeAddr(addr net.Addr) krpc.NodeAddr {
ipport, _ := tryIpPortFromNetAddr(addr)
return krpc.NodeAddr{IP: shortestIP(ipport.IP), Port: ipport.Port}
}
// mainly for the krpc marshallers
func shortestIP(ip net.IP) net.IP {
if ip4 := ip.To4(); ip4 != nil {
return ip4
}
return ip
}
// Per-torrent PEX state
type pexState struct {
ev []pexEvent // event feed, append-only
hold []pexEvent // delayed drops
nc int // net number of alive conns
}
func (s *pexState) Reset() {
s.ev = nil
s.hold = nil
s.nc = 0
}
func (s *pexState) Add(c *PeerConn) {
s.nc++
if s.nc >= pexTargAdded {
s.ev = append(s.ev, s.hold...)
s.hold = s.hold[:0]
}
e := c.pexEvent(pexAdd)
s.ev = append(s.ev, e)
c.pex.Listed = true
}
func (s *pexState) Drop(c *PeerConn) {
if !c.pex.Listed {
// skip connections which were not previously Added
return
}
e := c.pexEvent(pexDrop)
s.nc--
if s.nc < pexTargAdded && len(s.hold) < pexMaxHold {
s.hold = append(s.hold, e)
} else {
s.ev = append(s.ev, e)
}
}
// Generate a PEX message based on the event feed.
// Also returns an index to pass to the subsequent calls, producing incremental deltas.
func (s *pexState) Genmsg(start int) (*pp.PexMsg, int) {
m := new(pp.PexMsg)
n := start
for _, e := range s.ev[start:] {
if start > 0 && m.DeltaLen() >= pexMaxDelta {
break
}
e.put(m)
n++
}
return m, n
}