waku-org/go-waku

package peerstore

import (
    "errors"
    "sync"

    "github.com/ethereum/go-ethereum/p2p/enode"
    "github.com/libp2p/go-libp2p/core/network"
    "github.com/libp2p/go-libp2p/core/peer"
    "github.com/libp2p/go-libp2p/core/peerstore"
    "github.com/waku-org/go-waku/waku/v2/protocol"
    "golang.org/x/exp/maps"
)

// Origin is used to determine how the peer is identified,
// either it is statically added or discovered via one of the discovery protocols
type Origin int64

const (
    Unknown Origin = iota
    Discv5
    Static
    PeerExchange
    DNSDiscovery
    Rendezvous
    PeerManager
)

const peerOrigin = "origin"
const peerENR = "enr"
const peerDirection = "direction"
const peerPubSubTopics = "pubSubTopics"
const peerScore = "score"

// ConnectionFailures contains connection failure information towards all peers
type ConnectionFailures struct {
    sync.RWMutex
    failures map[peer.ID]int
}

// WakuPeerstoreImpl is a implementation of WakuPeerStore
type WakuPeerstoreImpl struct {
    peerStore    peerstore.Peerstore
    connFailures ConnectionFailures
}

// WakuPeerstore is an interface for implementing WakuPeerStore
type WakuPeerstore interface {
    SetOrigin(p peer.ID, origin Origin) error
    Origin(p peer.ID) (Origin, error)
    PeersByOrigin(origin Origin) peer.IDSlice
    SetENR(p peer.ID, enr *enode.Node) error
    ENR(p peer.ID) (*enode.Node, error)
    AddConnFailure(pID peer.ID)
    ResetConnFailures(pID peer.ID)
    ConnFailures(pID peer.ID) int

    SetDirection(p peer.ID, direction network.Direction) error
    Direction(p peer.ID) (network.Direction, error)

    AddPubSubTopic(p peer.ID, topic string) error
    RemovePubSubTopic(p peer.ID, topic string) error
    PubSubTopics(p peer.ID) (protocol.TopicSet, error)
    SetPubSubTopics(p peer.ID, topics []string) error
    PeersByPubSubTopics(pubSubTopics []string, specificPeers ...peer.ID) peer.IDSlice
    PeersByPubSubTopic(pubSubTopic string, specificPeers ...peer.ID) peer.IDSlice

    SetScore(peer.ID, float64) error
    Score(peer.ID) (float64, error)
}

// NewWakuPeerstore creates a new WakuPeerStore object
func NewWakuPeerstore(p peerstore.Peerstore) peerstore.Peerstore {
    return &WakuPeerstoreImpl{
        peerStore: p,
        connFailures: ConnectionFailures{
            failures: make(map[peer.ID]int),
        },
    }
}

// SetOrigin sets origin for a specific peer.
func (ps *WakuPeerstoreImpl) SetOrigin(p peer.ID, origin Origin) error {
    return ps.peerStore.Put(p, peerOrigin, origin)
}

// Origin fetches the origin for a specific peer.
func (ps *WakuPeerstoreImpl) Origin(p peer.ID) (Origin, error) {
    result, err := ps.peerStore.Get(p, peerOrigin)
    if err != nil {
        return Unknown, err
    }

    return result.(Origin), nil
}

// SetScore sets score for a specific peer.
func (ps *WakuPeerstoreImpl) SetScore(p peer.ID, score float64) error {
    return ps.peerStore.Put(p, peerScore, score)
}

// Score fetches the peerScore for a specific peer.
func (ps *WakuPeerstoreImpl) Score(p peer.ID) (float64, error) {
    result, err := ps.peerStore.Get(p, peerScore)
    if err != nil {
        return -1, err
    }

    return result.(float64), nil
}

// PeersByOrigin returns the list of peers for a specific origin
func (ps *WakuPeerstoreImpl) PeersByOrigin(expectedOrigin Origin) peer.IDSlice {
    var result peer.IDSlice
    for _, p := range ps.Peers() {
        actualOrigin, err := ps.Origin(p)
        if err == nil && actualOrigin == expectedOrigin {
            result = append(result, p)
        }
    }
    return result
}

// SetENR sets the ENR record a peer
func (ps *WakuPeerstoreImpl) SetENR(p peer.ID, enr *enode.Node) error {
    return ps.peerStore.Put(p, peerENR, enr)
}

// ENR fetches the ENR record for a peer
func (ps *WakuPeerstoreImpl) ENR(p peer.ID) (*enode.Node, error) {
    result, err := ps.peerStore.Get(p, peerENR)
    if err != nil {
        return nil, err
    }
    return result.(*enode.Node), nil
}

