Implement SDK handler to drop messages from non-validators

network/p2p/handler.go

@@ -5,6 +5,7 @@ package p2p
 
 import (
 	"context"
+	"errors"
 	"time"
 
 	"go.uber.org/zap"
@@ -15,7 +16,12 @@ import (
 	"github.com/ava-labs/avalanchego/utils/logging"
 )
 
-var _ Handler = (*NoOpHandler)(nil)
+var (
+	ErrNotValidator = errors.New("not a validator")
+
+	_ Handler = (*NoOpHandler)(nil)
+	_ Handler = (*ValidatorHandler)(nil)
+)
 
 // Handler is the server-side logic for virtual machine application protocols.
 type Handler interface {
@@ -58,6 +64,28 @@ func (NoOpHandler) CrossChainAppRequest(context.Context, ids.ID, time.Time, []by
 	return nil, nil
 }
 
+// ValidatorHandler drops messages from non-validators
+type ValidatorHandler struct {
+	Handler
+	ValidatorSet ValidatorSet
+}
+
+func (v ValidatorHandler) AppGossip(ctx context.Context, nodeID ids.NodeID, gossipBytes []byte) error {
+	if !v.ValidatorSet.Has(ctx, nodeID) {
+		return ErrNotValidator
+	}
+
+	return v.Handler.AppGossip(ctx, nodeID, gossipBytes)
+}
+
+func (v ValidatorHandler) AppRequest(ctx context.Context, nodeID ids.NodeID, deadline time.Time, requestBytes []byte) ([]byte, error) {
+	if !v.ValidatorSet.Has(ctx, nodeID) {
+		return nil, ErrNotValidator
+	}
+
+	return v.Handler.AppRequest(ctx, nodeID, deadline, requestBytes)
+}
+
 // responder automatically sends the response for a given request
 type responder struct {
 	handlerID uint64

network/p2p/handler_test.go

@@ -0,0 +1,105 @@
+// Copyright (C) 2019-2023, Ava Labs, Inc. All rights reserved.
+// See the file LICENSE for licensing terms.
+
+package p2p
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/ava-labs/avalanchego/ids"
+	"github.com/ava-labs/avalanchego/utils/set"
+)
+
+var _ ValidatorSet = (*testValidatorSet)(nil)
+
+type testValidatorSet struct {
+	validators set.Set[ids.NodeID]
+}
+
+func (t testValidatorSet) Has(_ context.Context, nodeID ids.NodeID) bool {
+	return t.validators.Contains(nodeID)
+}
+
+func TestValidatorHandlerAppGossip(t *testing.T) {
+	nodeID := ids.GenerateTestNodeID()
+	validatorSet := set.Of(nodeID)
+
+	tests := []struct {
+		name         string
+		validatorSet ValidatorSet
+		nodeID       ids.NodeID
+		expected     error
+	}{
+		{
+			name:         "message dropped",
+			validatorSet: testValidatorSet{},
+			nodeID:       nodeID,
+			expected:     ErrNotValidator,
+		},
+		{
+			name: "message handled",
+			validatorSet: testValidatorSet{
+				validators: validatorSet,
+			},
+			nodeID: nodeID,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			handler := ValidatorHandler{
+				Handler:      NoOpHandler{},
+				ValidatorSet: tt.validatorSet,
+			}
+
+			err := handler.AppGossip(context.Background(), tt.nodeID, []byte("foobar"))
+			require.ErrorIs(err, tt.expected)
+		})
+	}
+}
+
+func TestValidatorHandlerAppRequest(t *testing.T) {
+	nodeID := ids.GenerateTestNodeID()
+	validatorSet := set.Of(nodeID)
+
+	tests := []struct {
+		name         string
+		validatorSet ValidatorSet
+		nodeID       ids.NodeID
+		expected     error
+	}{
+		{
+			name:         "message dropped",
+			validatorSet: testValidatorSet{},
+			nodeID:       nodeID,
+			expected:     ErrNotValidator,
+		},
+		{
+			name: "message handled",
+			validatorSet: testValidatorSet{
+				validators: validatorSet,
+			},
+			nodeID: nodeID,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require := require.New(t)
+
+			handler := ValidatorHandler{
+				Handler:      NoOpHandler{},
+				ValidatorSet: tt.validatorSet,
+			}
+
+			_, err := handler.AppRequest(context.Background(), tt.nodeID, time.Time{}, []byte("foobar"))
+			require.ErrorIs(err, tt.expected)
+		})
+	}
+}

network/p2p/validators.go

@@ -16,7 +16,14 @@ import (
 	"github.com/ava-labs/avalanchego/utils/set"
 )
 
-var _ NodeSampler = (*Validators)(nil)
+var (
+	_ ValidatorSet = (*Validators)(nil)
+	_ NodeSampler  = (*Validators)(nil)
+)
+
+type ValidatorSet interface {
+	Has(ctx context.Context, nodeID ids.NodeID) bool
+}
 
 func NewValidators(log logging.Logger, subnetID ids.ID, validators validators.State, maxValidatorSetStaleness time.Duration) *Validators {
 	return &Validators{