torrent/pex.go

package torrent

import (
	"net"
	"sync"
	"time"

	"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 PeerRemoteAddr
	f    pp.PexPeerFlags
}

// facilitates efficient de-duplication while generating PEX messages
type pexMsgFactory struct {
	msg     pp.PexMsg
	added   map[addrKey]struct{}
	dropped map[addrKey]struct{}
}

func (me *pexMsgFactory) DeltaLen() int {
	return int(max(
		int64(len(me.added)),
		int64(len(me.dropped))))
}

type addrKey string

// Returns the key to use to identify a given addr in the factory.
func (me *pexMsgFactory) addrKey(addr PeerRemoteAddr) addrKey {
	return addrKey(addr.String())
}

func addrEqual(a, b *krpc.NodeAddr) bool {
	return a.IP.Equal(b.IP) && a.Port == b.Port
}

func addrIndex(v []krpc.NodeAddr, a *krpc.NodeAddr) int {
	for i := 0; i < len(v); i += 1 {
		if addrEqual(&v[i], a) {
			return i
		}
	}
	return -1
}

// Returns whether the entry was added (we can check if we're cancelling out another entry and so
// won't hit the limit consuming this event).
func (me *pexMsgFactory) add(e pexEvent) {
	key := me.addrKey(e.addr)
	if _, ok := me.added[key]; ok {
		return
	}
	if me.added == nil {
		me.added = make(map[addrKey]struct{}, pexMaxDelta)
	}
	addr, ok := nodeAddr(e.addr)
	if !ok {
		return
	}
	m := &me.msg
	switch {
	case addr.IP.To4() != nil:
		if _, ok := me.dropped[key]; ok {
			if i := addrIndex(m.Dropped.NodeAddrs(), &addr); i >= 0 {
				m.Dropped = append(m.Dropped[:i], m.Dropped[i+1:]...)
			}
			delete(me.dropped, key)
			return
		}
		m.Added = append(m.Added, addr)
		m.AddedFlags = append(m.AddedFlags, e.f)
	case len(addr.IP) == net.IPv6len:
		if _, ok := me.dropped[key]; ok {
			if i := addrIndex(m.Dropped6.NodeAddrs(), &addr); i >= 0 {
				m.Dropped6 = append(m.Dropped6[:i], m.Dropped6[i+1:]...)
			}
			delete(me.dropped, key)
			return
		}
		m.Added6 = append(m.Added6, addr)
		m.Added6Flags = append(m.Added6Flags, e.f)
	default:
		panic(addr)
	}
	me.added[key] = struct{}{}
}

// Returns whether the entry was added (we can check if we're cancelling out another entry and so
// won't hit the limit consuming this event).
func (me *pexMsgFactory) drop(e pexEvent) {
	addr, ok := nodeAddr(e.addr)
	if !ok {
		return
	}
	key := me.addrKey(e.addr)
	if me.dropped == nil {
		me.dropped = make(map[addrKey]struct{}, pexMaxDelta)
	}
	if _, ok := me.dropped[key]; ok {
		return
	}
	m := &me.msg
	switch {
	case addr.IP.To4() != nil:
		if _, ok := me.added[key]; ok {
			if i := addrIndex(m.Added.NodeAddrs(), &addr); i >= 0 {
				m.Added = append(m.Added[:i], m.Added[i+1:]...)
				m.AddedFlags = append(m.AddedFlags[:i], m.AddedFlags[i+1:]...)
			}
			delete(me.added, key)
			return
		}
		m.Dropped = append(m.Dropped, addr)
	case len(addr.IP) == net.IPv6len:
		if _, ok := me.added[key]; ok {
			if i := addrIndex(m.Added6.NodeAddrs(), &addr); i >= 0 {
				m.Added6 = append(m.Added6[:i], m.Added6[i+1:]...)
				m.Added6Flags = append(m.Added6Flags[:i], m.Added6Flags[i+1:]...)
			}
			delete(me.added, key)
			return
		}
		m.Dropped6 = append(m.Dropped6, addr)
	}
	me.dropped[key] = struct{}{}
}

func (me *pexMsgFactory) addEvent(event pexEvent) {
	switch event.t {
	case pexAdd:
		me.add(event)
	case pexDrop:
		me.drop(event)
	default:
		panic(event.t)
	}
}

func (me *pexMsgFactory) PexMsg() pp.PexMsg {
	return me.msg
}

// Convert an arbitrary torrent peer Addr into one that can be represented by the compact addr
// format.
func nodeAddr(addr PeerRemoteAddr) (krpc.NodeAddr, bool) {
	ipport, _ := tryIpPortFromNetAddr(addr)
	ok := ipport.IP != nil
	return krpc.NodeAddr{IP: shortestIP(ipport.IP), Port: ipport.Port}, ok
}

// 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
	rest      time.Time     // cooldown deadline on inbound
	nc        int           // net number of alive conns
	initCache pexMsgFactory // last generated initial message
	initSeq   int           // number of events which went into initCache
	initLock  sync.RWMutex  // serialise access to initCache and initSeq
}

// Reset wipes the state clean, releasing resources. Called from Torrent.Close().
func (s *pexState) Reset() {
	s.ev = nil
	s.hold = nil
	s.nc = 0
	s.rest = time.Time{}
	s.initLock.Lock()
	s.initCache = pexMsgFactory{}
	s.initSeq = 0
	s.initLock.Unlock()
}

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) {
	if start == 0 {
		return s.genmsg0()
	}

	var factory pexMsgFactory
	n := start
	for _, e := range s.ev[start:] {
		if start > 0 && factory.DeltaLen() >= pexMaxDelta {
			break
		}
		factory.addEvent(e)
		n++
	}
	return factory.PexMsg(), n
}

func (s *pexState) genmsg0() (pp.PexMsg, int) {
	s.initLock.Lock()
	for _, e := range s.ev[s.initSeq:] {
		s.initCache.addEvent(e)
		s.initSeq++
	}
	s.initLock.Unlock()
	s.initLock.RLock()
	n := s.initSeq
	msg := s.initCache.PexMsg()
	s.initLock.RUnlock()
	return msg, n
}