fix to ProposalPOLMessage poisoning (CON-222)

sei-ibc-go/modules/core/02-client/keeper/client_test.go

@@ -463,7 +463,8 @@ func (suite *KeeperTestSuite) TestCheckMisbehaviourAndUpdateState() {
 	altValSet := tmtypes.NewValidatorSet([]*tmtypes.Validator{altVal})
 
 	// Create signer array and ensure it is in same order as bothValSet
-	_, suiteVal := suite.valSet.GetByIndex(0)
+	_, suiteVal, ok := suite.valSet.GetByIndex(0)
+	suite.Require().True(ok)
 	bothSigners := ibctesting.CreateSortedSignerArray(altPrivVal, suite.privVal, altVal, suiteVal)
 
 	altSigners := []tmtypes.PrivValidator{altPrivVal}

sei-ibc-go/modules/light-clients/07-tendermint/types/misbehaviour_handle_test.go

@@ -28,7 +28,8 @@ func (suite *TendermintTestSuite) TestCheckMisbehaviourAndUpdateState() {
 	// Create alternative validator set with only altVal
 	altValSet := tmtypes.NewValidatorSet([]*tmtypes.Validator{altVal})
 
-	_, suiteVal := suite.valSet.GetByIndex(0)
+	_, suiteVal, ok := suite.valSet.GetByIndex(0)
+	suite.Require().True(ok)
 
 	// Create signer array and ensure it is in same order as bothValSet
 	bothSigners := ibctesting.CreateSortedSignerArray(altPrivVal, suite.privVal, altVal, suiteVal)

sei-ibc-go/modules/light-clients/07-tendermint/types/misbehaviour_test.go

@@ -44,7 +44,8 @@ func (suite *TendermintTestSuite) TestMisbehaviourValidateBasic() {
 	signers := []tmtypes.PrivValidator{suite.privVal}
 
 	// Create signer array and ensure it is in same order as bothValSet
-	_, suiteVal := suite.valSet.GetByIndex(0)
+	_, suiteVal, ok := suite.valSet.GetByIndex(0)
+	suite.Require().True(ok)
 	bothSigners := ibctesting.CreateSortedSignerArray(altPrivVal, suite.privVal, altVal, suiteVal)
 
 	altSigners := []tmtypes.PrivValidator{altPrivVal}

sei-ibc-go/modules/light-clients/07-tendermint/types/tendermint_test.go

@@ -98,7 +98,8 @@ func getBothSigners(suite *TendermintTestSuite, altVal *tmtypes.Validator, altPr
 	// Create bothValSet with both suite validator and altVal. Would be valid update
 	bothValSet := tmtypes.NewValidatorSet(append(suite.valSet.Validators, altVal))
 	// Create signer array and ensure it is in same order as bothValSet
-	_, suiteVal := suite.valSet.GetByIndex(0)
+	_, suiteVal, ok := suite.valSet.GetByIndex(0)
+	suite.Require().True(ok)
 	bothSigners := ibctesting.CreateSortedSignerArray(altPrivVal, suite.privVal, altVal, suiteVal)
 	return bothValSet, bothSigners
 }

sei-ibc-go/modules/light-clients/07-tendermint/types/update_test.go

@@ -318,7 +318,8 @@ func (suite *TendermintTestSuite) TestCheckHeaderAndUpdateState() {
 			signers = []tmtypes.PrivValidator{suite.privVal}
 
 			// Create signer array and ensure it is in same order as bothValSet
-			_, suiteVal := suite.valSet.GetByIndex(0)
+			_, suiteVal, ok := suite.valSet.GetByIndex(0)
+			suite.Require().True(ok)
 			bothSigners = ibctesting.CreateSortedSignerArray(altPrivVal, suite.privVal, altVal, suiteVal)
 
 			consStateHeight = height // must be explicitly changed

sei-ibc-go/testing/chain.go

@@ -471,7 +471,7 @@ func (chain *TestChain) CreateTMClientHeader(chainID string, blockHeight int64,
 		ChainID:            chainID,
 		Height:             blockHeight,
 		Time:               timestamp,
-		LastBlockID:        MakeBlockID(make([]byte, tmhash.Size), 10_000, make([]byte, tmhash.Size)),
+		LastBlockID:        MakeBlockID(make([]byte, tmhash.Size), tmtypes.MaxBlockPartsCount, make([]byte, tmhash.Size)),
 		LastCommitHash:     chain.App.LastCommitID().Hash,
 		DataHash:           tmhash.Sum([]byte("data_hash")),
 		ValidatorsHash:     vsetHash,

sei-tendermint/internal/consensus/common_test.go

@@ -301,8 +301,8 @@ func validatePrevote(
 
 	address := pubKey.Address()
 
-	vote := prevotes.GetByAddress(address)
-	require.NotNil(t, vote, "Failed to find prevote from validator")
+	vote, ok := prevotes.GetByAddress(address)
+	require.True(t, ok, "Failed to find prevote from validator")
 
 	if blockHash == nil {
 		require.Nil(t, vote.BlockID.Hash, "Expected prevote to be for nil, got %X", vote.BlockID.Hash)
@@ -319,8 +319,8 @@ func validateLastPrecommit(ctx context.Context, t *testing.T, cs *State, privVal
 	require.NoError(t, err)
 	address := pv.Address()
 
-	vote := votes.GetByAddress(address)
-	require.NotNil(t, vote)
+	vote, ok := votes.GetByAddress(address)
+	require.True(t, ok)
 
 	require.True(t, bytes.Equal(vote.BlockID.Hash, blockHash),
 		"Expected precommit to be for %X, got %X", blockHash, vote.BlockID.Hash)
@@ -343,8 +343,8 @@ func validatePrecommit(
 	require.NoError(t, err)
 	address := pv.Address()
 
-	vote := precommits.GetByAddress(address)
-	require.NotNil(t, vote, "Failed to find precommit from validator")
+	vote, ok := precommits.GetByAddress(address)
+	require.True(t, ok, "Failed to find precommit from validator")
 
 	if votedBlockHash == nil {
 		require.Nil(t, vote.BlockID.Hash, "Expected precommit to be for nil")

sei-tendermint/internal/consensus/invalid_test.go

@@ -11,17 +11,147 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"github.com/sei-protocol/sei-chain/sei-tendermint/crypto"
+	cstypes "github.com/sei-protocol/sei-chain/sei-tendermint/internal/consensus/types"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/eventbus"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/internal/p2p"
+	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/bits"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/bytes"
 	tmrand "github.com/sei-protocol/sei-chain/sei-tendermint/libs/rand"
 	tmtime "github.com/sei-protocol/sei-chain/sei-tendermint/libs/time"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/libs/utils"
 	tmcons "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/consensus"
 	tmproto "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/types"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/types"
+	"github.com/sei-protocol/sei-chain/sei-tendermint/version"
 )
 
+// Test checking that if peer sends a ProposalPOLMessage with a bitarray with bad length,
+// the node will handle it gracefully.
+func TestGossipVotesForHeightPoisonedProposalPOL(t *testing.T) {
+	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
+	defer cancel()
+
+	cfg := configSetup(t)
+	states, cleanup := makeConsensusState(ctx, t, cfg, 2, "consensus_reactor_test", newMockTickerFunc(true))
+	t.Cleanup(cleanup)
+
+	rts := setup(ctx, t, 2, states, 1)
+
+	var nodeIDs []types.NodeID
+	for _, node := range rts.network.Nodes() {
+		nodeIDs = append(nodeIDs, node.NodeID)
+	}
+	require.Len(t, nodeIDs, 2)
+
+	reactor := rts.reactors[nodeIDs[0]]
+	peerID := nodeIDs[1]
+	state := reactor.state.GetState()
+	reactor.SwitchToConsensus(ctx, state, false)
+
+	require.Eventually(t, func() bool {
+		_, ok := reactor.GetPeerState(peerID)
+		return ok
+	}, time.Hour, 50*time.Millisecond)
+
+	valSet, privVals := types.RandValidatorSet(4, 1)
+	proposerPubKey, err := privVals[0].GetPubKey(ctx)
+	require.NoError(t, err)
+
+	proposal := types.NewProposal(
+		1,
+		1,
+		0,
+		types.BlockID{
+			Hash: crypto.CRandBytes(crypto.HashSize),
+			PartSetHeader: types.PartSetHeader{
+				Total: 1,
+				Hash:  crypto.CRandBytes(crypto.HashSize),
+			},
+		},
+		time.Now(),
+		nil,
+		types.Header{
+			Version:         version.Consensus{Block: version.BlockProtocol},
+			Height:          1,
+			ProposerAddress: proposerPubKey.Address(),
+		},
+		&types.Commit{},
+		nil,
+		proposerPubKey.Address(),
+	)
+	proposal.Signature = makeSig("invalid-signature")
+
+	require.NoError(t, reactor.handleStateMessage(p2p.RecvMsg[*tmcons.Message]{
+		From: peerID,
+		Message: MsgToProto(&NewRoundStepMessage{
+			HRS: cstypes.HRS{
+				Height: 1,
+				Round:  1,
+				Step:   cstypes.RoundStepPrevote,
+			},
+			SecondsSinceStartTime: 1,
+			LastCommitRound:       -1,
+		}),
+	}))
+
+	require.NoError(t, reactor.handleDataMessage(ctx, p2p.RecvMsg[*tmcons.Message]{
+		From: peerID,
+		Message: MsgToProto(&ProposalMessage{
+			Proposal: proposal,
+		}),
+	}))
+
+	require.NoError(t, reactor.handleDataMessage(ctx, p2p.RecvMsg[*tmcons.Message]{
+		From: peerID,
+		Message: MsgToProto(&ProposalPOLMessage{
+			Height:           1,
+			ProposalPOLRound: 0,
+			ProposalPOL:      bits.NewBitArray(1),
+		}),
+	}))
+
+	ps, ok := reactor.GetPeerState(peerID)
+	require.True(t, ok)
+	prs := ps.GetRoundState()
+	require.Equal(t, int64(1), prs.Height)
+	require.Equal(t, int32(1), prs.Round)
+	require.Equal(t, int32(0), prs.ProposalPOLRound)
+	require.Equal(t, 1, prs.ProposalPOL.Size())
+
+	voteSet := cstypes.NewHeightVoteSet("test-chain", 1, valSet)
+	voteSet.SetRound(1)
+
+	voter := newValidatorStub(privVals[1], 1)
+	voter.Height = 1
+	voter.Round = 0
+
+	vote := signVote(ctx, t, voter, tmproto.PrevoteType, "test-chain", types.BlockID{
+		Hash: crypto.CRandBytes(crypto.HashSize),
+		PartSetHeader: types.PartSetHeader{
+			Total: 1,
+			Hash:  crypto.CRandBytes(crypto.HashSize),
+		},
+	})
+	added, err := voteSet.AddVote(vote, "")
+	require.NoError(t, err)
+	require.True(t, added)
+
+	rs := &cstypes.RoundState{
+		HRS: cstypes.HRS{
+			Height: 1,
+			Round:  1,
+			Step:   cstypes.RoundStepPrevote,
+		},
+		Votes: voteSet,
+	}
+
+	// Gossip should take into consideration that PeerState might contain
+	// invalid length bitarrays.
+	for range 10 {
+		reactor.gossipVotesForHeight(rs, ps.GetRoundState(), ps)
+	}
+}
+
 func TestReactorInvalidPrecommit(t *testing.T) {
 	t.Skip("test doesn't check anything useful")
 	ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
@@ -128,7 +258,10 @@ func invalidDoPrevoteFunc(
 		require.NoError(t, err)
 
 		addr := pubKey.Address()
-		valIndex, _ := cs.roundState.Validators().GetByAddress(addr)
+		valIndex, _, ok := cs.roundState.Validators().GetByAddress(addr)
+		if !ok {
+			panic("missing validator")
+		}
 
 		// precommit a random block
 		blockHash := bytes.HexBytes(tmrand.Bytes(32))

sei-tendermint/internal/consensus/memory_limit_test.go

@@ -2,11 +2,9 @@ package consensus
 
 import (
 	"testing"
-	"time"
 
 	"github.com/sei-protocol/sei-chain/sei-tendermint/crypto"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/crypto/ed25519"
-	tmproto "github.com/sei-protocol/sei-chain/sei-tendermint/proto/tendermint/types"
 	"github.com/sei-protocol/sei-chain/sei-tendermint/types"
 	"github.com/stretchr/testify/require"
 )
@@ -32,40 +30,6 @@ func TestPeerStateMemoryLimits(t *testing.T) {
 		{"very_large_total", 4294967295, true},
 	}
 
-	// Test SetHasProposal memory limits
-	t.Run("SetHasProposal", func(t *testing.T) {
-		for _, tc := range testCases {
-			t.Run(tc.name, func(t *testing.T) {
-				ps := NewPeerState(peerID)
-				ps.PRS.Height = 1
-				ps.PRS.Round = 0
-				blockID := types.BlockID{
-					Hash: make([]byte, 32),
-					PartSetHeader: types.PartSetHeader{
-						Total: tc.total,
-						Hash:  make([]byte, 32),
-					},
-				}
-				// Create a minimal proposal with basic required fields
-				proposal := &types.Proposal{
-					Type:      tmproto.ProposalType,
-					Height:    1,
-					Round:     0,
-					POLRound:  -1,
-					BlockID:   blockID,
-					Timestamp: time.Now(),
-					Signature: makeSig("test-signature"),
-				}
-				ps.SetHasProposal(proposal)
-				if tc.expectError {
-					require.False(t, ps.PRS.Proposal, "Expected proposal to be silently ignored for excessive Total")
-				} else {
-					require.True(t, ps.PRS.Proposal, "Expected proposal to be accepted for valid Total")
-				}
-			})
-		}
-	})
-
 	// Test InitProposalBlockParts memory limits
 	t.Run("InitProposalBlockParts", func(t *testing.T) {
 		for _, tc := range testCases {

sei-tendermint/internal/consensus/msgs.go

@@ -131,7 +131,7 @@ func (m *ProposalMessage) String() string {
 	return fmt.Sprintf("[Proposal %v]", m.Proposal)
 }
 
-// ProposalPOLMessage is sent when a previous proposal is re-proposed.
+// ProposalPOLMessage is sent when a node needs POL round votes.
 type ProposalPOLMessage struct {
 	Height           int64
 	ProposalPOLRound int32

sei-tendermint/internal/consensus/msgs_test.go

@@ -65,8 +65,8 @@ func TestMsgToProto(t *testing.T) {
 
 	proposal := types.Proposal{
 		Type:            tmproto.ProposalType,
-		Height:          1,
-		Round:           1,
+		Height:          3,
+		Round:           2,
 		POLRound:        1,
 		BlockID:         bi,
 		Timestamp:       time.Now(),

sei-tendermint/internal/consensus/peer_state.go

@@ -93,12 +93,6 @@ func (ps *PeerState) SetHasProposal(proposal *types.Proposal) {
 		return
 	}
 
-	// Check memory limits before acquiring lock or setting any state
-	if proposal.BlockID.PartSetHeader.Total > types.MaxBlockPartsCount {
-		logger.Debug("PartSetHeader.Total exceeds maximum", "total", proposal.BlockID.PartSetHeader.Total, "max", types.MaxBlockPartsCount)
-		return
-	}
-
 	ps.mtx.Lock()
 	defer ps.mtx.Unlock()
 
@@ -187,12 +181,8 @@ func (ps *PeerState) PickVoteToSend(votes types.VoteSetReader) (*types.Vote, boo
 	}
 
 	if index, ok := votes.BitArray().Sub(psVotes).PickRandom(); ok {
-		vote := votes.GetByIndex(int32(index)) //nolint:gosec // index is bounded by validator set size which fits in int32
-		if vote != nil {
-			return vote, true
-		}
+		return votes.GetByIndex(int32(index)) //nolint:gosec // index is bounded by validator set size which fits in int32
 	}
-
 	return nil, false
 }