201 lines
4.1 KiB
Go
201 lines
4.1 KiB
Go
package zcash
|
|
|
|
import (
|
|
mrand "math/rand"
|
|
"net"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/btcsuite/btcd/wire"
|
|
)
|
|
|
|
type Address struct {
|
|
netaddr *wire.NetAddress
|
|
blacklisted bool
|
|
lastUpdate time.Time
|
|
}
|
|
|
|
func (a *Address) String() string {
|
|
portString := strconv.Itoa(int(a.netaddr.Port))
|
|
return net.JoinHostPort(a.netaddr.IP.String(), portString)
|
|
}
|
|
|
|
func (a *Address) asPeerKey() PeerKey {
|
|
return PeerKey(a.String())
|
|
}
|
|
|
|
func (a *Address) fromPeerKey(s PeerKey) (*Address, error) {
|
|
host, portString, err := net.SplitHostPort(s.String())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
portInt, err := strconv.ParseUint(portString, 10, 16)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
na := wire.NewNetAddressTimestamp(
|
|
time.Now(),
|
|
0,
|
|
net.ParseIP(host),
|
|
uint16(portInt),
|
|
)
|
|
|
|
a.netaddr = na
|
|
a.blacklisted = false
|
|
a.lastUpdate = na.Timestamp
|
|
return a, nil
|
|
}
|
|
|
|
func (a *Address) asNetAddress() *wire.NetAddress {
|
|
newNA := *a.netaddr
|
|
newNA.Timestamp = a.lastUpdate
|
|
return &newNA
|
|
}
|
|
|
|
func (a *Address) fromNetAddress(na *wire.NetAddress) (*Address, error) {
|
|
a.netaddr = na
|
|
a.blacklisted = false
|
|
a.lastUpdate = na.Timestamp
|
|
return a, nil
|
|
}
|
|
|
|
func (a *Address) MarshalText() (text []byte, err error) {
|
|
return []byte(a.String()), nil
|
|
}
|
|
|
|
type AddressBook struct {
|
|
addrs map[PeerKey]*Address
|
|
addrState sync.RWMutex
|
|
addrRecvCond *sync.Cond
|
|
}
|
|
|
|
func NewAddressBook() *AddressBook {
|
|
addrBook := &AddressBook{
|
|
addrs: make(map[PeerKey]*Address),
|
|
}
|
|
addrBook.addrRecvCond = sync.NewCond(&addrBook.addrState)
|
|
return addrBook
|
|
}
|
|
|
|
func (bk *AddressBook) Add(s PeerKey) {
|
|
newAddr, err := (&Address{}).fromPeerKey(s)
|
|
if err != nil {
|
|
// XXX effectively NOP bogus peer strings
|
|
return
|
|
}
|
|
|
|
bk.addrState.Lock()
|
|
bk.addrs[s] = newAddr
|
|
bk.addrState.Unlock()
|
|
|
|
// Wake anyone who was waiting on us to receive an address.
|
|
bk.addrRecvCond.Broadcast()
|
|
}
|
|
|
|
func (bk *AddressBook) Remove(s PeerKey) {
|
|
bk.addrState.Lock()
|
|
defer bk.addrState.Unlock()
|
|
|
|
if _, ok := bk.addrs[s]; ok {
|
|
delete(bk.addrs, s)
|
|
}
|
|
}
|
|
|
|
func (bk *AddressBook) Blacklist(s PeerKey) {
|
|
bk.addrState.Lock()
|
|
defer bk.addrState.Unlock()
|
|
|
|
if target, ok := bk.addrs[s]; ok {
|
|
target.blacklisted = true
|
|
target.lastUpdate = time.Now()
|
|
} else {
|
|
// Create a new Address just to be blacklisted
|
|
addr, err := (&Address{}).fromPeerKey(s)
|
|
if err != nil {
|
|
// XXX effectively NOP bogus peer strings
|
|
return
|
|
}
|
|
addr.blacklisted = true
|
|
bk.addrs[s] = addr
|
|
}
|
|
}
|
|
|
|
// Touch updates the last-seen timestamp if the peer is in the address book or does nothing if not.
|
|
func (bk *AddressBook) Touch(s PeerKey) {
|
|
bk.addrState.Lock()
|
|
defer bk.addrState.Unlock()
|
|
|
|
if target, ok := bk.addrs[s]; ok {
|
|
target.lastUpdate = time.Now()
|
|
}
|
|
}
|
|
|
|
// IsKnown returns true if the peer is already in our address book, false if not.
|
|
func (bk *AddressBook) IsKnown(s PeerKey) bool {
|
|
bk.addrState.RLock()
|
|
defer bk.addrState.RUnlock()
|
|
|
|
_, known := bk.addrs[s]
|
|
return known
|
|
}
|
|
|
|
func (bk *AddressBook) IsBlacklisted(s PeerKey) bool {
|
|
bk.addrState.RLock()
|
|
defer bk.addrState.RUnlock()
|
|
|
|
if target, ok := bk.addrs[s]; ok {
|
|
return target.blacklisted
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
// WaitForAddresses waits for n addresses to be received and their initial
|
|
// connection attempts to resolve. There is no escape if that does not happen -
|
|
// this is intended for test runners or goroutines with a timeout.
|
|
func (bk *AddressBook) waitForAddresses(n int, done chan struct{}) {
|
|
bk.addrState.Lock()
|
|
for {
|
|
addrCount := len(bk.addrs)
|
|
if addrCount < n {
|
|
bk.addrRecvCond.Wait()
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
bk.addrState.Unlock()
|
|
done <- struct{}{}
|
|
return
|
|
}
|
|
|
|
// GetAddressList returns a slice of n valid addresses in random order.
|
|
// If there aren't enough known addresses, it returns as many as we have.
|
|
func (bk *AddressBook) shuffleAddressList(n int) []net.IP {
|
|
bk.addrState.RLock()
|
|
defer bk.addrState.RUnlock()
|
|
|
|
resp := make([]net.IP, 0, len(bk.addrs))
|
|
|
|
for _, v := range bk.addrs {
|
|
if v.blacklisted {
|
|
continue
|
|
}
|
|
|
|
resp = append(resp, v.netaddr.IP)
|
|
}
|
|
|
|
mrand.Seed(time.Now().UnixNano())
|
|
mrand.Shuffle(len(resp), func(i, j int) {
|
|
resp[i], resp[j] = resp[j], resp[i]
|
|
})
|
|
|
|
if len(resp) > n {
|
|
return resp[:n]
|
|
}
|
|
|
|
return resp
|
|
}
|