gossipsub: implement extensions

Author: MarcoPolo

extensions.go

@@ -0,0 +1,131 @@
+package pubsub
+
+import (
+	pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
+	"github.com/libp2p/go-libp2p/core/peer"
+)
+
+type PeerExtensions struct {
+	TestExtension bool
+}
+
+func WithTestExtension() Option {
+	return func(ps *PubSub) error {
+		if rt, ok := ps.rt.(*GossipSubRouter); ok {
+			rt.extensions.myExtensions.TestExtension = true
+		}
+		return nil
+	}
+}
+
+func hasPeerExtensions(rpc *RPC) bool {
+	if rpc != nil && rpc.Control != nil && rpc.Control.Extensions != nil {
+		return true
+	}
+	return false
+}
+
+func peerExtensionsFromRPC(rpc *RPC) PeerExtensions {
+	out := PeerExtensions{}
+	if hasPeerExtensions(rpc) {
+		out.TestExtension = rpc.Control.Extensions.GetTestExtension()
+	}
+	return out
+}
+
+func (pe *PeerExtensions) ExtendRPC(rpc *RPC) {
+	if pe.TestExtension {
+		if rpc.Control == nil {
+			rpc.Control = &pubsub_pb.ControlMessage{}
+		}
+		rpc.Control.Extensions = &pubsub_pb.ControlExtensions{
+			TestExtension: &pe.TestExtension,
+		}
+	}
+}
+
+type extensionsState struct {
+	myExtensions      PeerExtensions
+	peerExtensions    map[peer.ID]PeerExtensions // peer's extensions
+	sentExtensions    map[peer.ID]struct{}
+	reportMisbehavior func(peer.ID)
+	sendRPC           func(p peer.ID, r *RPC, urgent bool)
+	testExtension     testExtension
+}
+
+func newExtensionsState(myExtensions PeerExtensions, reportMisbehavior func(peer.ID), sendRPC func(peer.ID, *RPC, bool)) *extensionsState {
+	return &extensionsState{
+		myExtensions:      myExtensions,
+		peerExtensions:    make(map[peer.ID]PeerExtensions),
+		sentExtensions:    make(map[peer.ID]struct{}),
+		reportMisbehavior: reportMisbehavior,
+		sendRPC:           sendRPC,
+		testExtension:     testExtension{sendRPC: sendRPC},
+	}
+}
+
+func (es *extensionsState) HandleRPC(rpc *RPC) {
+	if _, ok := es.peerExtensions[rpc.from]; !ok {
+		// We know this is the first message because we didn't have extensions
+		// for this peer, and we always set extensions on the first rpc.
+		es.peerExtensions[rpc.from] = peerExtensionsFromRPC(rpc)
+		if _, ok := es.sentExtensions[rpc.from]; ok {
+			// We just finished both sending and receiving the extensions
+			// control message.
+			es.extensionsAddPeer(rpc.from)
+		}
+	} else {
+		// We already have an extension for this peer. If they send us another
+		// extensions control message, that is a protocol error. We should
+		// down score them because they are misbehaving.
+		if hasPeerExtensions(rpc) {
+			es.reportMisbehavior(rpc.from)
+		}
+	}
+
+	es.extensionsHandleRPC(rpc)
+}
+
+func (es *extensionsState) AddPeer(id peer.ID, helloPacket *RPC) {
+	// Send our extensions as the first message.
+	es.myExtensions.ExtendRPC(helloPacket)
+
+	es.sentExtensions[id] = struct{}{}
+	if _, ok := es.peerExtensions[id]; ok {
+		// We've just finished sending and receiving the extensions control
+		// message.
+		es.extensionsAddPeer(id)
+	}
+}
+
+func (es *extensionsState) RemovePeer(id peer.ID) {
+	_, recvdExt := es.peerExtensions[id]
+	_, sentExt := es.sentExtensions[id]
+	if recvdExt && sentExt {
+		// Add peer was previously called, so we need to call remove peer
+		es.extensionsRemovePeer(id)
+	}
+	delete(es.peerExtensions, id)
+	if len(es.peerExtensions) == 0 {
+		es.peerExtensions = make(map[peer.ID]PeerExtensions)
+	}
+	delete(es.sentExtensions, id)
+	if len(es.sentExtensions) == 0 {
+		es.sentExtensions = make(map[peer.ID]struct{})
+	}
+}
+
+// extensionsAddPeer is only called once we've both sent and received the
+// extensions control message.
+func (es *extensionsState) extensionsAddPeer(id peer.ID) {
+	if es.myExtensions.TestExtension && es.peerExtensions[id].TestExtension {
+		es.testExtension.AddPeer(id)
+	}
+}
+
+// extensionsRemovePeer is always called after extensionsAddPeer.
+func (es *extensionsState) extensionsRemovePeer(id peer.ID) {
+}
+
+func (es *extensionsState) extensionsHandleRPC(rpc *RPC) {
+}

