Abstracts GossipSubRouter interface

discovery_test.go

@@ -244,7 +244,7 @@ func TestGossipSubDiscoveryAfterBootstrap(t *testing.T) {
 			s = server2
 		}
 		disc := &mockDiscoveryClient{h, s}
-		ps := getGossipsub(ctx, h, WithDiscovery(disc, WithDiscoveryOpts(discOpts...)))
+		ps := getGossipSub(ctx, h, WithDiscovery(disc, WithDiscoveryOpts(discOpts...)))
 		psubs[i] = ps
 		topicHandlers[i], _ = ps.Join(topic)
 	}

gossip_tracer.go

@@ -9,10 +9,10 @@ import (
 	"github.com/libp2p/go-libp2p/core/protocol"
 )
 
-// gossipTracer is an internal tracer that tracks IWANT requests in order to penalize
+// GossipTracer is an internal tracer that tracks IWANT requests in order to penalize
 // peers who don't follow up on IWANT requests after an IHAVE advertisement.
 // The tracking of promises is probabilistic to avoid using too much memory.
-type gossipTracer struct {
+type GossipTracer struct {
 	sync.Mutex
 
 	idGen *msgIDGenerator
@@ -27,15 +27,15 @@ type gossipTracer struct {
 	peerPromises map[peer.ID]map[string]struct{}
 }
 
-func newGossipTracer() *gossipTracer {
-	return &gossipTracer{
+func newGossipTracer() *GossipTracer {
+	return &GossipTracer{
 		idGen:        newMsgIdGenerator(),
 		promises:     make(map[string]map[peer.ID]time.Time),
 		peerPromises: make(map[peer.ID]map[string]struct{}),
 	}
 }
 
-func (gt *gossipTracer) Start(gs *GossipSubRouter) {
+func (gt *GossipTracer) Start(gs *GossipSubRouter) {
 	if gt == nil {
 		return
 	}
@@ -45,7 +45,7 @@ func (gt *gossipTracer) Start(gs *GossipSubRouter) {
 }
 
 // track a promise to deliver a message from a list of msgIDs we are requesting
-func (gt *gossipTracer) AddPromise(p peer.ID, msgIDs []string) {
+func (gt *GossipTracer) AddPromise(p peer.ID, msgIDs []string) {
 	if gt == nil {
 		return
 	}
@@ -76,7 +76,7 @@ func (gt *gossipTracer) AddPromise(p peer.ID, msgIDs []string) {
 
 // returns the number of broken promises for each peer who didn't follow up
 // on an IWANT request.
-func (gt *gossipTracer) GetBrokenPromises() map[peer.ID]int {
+func (gt *GossipTracer) GetBrokenPromises() map[peer.ID]int {
 	if gt == nil {
 		return nil
 	}
@@ -114,9 +114,9 @@ func (gt *gossipTracer) GetBrokenPromises() map[peer.ID]int {
 	return res
 }
 
-var _ RawTracer = (*gossipTracer)(nil)
+var _ RawTracer = (*GossipTracer)(nil)
 
-func (gt *gossipTracer) fulfillPromise(msg *Message) {
+func (gt *GossipTracer) fulfillPromise(msg *Message) {
 	mid := gt.idGen.ID(msg)
 
 	gt.Lock()
@@ -140,12 +140,12 @@ func (gt *gossipTracer) fulfillPromise(msg *Message) {
 	}
 }
 
-func (gt *gossipTracer) DeliverMessage(msg *Message) {
+func (gt *GossipTracer) DeliverMessage(msg *Message) {
 	// someone delivered a message, fulfill promises for it
 	gt.fulfillPromise(msg)
 }
 
-func (gt *gossipTracer) RejectMessage(msg *Message, reason string) {
+func (gt *GossipTracer) RejectMessage(msg *Message, reason string) {
 	// A message got rejected, so we can fulfill promises and let the score penalty apply
 	// from invalid message delivery.
 	// We do take exception and apply promise penalty regardless in the following cases, where
@@ -160,26 +160,26 @@ func (gt *gossipTracer) RejectMessage(msg *Message, reason string) {
 	gt.fulfillPromise(msg)
 }
 
-func (gt *gossipTracer) ValidateMessage(msg *Message) {
+func (gt *GossipTracer) ValidateMessage(msg *Message) {
 	// we consider the promise fulfilled as soon as the message begins validation
 	// if it was a case of signature issue it would have been rejected immediately
 	// without triggering the Validate trace
 	gt.fulfillPromise(msg)
 }
 
-func (gt *gossipTracer) AddPeer(p peer.ID, proto protocol.ID) {}
-func (gt *gossipTracer) RemovePeer(p peer.ID)                 {}
-func (gt *gossipTracer) Join(topic string)                    {}
-func (gt *gossipTracer) Leave(topic string)                   {}
-func (gt *gossipTracer) Graft(p peer.ID, topic string)        {}
-func (gt *gossipTracer) Prune(p peer.ID, topic string)        {}
-func (gt *gossipTracer) DuplicateMessage(msg *Message)        {}
-func (gt *gossipTracer) RecvRPC(rpc *RPC)                     {}
-func (gt *gossipTracer) SendRPC(rpc *RPC, p peer.ID)          {}
-func (gt *gossipTracer) DropRPC(rpc *RPC, p peer.ID)          {}
-func (gt *gossipTracer) UndeliverableMessage(msg *Message)    {}
-
-func (gt *gossipTracer) ThrottlePeer(p peer.ID) {
+func (gt *GossipTracer) AddPeer(p peer.ID, proto protocol.ID) {}
+func (gt *GossipTracer) RemovePeer(p peer.ID)                 {}
+func (gt *GossipTracer) Join(topic string)                    {}
+func (gt *GossipTracer) Leave(topic string)                   {}
+func (gt *GossipTracer) Graft(p peer.ID, topic string)        {}
+func (gt *GossipTracer) Prune(p peer.ID, topic string)        {}
+func (gt *GossipTracer) DuplicateMessage(msg *Message)        {}
+func (gt *GossipTracer) RecvRPC(rpc *RPC)                     {}
+func (gt *GossipTracer) SendRPC(rpc *RPC, p peer.ID)          {}
+func (gt *GossipTracer) DropRPC(rpc *RPC, p peer.ID)          {}
+func (gt *GossipTracer) UndeliverableMessage(msg *Message)    {}
+
+func (gt *GossipTracer) ThrottlePeer(p peer.ID) {
 	gt.Lock()
 	defer gt.Unlock()
 

gossipsub.go

@@ -206,7 +206,10 @@ type GossipSubParams struct {
 
 // NewGossipSub returns a new PubSub object using the default GossipSubRouter as the router.
 func NewGossipSub(ctx context.Context, h host.Host, opts ...Option) (*PubSub, error) {
-	rt := DefaultGossipSubRouter(h)
+	rt, err := DefaultGossipSubRouter(h)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create default gossipsub router: %w", err)
+	}
 	opts = append(opts, WithRawTracer(rt.tagTracer))
 	return NewGossipSubWithRouter(ctx, h, rt, opts...)
 }
@@ -216,10 +219,12 @@ func NewGossipSubWithRouter(ctx context.Context, h host.Host, rt PubSubRouter, o
 	return NewPubSub(ctx, h, rt, opts...)
 }
 
+type GossipSubRouterOption func(*GossipSubRouter) error
+
 // DefaultGossipSubRouter returns a new GossipSubRouter with default parameters.
-func DefaultGossipSubRouter(h host.Host) *GossipSubRouter {
+func DefaultGossipSubRouter(h host.Host, opts ...GossipSubRouterOption) (*GossipSubRouter, error) {
 	params := DefaultGossipSubParams()
-	return &GossipSubRouter{
+	rt := &GossipSubRouter{
 		peers:     make(map[peer.ID]protocol.ID),
 		mesh:      make(map[string]map[peer.ID]struct{}),
 		fanout:    make(map[string]map[peer.ID]struct{}),
@@ -237,6 +242,14 @@ func DefaultGossipSubRouter(h host.Host) *GossipSubRouter {
 		tagTracer: newTagTracer(h.ConnManager()),
 		params:    params,
 	}
+
+	for _, opt := range opts {
+		if err := opt(rt); err != nil {
+			return nil, fmt.Errorf("failed to apply gossipsub router option: %w", err)
+		}
+	}
+
+	return rt, nil
 }
 
 // DefaultGossipSubParams returns the default gossip sub parameters
@@ -277,7 +290,7 @@ func DefaultGossipSubParams() GossipSubParams {
 // WithPeerScore is a gossipsub router option that enables peer scoring.
 func WithPeerScore(params *PeerScoreParams, thresholds *PeerScoreThresholds) Option {
 	return func(ps *PubSub) error {
-		gs, ok := ps.rt.(*GossipSubRouter)
+		gs, ok := ps.rt.(GossipPubSubRouter)
 		if !ok {
 			return fmt.Errorf("pubsub router is not gossipsub")
 		}
@@ -294,21 +307,17 @@ func WithPeerScore(params *PeerScoreParams, thresholds *PeerScoreThresholds) Opt
 			return err
 		}
 
-		gs.score = newPeerScore(params)
-		gs.gossipThreshold = thresholds.GossipThreshold
-		gs.publishThreshold = thresholds.PublishThreshold
-		gs.graylistThreshold = thresholds.GraylistThreshold
-		gs.acceptPXThreshold = thresholds.AcceptPXThreshold
-		gs.opportunisticGraftThreshold = thresholds.OpportunisticGraftThreshold
+		gs.SetPeerScore(newPeerScore(params))
+		gs.SetPeerScoreThresholds(thresholds)
 
-		gs.gossipTracer = newGossipTracer()
+		gs.SetGossipTracer(newGossipTracer())
 
 		// hook the tracer
 		if ps.tracer != nil {
-			ps.tracer.raw = append(ps.tracer.raw, gs.score, gs.gossipTracer)
+			ps.tracer.raw = append(ps.tracer.raw, gs.GetPeerScore(), gs.GetGossipTracer())
 		} else {
 			ps.tracer = &pubsubTracer{
-				raw:   []RawTracer{gs.score, gs.gossipTracer},
+				raw:   []RawTracer{gs.GetPeerScore(), gs.GetGossipTracer()},
 				pid:   ps.host.ID(),
 				idGen: ps.idGen,
 			}
@@ -321,31 +330,19 @@ func WithPeerScore(params *PeerScoreParams, thresholds *PeerScoreThresholds) Opt
 // WithFloodPublish is a gossipsub router option that enables flood publishing.
 // When this is enabled, published messages are forwarded to all peers with score >=
 // to publishThreshold
-func WithFloodPublish(floodPublish bool) Option {
-	return func(ps *PubSub) error {
-		gs, ok := ps.rt.(*GossipSubRouter)
-		if !ok {
-			return fmt.Errorf("pubsub router is not gossipsub")
-		}
-
+func WithFloodPublish(floodPublish bool) GossipSubRouterOption {
+	return func(gs *GossipSubRouter) error {
 		gs.floodPublish = floodPublish
-
 		return nil
 	}
 }
 
 // WithPeerExchange is a gossipsub router option that enables Peer eXchange on PRUNE.
 // This should generally be enabled in bootstrappers and well connected/trusted nodes
 // used for bootstrapping.
-func WithPeerExchange(doPX bool) Option {
-	return func(ps *PubSub) error {
-		gs, ok := ps.rt.(*GossipSubRouter)
-		if !ok {
-			return fmt.Errorf("pubsub router is not gossipsub")
-		}
-
+func WithPeerExchange(doPX bool) GossipSubRouterOption {
+	return func(gs *GossipSubRouter) error {
 		gs.doPX = doPX
-
 		return nil
 	}
 }
@@ -357,7 +354,7 @@ func WithPeerExchange(doPX bool) Option {
 // symmetrically configured at both ends.
 func WithDirectPeers(pis []peer.AddrInfo) Option {
 	return func(ps *PubSub) error {
-		gs, ok := ps.rt.(*GossipSubRouter)
+		gs, ok := ps.rt.(GossipPubSubRouter)
 		if !ok {
 			return fmt.Errorf("pubsub router is not gossipsub")
 		}
@@ -368,10 +365,10 @@ func WithDirectPeers(pis []peer.AddrInfo) Option {
 			ps.host.Peerstore().AddAddrs(pi.ID, pi.Addrs, peerstore.PermanentAddrTTL)
 		}
 
-		gs.direct = direct
+		gs.SetDirectPeers(direct)
 
-		if gs.tagTracer != nil {
-			gs.tagTracer.direct = direct
+		if gs.GetTagTracer() != nil {
+			gs.GetTagTracer().direct = direct
 		}
 
 		return nil
@@ -382,25 +379,17 @@ func WithDirectPeers(pis []peer.AddrInfo) Option {
 // heartbeat ticks between attempting to reconnect direct peers that are not
 // currently connected. A "tick" is based on the heartbeat interval, which is
 // 1s by default. The default value for direct connect ticks is 300.
-func WithDirectConnectTicks(t uint64) Option {
-	return func(ps *PubSub) error {
-		gs, ok := ps.rt.(*GossipSubRouter)
-		if !ok {
-			return fmt.Errorf("pubsub router is not gossipsub")
-		}
+func WithDirectConnectTicks(t uint64) GossipSubRouterOption {
+	return func(gs *GossipSubRouter) error {
 		gs.params.DirectConnectTicks = t
 		return nil
 	}
 }
 
 // WithGossipSubParams is a gossip sub router option that allows a custom
 // config to be set when instantiating the gossipsub router.
-func WithGossipSubParams(cfg GossipSubParams) Option {
-	return func(ps *PubSub) error {
-		gs, ok := ps.rt.(*GossipSubRouter)
-		if !ok {
-			return fmt.Errorf("pubsub router is not gossipsub")
-		}
+func WithGossipSubParams(cfg GossipSubParams) GossipSubRouterOption {
+	return func(gs *GossipSubRouter) error {
 		// Overwrite current config and associated variables in the router.
 		gs.params = cfg
 		gs.connect = make(chan connectInfo, cfg.MaxPendingConnections)
@@ -410,6 +399,25 @@ func WithGossipSubParams(cfg GossipSubParams) Option {
 	}
 }
 
+type GossipPubSubRouter interface {
+	PubSubRouter
+
+	SetPeerScore(*PeerScore)
+	GetPeerScore() *PeerScore
+
+	SetPeerScoreThresholds(*PeerScoreThresholds)
+
+	SetGossipTracer(*GossipTracer)
+	GetGossipTracer() *GossipTracer
+
+	GetTagTracer() *TagTracer
+
+	SetDirectPeers(map[peer.ID]struct{})
+
+	SetPeerGater(*PeerGater)
+	GetPeerGater() *PeerGater
+}
+
 // GossipSubRouter is a router that implements the gossipsub protocol.
 // For each topic we have joined, we maintain an overlay through which
 // messages flow; this is the mesh map.
@@ -437,10 +445,10 @@ type GossipSubRouter struct {
 
 	mcache       *MessageCache
 	tracer       *pubsubTracer
-	score        *peerScore
-	gossipTracer *gossipTracer
-	tagTracer    *tagTracer
-	gate         *peerGater
+	score        *PeerScore
+	gossipTracer *GossipTracer
+	tagTracer    *TagTracer
+	gate         *PeerGater
 
 	// config for gossipsub parameters
 	params GossipSubParams
@@ -476,6 +484,8 @@ type GossipSubRouter struct {
 	heartbeatTicks uint64
 }
 
+var _ GossipPubSubRouter = (*GossipSubRouter)(nil)
+
 type connectInfo struct {
 	p   peer.ID
 	spr *record.Envelope
@@ -1917,6 +1927,46 @@ func (gs *GossipSubRouter) getPeers(topic string, count int, filter func(peer.ID
 	return peers
 }
 
+func (gs *GossipSubRouter) SetPeerScore(score *PeerScore) {
+	gs.score = score
+}
+
+func (gs *GossipSubRouter) GetPeerScore() *PeerScore {
+	return gs.score
+}
+
+func (gs *GossipSubRouter) SetPeerScoreThresholds(thresholds *PeerScoreThresholds) {
+	gs.gossipThreshold = thresholds.GossipThreshold
+	gs.publishThreshold = thresholds.PublishThreshold
+	gs.graylistThreshold = thresholds.GraylistThreshold
+	gs.acceptPXThreshold = thresholds.AcceptPXThreshold
+	gs.opportunisticGraftThreshold = thresholds.OpportunisticGraftThreshold
+}
+
+func (gs *GossipSubRouter) SetGossipTracer(tracer *GossipTracer) {
+	gs.gossipTracer = tracer
+}
+
+func (gs *GossipSubRouter) GetGossipTracer() *GossipTracer {
+	return gs.gossipTracer
+}
+
+func (gs *GossipSubRouter) GetTagTracer() *TagTracer {
+	return gs.tagTracer
+}
+
+func (gs *GossipSubRouter) SetDirectPeers(direct map[peer.ID]struct{}) {
+	gs.direct = direct
+}
+
+func (gs *GossipSubRouter) SetPeerGater(gater *PeerGater) {
+	gs.gate = gater
+}
+
+func (gs *GossipSubRouter) GetPeerGater() *PeerGater {
+	return gs.gate
+}
+
 // WithDefaultTagTracer returns the tag tracer of the GossipSubRouter as a PubSub option.
 // This is useful for cases where the GossipSubRouter is instantiated externally, and is
 // injected into the GossipSub constructor as a dependency. This allows the tag tracer to be