fix(p2p): do proper msg.ValidatorData handling (#281)

Author: cristaloleg

p2p/subscriber.go

@@ -3,6 +3,7 @@ package p2p
 import (
 	"context"
 	"errors"
+	"fmt"
 	"sync"
 
 	pubsub "github.com/libp2p/go-libp2p-pubsub"
@@ -150,12 +151,6 @@ func (s *Subscriber[H]) verifyMessage(
 	p peer.ID,
 	msg *pubsub.Message,
 ) (res pubsub.ValidationResult) {
-	if msg.ValidatorData != nil {
-		// means the message is local and was already validated
-		// so simply accept it
-		return pubsub.ValidationAccept
-	}
-
 	defer func() {
 		err := recover()
 		if err != nil {
@@ -164,17 +159,7 @@ func (s *Subscriber[H]) verifyMessage(
 		}
 	}()
 
-	hdr := header.New[H]()
-	err := hdr.UnmarshalBinary(msg.Data)
-	if err != nil {
-		log.Errorw("unmarshalling header",
-			"from", p.ShortString(),
-			"err", err)
-		s.metrics.reject(ctx)
-		return pubsub.ValidationReject
-	}
-	// ensure header validity
-	err = hdr.Validate()
+	hdr, err := s.extractHeader(msg)
 	if err != nil {
 		log.Errorw("invalid header",
 			"from", p.ShortString(),
@@ -197,20 +182,37 @@ func (s *Subscriber[H]) verifyMessage(
 	}
 
 	var verErr *header.VerifyError
-	err = s.verifier(ctx, hdr)
-	switch {
+	switch err := s.verifier(ctx, hdr); {
 	case errors.As(err, &verErr) && verErr.SoftFailure:
 		s.metrics.ignore(ctx)
 		return pubsub.ValidationIgnore
 	case err != nil:
 		s.metrics.reject(ctx)
 		return pubsub.ValidationReject
 	default:
+		// keep the valid header in the msg so Subscriptions can access it without
+		// additional unmarshalling
+		msg.ValidatorData = hdr
+		s.metrics.accept(ctx, len(msg.Data))
+		return pubsub.ValidationAccept
+	}
+}
+
+func (s *Subscriber[H]) extractHeader(msg *pubsub.Message) (H, error) {
+	if msg.ValidatorData != nil {
+		hdr, ok := msg.ValidatorData.(H)
+		if !ok {
+			panic(fmt.Sprintf("msg ValidatorData is of type %T", msg.ValidatorData))
+		}
+		return hdr, nil
 	}
 
-	// keep the valid header in the msg so Subscriptions can access it without
-	// additional unmarshalling
-	msg.ValidatorData = hdr
-	s.metrics.accept(ctx, len(msg.Data))
-	return pubsub.ValidationAccept
+	hdr := header.New[H]()
+	if err := hdr.UnmarshalBinary(msg.Data); err != nil {
+		return hdr, err
+	}
+	if err := hdr.Validate(); err != nil {
+		return hdr, err
+	}
+	return hdr, nil
 }

sync/syncsub_test.go

@@ -0,0 +1,72 @@
+package sync
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	pubsub "github.com/libp2p/go-libp2p-pubsub"
+	mocknet "github.com/libp2p/go-libp2p/p2p/net/mock"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/celestiaorg/go-header/headertest"
+	"github.com/celestiaorg/go-header/p2p"
+)
+
+func TestSyncerWithSubscriber(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	t.Cleanup(cancel)
+
+	suite := headertest.NewTestSuite(t)
+
+	netw, err := mocknet.FullMeshLinked(1)
+	require.NoError(t, err)
+
+	gossipSub, err := pubsub.NewGossipSub(
+		ctx,
+		netw.Hosts()[0],
+		pubsub.WithMessageSignaturePolicy(pubsub.StrictNoSign),
+	)
+	require.NoError(t, err)
+
+	p2pSub, err := p2p.NewSubscriber[*headertest.DummyHeader](
+		gossipSub,
+		pubsub.DefaultMsgIdFn,
+	)
+	require.NoError(t, err)
+	err = p2pSub.Start(context.Background())
+	require.NoError(t, err)
+
+	sub, err := p2pSub.Subscribe()
+	require.NoError(t, err)
+
+	head := suite.Head()
+	syncer, err := NewSyncer(
+		newTestStore(t, ctx, head),
+		newTestStore(t, ctx, head),
+		p2pSub,
+	)
+	require.NoError(t, err)
+	err = syncer.Start(ctx)
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		err = syncer.Stop(ctx)
+		require.NoError(t, err)
+	})
+
+	expectedHeader := suite.GenDummyHeaders(1)[0]
+
+	err = p2pSub.Broadcast(ctx, expectedHeader)
+	require.NoError(t, err)
+
+	header, err := sub.NextHeader(ctx)
+	require.NoError(t, err)
+	assert.Equal(t, expectedHeader.Height(), header.Height())
+	assert.Equal(t, expectedHeader.Hash(), header.Hash())
+
+	state := syncer.State()
+	require.NoError(t, err)
+	assert.Equal(t, expectedHeader.Height(), state.Height)
+}