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 }