FedP2P/peer_protocol/pex.go

package peer_protocol

import (
	"net"

	"github.com/anacrolix/dht/v2/krpc"
	"github.com/anacrolix/torrent/bencode"
)

type PexMsg struct {
	Added       krpc.CompactIPv4NodeAddrs `bencode:"added"`
	AddedFlags  []PexPeerFlags            `bencode:"added.f"`
	Added6      krpc.CompactIPv6NodeAddrs `bencode:"added6"`
	Added6Flags []PexPeerFlags            `bencode:"added6.f"`
	Dropped     krpc.CompactIPv4NodeAddrs `bencode:"dropped"`
	Dropped6    krpc.CompactIPv6NodeAddrs `bencode:"dropped6"`
}

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 := range v {
		if addrEqual(&v[i], a) {
			return i
		}
	}
	return -1
}

func (m *PexMsg) Add(addr krpc.NodeAddr, f PexPeerFlags) {
	if addr.IP.To4() != nil {
		if addrIndex(m.Added.NodeAddrs(), &addr) >= 0 {
			// already added
			return
		}
		if i := addrIndex(m.Dropped.NodeAddrs(), &addr); i >= 0 {
			// on the dropped list - cancel out
			m.Dropped = append(m.Dropped[:i], m.Dropped[i+1:]...)
			return
		}
		m.Added = append(m.Added, addr)
		m.AddedFlags = append(m.AddedFlags, f)
	} else if len(addr.IP) == net.IPv6len {
		if addrIndex(m.Added6.NodeAddrs(), &addr) >= 0 {
			// already added
			return
		}
		if i := addrIndex(m.Dropped6.NodeAddrs(), &addr); i >= 0 {
			// on the dropped list - cancel out
			m.Dropped6 = append(m.Dropped6[:i], m.Dropped6[i+1:]...)
			return
		}
		m.Added6 = append(m.Added6, addr)
		m.Added6Flags = append(m.Added6Flags, f)
	}
}

func (m *PexMsg) Drop(addr krpc.NodeAddr) {
	if addr.IP.To4() != nil {
		if addrIndex(m.Dropped.NodeAddrs(), &addr) >= 0 {
			// already dropped
			return
		}
		if i := addrIndex(m.Added.NodeAddrs(), &addr); i >= 0 {
			// on the added list - cancel out
			m.Added = append(m.Added[:i], m.Added[i+1:]...)
			m.AddedFlags = append(m.AddedFlags[:i], m.AddedFlags[i+1:]...)
			return
		}
		m.Dropped = append(m.Dropped, addr)
	} else if len(addr.IP) == net.IPv6len {
		if addrIndex(m.Dropped6.NodeAddrs(), &addr) >= 0 {
			// already dropped
			return
		}
		if i := addrIndex(m.Added6.NodeAddrs(), &addr); i >= 0 {
			// on the added list - cancel out
			m.Added6 = append(m.Added6[:i], m.Added6[i+1:]...)
			m.Added6Flags = append(m.Added6Flags[:i], m.Added6Flags[i+1:]...)
			return
		}
		m.Dropped6 = append(m.Dropped6, addr)
	}
}

func (m *PexMsg) Len() int {
	return len(m.Added)+len(m.Added6)+len(m.Dropped)+len(m.Dropped6)
}

// DeltaLen returns max of {added+added6, dropped+dropped6}
func (m *PexMsg) DeltaLen() int {
	lenAdded := len(m.Added)+len(m.Added6)
	lenDropped := len(m.Dropped)+len(m.Dropped6)
	if lenAdded > lenDropped {
		return lenAdded
	}
	return lenDropped
}

func (m *PexMsg) Message(pexExtendedId ExtensionNumber) Message {
	payload := bencode.MustMarshal(m)
	return Message{
		Type:            Extended,
		ExtendedID:      pexExtendedId,
		ExtendedPayload: payload,
	}
}

func LoadPexMsg(b []byte) (*PexMsg, error) {
	m := new(PexMsg) 
	if err := bencode.Unmarshal(b, m); err != nil {
		return nil, err
	}
	return m, nil
}


type PexPeerFlags byte

func (me PexPeerFlags) Get(f PexPeerFlags) bool {
	return me&f == f
}

const (
	PexPrefersEncryption PexPeerFlags = 1 << iota
	PexSeedUploadOnly
	PexSupportsUtp
	PexHolepunchSupport
	PexOutgoingConn
)