gossipsub.go

@@ -260,7 +260,7 @@ func NewGossipSubWithRouter(ctx context.Context, h host.Host, rt PubSubRouter, o
 // DefaultGossipSubRouter returns a new GossipSubRouter with default parameters.
 func DefaultGossipSubRouter(h host.Host) *GossipSubRouter {
 	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{}),
@@ -281,6 +281,14 @@ func DefaultGossipSubRouter(h host.Host) *GossipSubRouter {
 		tagTracer:    newTagTracer(h.ConnManager()),
 		params:       params,
 	}
+
+	rt.extensions = newExtensionsState(PeerExtensions{}, func(p peer.ID) {
+		if rt.score != nil {
+			rt.score.AddPenalty(p, 10)
+		}
+	}, rt.sendRPC)
+
+	return rt
 }
 
 // DefaultGossipSubParams returns the default gossip sub parameters
@@ -466,8 +474,10 @@ func WithGossipSubParams(cfg GossipSubParams) Option {
 // is the fanout map. Fanout peer lists are expired if we don't publish any
 // messages to their topic for GossipSubFanoutTTL.
 type GossipSubRouter struct {
-	p            *PubSub
-	peers        map[peer.ID]protocol.ID          // peer protocols
+	p          *PubSub
+	peers      map[peer.ID]protocol.ID // peer protocols
+	extensions *extensionsState
+
 	direct       map[peer.ID]struct{}             // direct peers
 	mesh         map[string]map[peer.ID]struct{}  // topic meshes
 	fanout       map[string]map[peer.ID]struct{}  // topic fanout
@@ -652,11 +662,18 @@ loop:
 		}
 	}
 	gs.outbound[p] = outbound
