Improved checks for vote extension signing in the unit tests (#9895)

* Added checks to vote signing. Need to refactor makevote and signaddvote

* Merged a number of similar versions of `makeVote`

* refactor `signAddVote` so `MakeVote` can reuse the checks

blocksync/reactor_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/tendermint/tendermint/libs/log"
 	mpmocks "github.com/tendermint/tendermint/mempool/mocks"
 	"github.com/tendermint/tendermint/p2p"
+	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 	"github.com/tendermint/tendermint/proxy"
 	sm "github.com/tendermint/tendermint/state"
 	"github.com/tendermint/tendermint/store"
@@ -123,12 +124,20 @@ func newReactor(
 		blockID := types.BlockID{Hash: thisBlock.Hash(), PartSetHeader: thisParts.Header()}
 
 		// Simulate a commit for the current height
+		pubKey, err := privVals[0].GetPubKey()
+		if err != nil {
+			panic(err)
+		}
+		addr := pubKey.Address()
+		idx, _ := state.Validators.GetByAddress(addr)
 		vote, err := types.MakeVote(
-			thisBlock.Header.Height,
-			blockID,
-			state.Validators,
 			privVals[0],
 			thisBlock.Header.ChainID,
+			idx,
+			thisBlock.Header.Height,
+			0,
+			tmproto.PrecommitType,
+			blockID,
 			time.Now(),
 		)
 		if err != nil {

consensus/msgs_test.go

@@ -57,15 +57,17 @@ func TestMsgToProto(t *testing.T) {
 	}
 	pbProposal := proposal.ToProto()
 
-	pv := types.NewMockPV()
-	pk, err := pv.GetPubKey()
-	require.NoError(t, err)
-	val := types.NewValidator(pk, 100)
-
-	vote, err := types.MakeVote(
-		1, bi, &types.ValidatorSet{Proposer: val, Validators: []*types.Validator{val}},
-		pv, "chainID", time.Now())
-	require.NoError(t, err)
+	vote := types.MakeVoteNoError(
+		t,
+		types.NewMockPV(),
+		"chainID",
+		0,
+		1,
+		0,
+		tmproto.PrecommitType,
+		bi,
+		time.Now(),
+	)
 	pbVote := vote.ToProto()
 
 	testsCases := []struct {

consensus/types/height_vote_set_test.go

@@ -1,7 +1,6 @@
 package types
 
 import (
-	"fmt"
 	"os"
 	"testing"
 
@@ -56,33 +55,29 @@ func TestPeerCatchupRounds(t *testing.T) {
 
 }
 
-func makeVoteHR(t *testing.T, height int64, valIndex, round int32, privVals []types.PrivValidator) *types.Vote {
+func makeVoteHR(
+	t *testing.T,
+	height int64,
+	valIndex,
+	round int32,
+	privVals []types.PrivValidator,
+) *types.Vote {
 	privVal := privVals[valIndex]
-	pubKey, err := privVal.GetPubKey()
+	randBytes := tmrand.Bytes(tmhash.Size)
+
+	vote, err := types.MakeVote(
+		privVal,
+		test.DefaultTestChainID,
+		valIndex,
+		height,
+		round,
+		tmproto.PrecommitType,
+		types.BlockID{Hash: randBytes, PartSetHeader: types.PartSetHeader{}},
+		tmtime.Now(),
+	)
 	if err != nil {
 		panic(err)
 	}
 
-	randBytes := tmrand.Bytes(tmhash.Size)
-
-	vote := &types.Vote{
-		ValidatorAddress: pubKey.Address(),
-		ValidatorIndex:   valIndex,
-		Height:           height,
-		Round:            round,
-		Timestamp:        tmtime.Now(),
-		Type:             tmproto.PrecommitType,
-		BlockID:          types.BlockID{Hash: randBytes, PartSetHeader: types.PartSetHeader{}},
-	}
-
-	v := vote.ToProto()
-	err = privVal.SignVote(test.DefaultTestChainID, v)
-	if err != nil {
-		panic(fmt.Sprintf("Error signing vote: %v", err))
-	}
-
-	vote.Signature = v.Signature
-	vote.ExtensionSignature = v.ExtensionSignature
-
 	return vote
 }

evidence/verify_test.go

@@ -373,11 +373,12 @@ func TestVerifyDuplicateVoteEvidence(t *testing.T) {
 
 	const chainID = "mychain"
 
-	vote1 := makeVote(t, val, chainID, 0, 10, 2, 1, blockID, defaultEvidenceTime)
+	vote1 := types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID, defaultEvidenceTime)
+
 	v1 := vote1.ToProto()
 	err := val.SignVote(chainID, v1)
 	require.NoError(t, err)
-	badVote := makeVote(t, val, chainID, 0, 10, 2, 1, blockID, defaultEvidenceTime)
+	badVote := types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID, defaultEvidenceTime)
 	bv := badVote.ToProto()
 	err = val2.SignVote(chainID, bv)
 	require.NoError(t, err)
@@ -386,17 +387,17 @@ func TestVerifyDuplicateVoteEvidence(t *testing.T) {
 	badVote.Signature = bv.Signature
 
 	cases := []voteData{
-		{vote1, makeVote(t, val, chainID, 0, 10, 2, 1, blockID2, defaultEvidenceTime), true}, // different block ids
-		{vote1, makeVote(t, val, chainID, 0, 10, 2, 1, blockID3, defaultEvidenceTime), true},
-		{vote1, makeVote(t, val, chainID, 0, 10, 2, 1, blockID4, defaultEvidenceTime), true},
-		{vote1, makeVote(t, val, chainID, 0, 10, 2, 1, blockID, defaultEvidenceTime), false},     // wrong block id
-		{vote1, makeVote(t, val, "mychain2", 0, 10, 2, 1, blockID2, defaultEvidenceTime), false}, // wrong chain id
-		{vote1, makeVote(t, val, chainID, 0, 11, 2, 1, blockID2, defaultEvidenceTime), false},    // wrong height
-		{vote1, makeVote(t, val, chainID, 0, 10, 3, 1, blockID2, defaultEvidenceTime), false},    // wrong round
-		{vote1, makeVote(t, val, chainID, 0, 10, 2, 2, blockID2, defaultEvidenceTime), false},    // wrong step
-		{vote1, makeVote(t, val2, chainID, 0, 10, 2, 1, blockID2, defaultEvidenceTime), false},   // wrong validator
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID2, defaultEvidenceTime), true}, // different block ids
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID3, defaultEvidenceTime), true},
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID4, defaultEvidenceTime), true},
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID, defaultEvidenceTime), false},     // wrong block id
+		{vote1, types.MakeVoteNoError(t, val, "mychain2", 0, 10, 2, 1, blockID2, defaultEvidenceTime), false}, // wrong chain id
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 11, 2, 1, blockID2, defaultEvidenceTime), false},    // wrong height
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 3, 1, blockID2, defaultEvidenceTime), false},    // wrong round
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 2, blockID2, defaultEvidenceTime), false},    // wrong step
+		{vote1, types.MakeVoteNoError(t, val2, chainID, 0, 10, 2, 1, blockID2, defaultEvidenceTime), false},   // wrong validator
 		// a different vote time doesn't matter