// AddConnFailure increments connectionFailures for a peer
func (ps *WakuPeerstoreImpl) AddConnFailure(pID peer.ID) {
    ps.connFailures.Lock()
    defer ps.connFailures.Unlock()
    ps.connFailures.failures[pID]++
}

// ResetConnFailures resets connectionFailures for a peer to 0
func (ps *WakuPeerstoreImpl) ResetConnFailures(pID peer.ID) {
    ps.connFailures.Lock()
    defer ps.connFailures.Unlock()
    ps.connFailures.failures[pID] = 0
}

// ConnFailures fetches connectionFailures for a peer
func (ps *WakuPeerstoreImpl) ConnFailures(pID peer.ID) int {
    ps.connFailures.RLock()
    defer ps.connFailures.RUnlock()
    return ps.connFailures.failures[pID]
}

// SetDirection sets connection direction for a specific peer.
func (ps *WakuPeerstoreImpl) SetDirection(p peer.ID, direction network.Direction) error {
    return ps.peerStore.Put(p, peerDirection, direction)
}

// Direction fetches the connection direction (Inbound or outBound) for a specific peer
func (ps *WakuPeerstoreImpl) Direction(p peer.ID) (network.Direction, error) {
    result, err := ps.peerStore.Get(p, peerDirection)
    if err != nil {
        return network.DirUnknown, err
    }

    return result.(network.Direction), nil
}

// AddPubSubTopic adds a new pubSubTopic for a peer
func (ps *WakuPeerstoreImpl) AddPubSubTopic(p peer.ID, topic string) error {
    existingTopics, err := ps.PubSubTopics(p)
    if err != nil {
        return err
    }

    if _, found := existingTopics[topic]; found {
        return nil
    }
    existingTopics[topic] = struct{}{}
    return ps.peerStore.Put(p, peerPubSubTopics, maps.Keys(existingTopics))
}

// RemovePubSubTopic removes a pubSubTopic from the peer
func (ps *WakuPeerstoreImpl) RemovePubSubTopic(p peer.ID, topic string) error {
    existingTopics, err := ps.PubSubTopics(p)
    if err != nil {
        return err
    }

    if len(existingTopics) == 0 {
        return nil
    }

    delete(existingTopics, topic)

    err = ps.SetPubSubTopics(p, maps.Keys(existingTopics))
    if err != nil {
        return err
    }
    return nil
}

// SetPubSubTopics sets pubSubTopics for a peer, it also overrides existing ones that were set previously..
func (ps *WakuPeerstoreImpl) SetPubSubTopics(p peer.ID, topics []string) error {
    return ps.peerStore.Put(p, peerPubSubTopics, topics)
}

// PubSubTopics fetches list of pubSubTopics for a peer
func (ps *WakuPeerstoreImpl) PubSubTopics(p peer.ID) (protocol.TopicSet, error) {
    result, err := ps.peerStore.Get(p, peerPubSubTopics)
    if err != nil {
        if errors.Is(err, peerstore.ErrNotFound) {
            return protocol.NewTopicSet(), nil
        } else {
            return nil, err
        }
    }
    return protocol.NewTopicSet((result.([]string))...), nil
}

// PeersByPubSubTopic Returns list of peers that support list of pubSubTopics
// If specifiPeers are listed, filtering is done from them otherwise from all peers in peerstore
func (ps *WakuPeerstoreImpl) PeersByPubSubTopics(pubSubTopics []string, specificPeers ...peer.ID) peer.IDSlice {
    if specificPeers == nil {
        specificPeers = ps.Peers()
    }
    var result peer.IDSlice
    for _, p := range specificPeers {
        peerMatch := true
        peerTopics, err := ps.PubSubTopics(p)
        if err == nil {
            for _, t := range pubSubTopics {
                if _, ok := peerTopics[t]; !ok {
                    peerMatch = false
                    break
                }
            }
            if peerMatch {
                result = append(result, p)
            }
        } //Note: skipping a peer in case of an error as there would be others available.
    }
    return result
}

// PeersByPubSubTopic Returns list of peers that support a single pubSubTopic
// If specifiPeers are listed, filtering is done from them otherwise from all peers in peerstore
func (ps *WakuPeerstoreImpl) PeersByPubSubTopic(pubSubTopic string, specificPeers ...peer.ID) peer.IDSlice {
    if specificPeers == nil {
        specificPeers = ps.Peers()
    }
    var result peer.IDSlice
    for _, p := range specificPeers {
        topics, err := ps.PubSubTopics(p)
        if err == nil {
            for topic := range topics {
                if topic == pubSubTopic {
                    result = append(result, p)
                }
            }
        } //Note: skipping a peer in case of an error as there would be others available.
    }
    return result
}