+
+	if gs.feature(GossipSubFeatureExtensions, proto) {
+		gs.extensions.AddPeer(p, helloPacket)
+	}
 }
 
 func (gs *GossipSubRouter) RemovePeer(p peer.ID) {
 	log.Debugf("PEERDOWN: Remove disconnected peer %s", p)
 	gs.tracer.RemovePeer(p)
+	if gs.feature(GossipSubFeatureExtensions, gs.peers[p]) {
+		gs.extensions.RemovePeer(p)
+	}
 	delete(gs.peers, p)
 	for _, peers := range gs.mesh {
 		delete(peers, p)
@@ -746,6 +763,8 @@ func (gs *GossipSubRouter) Preprocess(from peer.ID, msgs []*Message) {
 }
 
 func (gs *GossipSubRouter) HandleRPC(rpc *RPC) {
+	gs.extensions.HandleRPC(rpc)
+
 	ctl := rpc.GetControl()
 	if ctl == nil {
 		return

gossipsub_test.go

@@ -3961,3 +3961,106 @@ func newSkeletonGossipsub(ctx context.Context, h host.Host) *skeletonGossipsub {
 		inRPC:  sendRPC,
 	}
 }
+
+func TestExtensionsControlMessage(t *testing.T) {
+	for _, wellBehaved := range []bool{true, false} {
+		t.Run(fmt.Sprintf("wellBehaved=%t", wellBehaved), func(t *testing.T) {
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			hosts := getDefaultHosts(t, 2)
+			psub0 := getGossipsub(ctx, hosts[0],
+				WithPeerScore(
+					&PeerScoreParams{
+						AppSpecificScore:       func(peer.ID) float64 { return 0 },
+						BehaviourPenaltyWeight: -1,
+						BehaviourPenaltyDecay:  ScoreParameterDecay(time.Minute),
+						DecayInterval:          DefaultDecayInterval,
+						DecayToZero:            DefaultDecayToZero,
+					},
+					&PeerScoreThresholds{
+						GossipThreshold:   -100,
+						PublishThreshold:  -500,
+						GraylistThreshold: -1000,
+					}),
+				WithMessageIdFn(func(msg *pb.Message) string {
+					return string(msg.Data)
+				}))
+
+			psub1 := newSkeletonGossipsub(ctx, hosts[1])
+
+			connect(t, hosts[0], hosts[1])
+			time.Sleep(time.Second)
+
+			loopTimes := 3
+
+			for i := range loopTimes {
+				rpcToSend := &pb.RPC{
+					Control: &pb.ControlMessage{
+						Extensions: &pb.ControlExtensions{},
+					},
+				}
+				if wellBehaved && i > 0 {
+					// A well behaved node does not repeat the control
+					// extension message
+					rpcToSend.Control.Extensions = nil
+				}
+				psub1.inRPC <- rpcToSend
+			}
+
+			time.Sleep(time.Second)
+
+			peerScore := psub0.rt.(*GossipSubRouter).score.Score(hosts[1].ID())
+			t.Log("Peer score:", peerScore)
+			if wellBehaved {
+				if peerScore < 0 {
+					t.Fatal("Peer score should not be negative")
+				}
+			} else {
+				if peerScore >= 0 {
+					t.Fatal("Peer score should not be positive")
+				}
+			}
+		})
+	}
+}
+
+func TestTestExtension(t *testing.T) {
+	hosts := getDefaultHosts(t, 2)
+	psub := getGossipsub(context.Background(), hosts[0], WithTestExtension())
+	_ = psub
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	psub1 := newSkeletonGossipsub(ctx, hosts[1])
+
+	connect(t, hosts[0], hosts[1])
+
+	const timeout = 3 * time.Second
+	select {
+	case <-time.After(timeout):
+		t.Fatal("Timeout")
+	case r := <-psub1.outRPC:
+		if !*r.Control.Extensions.TestExtension {
+			t.Fatal("Unexpected RPC. First RPC should be the Extensions Control Message")
+		}
+	}
+
+	truePtr := true
+	psub1.inRPC <- &pb.RPC{
+		Control: &pb.ControlMessage{
+			Extensions: &pb.ControlExtensions{
+				TestExtension: &truePtr,
+			},
+		},
+	}
+
+	select {
+	case <-time.After(timeout):
+		t.Fatal("Timeout")
+	case r := <-psub1.outRPC:
+		if r.TestExtension == nil {
+			t.Fatal("Unexpected RPC. Next RPC should be the TestExtension Message")
+		}
+	}
+}

testextension.go

@@ -0,0 +1,18 @@
+package pubsub
+
+import (
+	pubsub_pb "github.com/libp2p/go-libp2p-pubsub/pb"
+	"github.com/libp2p/go-libp2p/core/peer"
+)
+
+type testExtension struct {
+	sendRPC func(p peer.ID, r *RPC, urgent bool)
+}
+
+func (e *testExtension) AddPeer(id peer.ID) {
+	e.sendRPC(id, &RPC{
+		RPC: pubsub_pb.RPC{
+			TestExtension: &pubsub_pb.TestExtension{},
+		},
+	}, false)
+}