-		{vote1, makeVote(t, val, chainID, 0, 10, 2, 1, blockID2, time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)), true},
+		{vote1, types.MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID2, time.Date(2020, 1, 1, 0, 0, 0, 0, time.UTC)), true},
 		{vote1, badVote, false}, // signed by wrong key
 	}
 
@@ -537,30 +538,6 @@ func makeLunaticEvidence(
 
 // }
 
-func makeVote(
-	t *testing.T, val types.PrivValidator, chainID string, valIndex int32, height int64,
-	round int32, step int, blockID types.BlockID, time time.Time) *types.Vote {
-	pubKey, err := val.GetPubKey()
-	require.NoError(t, err)
-	v := &types.Vote{
-		ValidatorAddress: pubKey.Address(),
-		ValidatorIndex:   valIndex,
-		Height:           height,
-		Round:            round,
-		Type:             tmproto.SignedMsgType(step),
-		BlockID:          blockID,
-		Timestamp:        time,
-	}
-
-	vpb := v.ToProto()
-	err = val.SignVote(chainID, vpb)
-	if err != nil {
-		panic(err)
-	}
-	v.Signature = vpb.Signature
-	return v
-}
-
 func makeHeaderRandom(height int64) *types.Header {
 	return &types.Header{
 		Version:            tmversion.Consensus{Block: version.BlockProtocol, App: 1},

internal/test/vote.go

@@ -1,43 +0,0 @@
-package test
-
-import (
-	"time"
-
-	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
-	"github.com/tendermint/tendermint/types"
-)
-
-func MakeVote(
-	val types.PrivValidator,
-	chainID string,
-	valIndex int32,
-	height int64,
-	round int32,
-	step int,
-	blockID types.BlockID,
-	time time.Time,
-) (*types.Vote, error) {
-	pubKey, err := val.GetPubKey()
-	if err != nil {
-		return nil, err
-	}
-
-	v := &types.Vote{
-		ValidatorAddress: pubKey.Address(),
-		ValidatorIndex:   valIndex,
-		Height:           height,
-		Round:            round,
-		Type:             tmproto.SignedMsgType(step),
-		BlockID:          blockID,
-		Timestamp:        time,
-	}
-
-	vpb := v.ToProto()
-	if err := val.SignVote(chainID, vpb); err != nil {
-		return nil, err
-	}
-
-	v.Signature = vpb.Signature
-	v.ExtensionSignature = vpb.ExtensionSignature
-	return v, nil
-}

light/helpers_test.go

@@ -120,9 +120,15 @@ func makeVote(header *types.Header, valset *types.ValidatorSet,
 	if err != nil {
 		panic(err)
 	}
-
 	vote.Signature = sig
 
+	extSignBytes := types.VoteExtensionSignBytes(header.ChainID, v)
+	extSig, err := key.Sign(extSignBytes)
+	if err != nil {
+		panic(err)
+	}
+	vote.ExtensionSignature = extSig
+
 	return vote
 }
 

state/execution_test.go

@@ -394,8 +394,7 @@ func TestProcessProposal(t *testing.T) {
 		pk, err := privVal.GetPubKey()
 		require.NoError(t, err)
 		idx, _ := state.Validators.GetByAddress(pk.Address())
-		vote, err := test.MakeVote(privVal, block0.Header.ChainID, idx, height-1, 0, 2, blockID, time.Now())
-		require.NoError(t, err)
+		vote := types.MakeVoteNoError(t, privVal, block0.Header.ChainID, idx, height-1, 0, 2, blockID, time.Now())
 		addr := pk.Address()
 		voteInfos = append(voteInfos,
 			abci.VoteInfo{

state/helpers_test.go

@@ -93,7 +93,16 @@ func makeValidCommit(
 	votes := make([]*types.Vote, vals.Size())
 	for i := 0; i < vals.Size(); i++ {
 		_, val := vals.GetByIndex(int32(i))
-		vote, err := types.MakeVote(height, blockID, vals, privVals[val.Address.String()], chainID, time.Now())
+		vote, err := types.MakeVote(
+			privVals[val.Address.String()],
+			chainID,
+			int32(i),
+			height,
+			0,
+			tmproto.PrecommitType,
+			blockID,
+			time.Now(),
+		)
 		if err != nil {
 			return nil, nil, err
 		}

state/validation_test.go

@@ -160,7 +160,8 @@ func TestValidateBlockCommit(t *testing.T) {
 			*/
 			// should be height-1 instead of height
 			idx, _ := state.Validators.GetByAddress(proposerAddr)
-			wrongHeightVote, err := test.MakeVote(
+			wrongHeightVote := types.MakeVoteNoError(
+				t,
 				privVals[proposerAddr.String()],
 				chainID,
 				idx,
@@ -170,7 +171,6 @@ func TestValidateBlockCommit(t *testing.T) {
 				state.LastBlockID,
 				time.Now(),
 			)
-			require.NoError(t, err, "height %d", height)
 			wrongHeightCommit := &types.Commit{
 				Height:     wrongHeightVote.Height,
 				Round:      wrongHeightVote.Round,
@@ -178,7 +178,7 @@ func TestValidateBlockCommit(t *testing.T) {
 				Signatures: []types.CommitSig{wrongHeightVote.CommitSig()},
 			}
 			block := makeBlock(state, height, wrongHeightCommit)
-			err = blockExec.ValidateBlock(state, block)
+			err := blockExec.ValidateBlock(state, block)
 			_, isErrInvalidCommitHeight := err.(types.ErrInvalidCommitHeight)
 			require.True(t, isErrInvalidCommitHeight, "expected ErrInvalidCommitHeight at height %d but got: %v", height, err)
 
@@ -215,14 +215,18 @@ func TestValidateBlockCommit(t *testing.T) {
 		/*
 			wrongSigsCommit is fine except for the extra bad precommit
 		*/
-		goodVote, err := types.MakeVote(height,
-			blockID,
-			state.Validators,
+		idx, _ := state.Validators.GetByAddress(proposerAddr)
+		goodVote := types.MakeVoteNoError(
+			t,
 			privVals[proposerAddr.String()],
 			chainID,
+			idx,
+			height,
+			0,
+			tmproto.PrecommitType,
+			blockID,
 			time.Now(),
 		)
-		require.NoError(t, err, "height %d", height)
 
 		bpvPubKey, err := badPrivVal.GetPubKey()
 		require.NoError(t, err)

test/e2e/runner/evidence.go

@@ -210,11 +210,11 @@ func generateDuplicateVoteEvidence(
 	if err != nil {
 		return nil, err
 	}
-	voteA, err := test.MakeVote(privVal, chainID, valIdx, height, 0, 2, makeRandomBlockID(), time)
+	voteA, err := types.MakeVote(privVal, chainID, valIdx, height, 0, 2, makeRandomBlockID(), time)
 	if err != nil {
 		return nil, err
 	}
-	voteB, err := test.MakeVote(privVal, chainID, valIdx, height, 0, 2, makeRandomBlockID(), time)
+	voteB, err := types.MakeVote(privVal, chainID, valIdx, height, 0, 2, makeRandomBlockID(), time)
 	if err != nil {
 		return nil, err
 	}

types/evidence_test.go

@@ -33,8 +33,8 @@ func randomDuplicateVoteEvidence(t *testing.T) *DuplicateVoteEvidence {
 	blockID2 := makeBlockID([]byte("blockhash2"), 1000, []byte("partshash"))
 	const chainID = "mychain"
 	return &DuplicateVoteEvidence{
-		VoteA:            makeVote(t, val, chainID, 0, 10, 2, 1, blockID, defaultVoteTime),
-		VoteB:            makeVote(t, val, chainID, 0, 10, 2, 1, blockID2, defaultVoteTime.Add(1*time.Minute)),
+		VoteA:            MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID, defaultVoteTime),
+		VoteB:            MakeVoteNoError(t, val, chainID, 0, 10, 2, 1, blockID2, defaultVoteTime.Add(1*time.Minute)),
 		TotalVotingPower: 30,
 		ValidatorPower:   10,
 		Timestamp:        defaultVoteTime,
@@ -69,7 +69,7 @@ func TestDuplicateVoteEvidenceValidation(t *testing.T) {
 			ev.VoteB = nil
 		}, true},
 		{"Invalid vote type", func(ev *DuplicateVoteEvidence) {
-			ev.VoteA = makeVote(t, val, chainID, math.MaxInt32, math.MaxInt64, math.MaxInt32, 0, blockID2, defaultVoteTime)
+			ev.VoteA = MakeVoteNoError(t, val, chainID, math.MaxInt32, math.MaxInt64, math.MaxInt32, 0, blockID2, defaultVoteTime)
 		}, true},
 		{"Invalid vote order", func(ev *DuplicateVoteEvidence) {
 			swap := ev.VoteA.Copy()
@@ -80,8 +80,8 @@ func TestDuplicateVoteEvidenceValidation(t *testing.T) {
 	for _, tc := range testCases {
 		tc := tc
 		t.Run(tc.testName, func(t *testing.T) {
-			vote1 := makeVote(t, val, chainID, math.MaxInt32, math.MaxInt64, math.MaxInt32, 0x02, blockID, defaultVoteTime)
-			vote2 := makeVote(t, val, chainID, math.MaxInt32, math.MaxInt64, math.MaxInt32, 0x02, blockID2, defaultVoteTime)
+			vote1 := MakeVoteNoError(t, val, chainID, math.MaxInt32, math.MaxInt64, math.MaxInt32, 0x02, blockID, defaultVoteTime)
+			vote2 := MakeVoteNoError(t, val, chainID, math.MaxInt32, math.MaxInt64, math.MaxInt32, 0x02, blockID2, defaultVoteTime)
 			valSet := NewValidatorSet([]*Validator{val.ExtractIntoValidator(10)})
 			ev, err := NewDuplicateVoteEvidence(vote1, vote2, defaultVoteTime, valSet)
 			require.NoError(t, err)
@@ -235,30 +235,6 @@ func TestMockEvidenceValidateBasic(t *testing.T) {
 	assert.Nil(t, goodEvidence.ValidateBasic())
 }
 
-func makeVote(
-	t *testing.T, val PrivValidator, chainID string, valIndex int32, height int64, round int32, step int, blockID BlockID,
-	time time.Time) *Vote {
-	pubKey, err := val.GetPubKey()
-	require.NoError(t, err)
-	v := &Vote{
-		ValidatorAddress: pubKey.Address(),
-		ValidatorIndex:   valIndex,
-		Height:           height,
-		Round:            round,
-		Type:             tmproto.SignedMsgType(step),
-		BlockID:          blockID,
-		Timestamp:        time,
-	}
-
-	vpb := v.ToProto()
-	err = val.SignVote(chainID, vpb)
-	require.NoError(t, err)
-
-	v.Signature = vpb.Signature
-	v.ExtensionSignature = vpb.ExtensionSignature
-	return v
-}
-
 func makeHeaderRandom() *Header {
 	return &Header{
 		Version:            tmversion.Consensus{Block: version.BlockProtocol, App: 1},
@@ -284,8 +260,8 @@ func TestEvidenceProto(t *testing.T) {
 	blockID := makeBlockID(tmhash.Sum([]byte("blockhash")), math.MaxInt32, tmhash.Sum([]byte("partshash")))
 	blockID2 := makeBlockID(tmhash.Sum([]byte("blockhash2")), math.MaxInt32, tmhash.Sum([]byte("partshash")))
 	const chainID = "mychain"
-	v := makeVote(t, val, chainID, math.MaxInt32, math.MaxInt64, 1, 0x01, blockID, defaultVoteTime)
-	v2 := makeVote(t, val, chainID, math.MaxInt32, math.MaxInt64, 2, 0x01, blockID2, defaultVoteTime)
+	v := MakeVoteNoError(t, val, chainID, math.MaxInt32, math.MaxInt64, 1, 0x01, blockID, defaultVoteTime)
+	v2 := MakeVoteNoError(t, val, chainID, math.MaxInt32, math.MaxInt64, 2, 0x01, blockID2, defaultVoteTime)
 
 	// -------- SignedHeaders --------
 	const height int64 = 37

types/test_util.go

@@ -2,8 +2,10 @@ package types
 
 import (
 	"fmt"
+	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
 	tmproto "github.com/tendermint/tendermint/proto/tendermint/types"
 	tmversion "github.com/tendermint/tendermint/proto/tendermint/version"
 	"github.com/tendermint/tendermint/version"
@@ -37,52 +39,101 @@ func MakeExtCommit(blockID BlockID, height int64, round int32,
 	return voteSet.MakeExtendedCommit(), nil
 }
 
-func signAddVote(privVal PrivValidator, vote *Vote, voteSet *VoteSet) (signed bool, err error) {
+func signAndCheckVote(
+	vote *Vote,
+	privVal PrivValidator,
+	chainID string,
+	extensionsEnabled bool,
+) error {
 	v := vote.ToProto()
-	err = privVal.SignVote(voteSet.ChainID(), v)
-	if err != nil {
-		return false, err
+	if err := privVal.SignVote(chainID, v); err != nil {
+		return err
 	}
 	vote.Signature = v.Signature
-	vote.ExtensionSignature = v.ExtensionSignature
+
+	isPrecommit := vote.Type == tmproto.PrecommitType
+	if !isPrecommit && extensionsEnabled {
+		return fmt.Errorf("only Precommit votes may have extensions enabled; vote type: %d", vote.Type)
+	}
+
+	isNil := vote.BlockID.IsZero()
+	extSignature := (len(v.ExtensionSignature) > 0)
+	if extSignature == (!isPrecommit || isNil) {
+		return fmt.Errorf(
+			"extensions must be present IFF vote is a non-nil Precommit; present %t, vote type %d, is nil %t",
+			extSignature,
+			vote.Type,
+			isNil,
+		)
+	}
+
+	vote.ExtensionSignature = nil
+	if extensionsEnabled {
+		vote.ExtensionSignature = v.ExtensionSignature
+	}
+
+	return nil
+}
+
+func signAddVote(privVal PrivValidator, vote *Vote, voteSet *VoteSet) (bool, error) {
+	if vote.Type != voteSet.signedMsgType {
+		return false, fmt.Errorf("vote and voteset are of different types; %d != %d", vote.Type, voteSet.signedMsgType)
+	}
+	if err := signAndCheckVote(vote, privVal, voteSet.ChainID(), voteSet.extensionsEnabled); err != nil {
+		return false, err
+	}
 	return voteSet.AddVote(vote)
 }
 
 func MakeVote(
-	height int64,
-	blockID BlockID,
-	valSet *ValidatorSet,
-	privVal PrivValidator,
+	val PrivValidator,
 	chainID string,
-	now time.Time,
+	valIndex int32,
+	height int64,
+	round int32,
+	step tmproto.SignedMsgType,
+	blockID BlockID,
+	time time.Time,
 ) (*Vote, error) {
-	pubKey, err := privVal.GetPubKey()
+	pubKey, err := val.GetPubKey()
 	if err != nil {
-		return nil, fmt.Errorf("can't get pubkey: %w", err)
-	}
-	addr := pubKey.Address()
-	idx, _ := valSet.GetByAddress(addr)
-	vote := &Vote{
-		ValidatorAddress: addr,
-		ValidatorIndex:   idx,
-		Height:           height,
-		Round:            0,
-		Timestamp:        now,
-		Type:             tmproto.PrecommitType,
-		BlockID:          blockID,
-	}
-	v := vote.ToProto()
-
-	if err := privVal.SignVote(chainID, v); err != nil {
 		return nil, err
 	}
 
-	vote.Signature = v.Signature
-	vote.ExtensionSignature = v.ExtensionSignature
+	vote := &Vote{
+		ValidatorAddress: pubKey.Address(),
+		ValidatorIndex:   valIndex,
+		Height:           height,
+		Round:            round,
+		Type:             step,
+		BlockID:          blockID,
+		Timestamp:        time,
+	}
+
+	extensionsEnabled := step == tmproto.PrecommitType
+	if err := signAndCheckVote(vote, val, chainID, extensionsEnabled); err != nil {
+		return nil, err
+	}
 
 	return vote, nil
 }
 
+func MakeVoteNoError(
+	t *testing.T,
+	val PrivValidator,
+	chainID string,
+	valIndex int32,
+	height int64,
+	round int32,
+	step tmproto.SignedMsgType,
+	blockID BlockID,
+	time time.Time,
+) *Vote {
+	vote, err := MakeVote(val, chainID, valIndex, height, round, step, blockID, time)
+	require.NoError(t, err)
+	return vote
+}
+
 // MakeBlock returns a new block with an empty header, except what can be
 // computed from itself.
 // It populates the same set of fields validated by ValidateBasic.

types/vote_set_test.go

@@ -567,7 +567,10 @@ func randVoteSet(
 	votingPower int64,
 ) (*VoteSet, *ValidatorSet, []PrivValidator) {
 	valSet, privValidators := RandValidatorSet(numValidators, votingPower)
-	return NewExtendedVoteSet("test_chain_id", height, round, signedMsgType, valSet), valSet, privValidators
+	if signedMsgType == tmproto.PrecommitType {
+		return NewExtendedVoteSet("test_chain_id", height, round, signedMsgType, valSet), valSet, privValidators
+	}
+	return NewVoteSet("test_chain_id", height, round, signedMsgType, valSet), valSet, privValidators
 }
 
 // Convenience: Return new vote with different validator address/index