From 71d36953c326f7057cb5102889f15739507daaea Mon Sep 17 00:00:00 2001
From: Thane Thomson <connect@thanethomson.com>
Date: Sun, 20 Mar 2022 13:29:29 -0400
Subject: [PATCH] Thread vote extensions through code and fix tests

Signed-off-by: Thane Thomson <connect@thanethomson.com>
---
 internal/consensus/byzantine_test.go |  2 +-
 internal/consensus/common_test.go    | 25 +++++++++++++------------
 internal/consensus/msgs_test.go      |  2 +-
 internal/consensus/state.go          |  5 ++++-
 internal/state/execution.go          |  8 +++++++-
 internal/state/execution_test.go     | 20 ++++++++++----------
 internal/state/helpers_test.go       | 12 +++++++-----
 node/node_test.go                    |  5 ++++-
 privval/file.go                      | 11 ++++++++++-
 privval/msgs_test.go                 |  4 ++--
 types/priv_validator.go              |  6 ++++++
 types/vote.go                        | 27 ++++++++++++++++++---------
 12 files changed, 83 insertions(+), 44 deletions(-)

diff --git a/internal/consensus/byzantine_test.go b/internal/consensus/byzantine_test.go
index dfeb556fe..053f1603d 100644
--- a/internal/consensus/byzantine_test.go
+++ b/internal/consensus/byzantine_test.go
@@ -201,7 +201,7 @@ func TestByzantinePrevoteEquivocation(t *testing.T) {
 		proposerAddr := lazyNodeState.privValidatorPubKey.Address()
 
 		block, err := lazyNodeState.blockExec.CreateProposalBlock(
-			ctx, lazyNodeState.Height, lazyNodeState.state, commit, proposerAddr, nil)
+			ctx, lazyNodeState.Height, lazyNodeState.state, commit, proposerAddr, lazyNodeState.LastCommit.GetVotes())
 		require.NoError(t, err)
 		blockParts, err := block.MakePartSet(types.BlockPartSizeBytes)
 		require.NoError(t, err)
diff --git a/internal/consensus/common_test.go b/internal/consensus/common_test.go
index 7498814f7..577315207 100644
--- a/internal/consensus/common_test.go
+++ b/internal/consensus/common_test.go
@@ -120,19 +120,17 @@ func (vs *validatorStub) signVote(
 	}
 
 	vote := &types.Vote{
-		Type:               voteType,
-		Height:             vs.Height,
-		Round:              vs.Round,
-		BlockID:            blockID,
-		Timestamp:          vs.clock.Now(),
-		ValidatorAddress:   pubKey.Address(),
-		ValidatorIndex:     vs.Index,
-		Signature:          []byte{},
-		Extension:          []byte{},
-		ExtensionSignature: []byte{},
+		Type:             voteType,
+		Height:           vs.Height,
+		Round:            vs.Round,
+		BlockID:          blockID,
+		Timestamp:        vs.clock.Now(),
+		ValidatorAddress: pubKey.Address(),
+		ValidatorIndex:   vs.Index,
+		Extension:        []byte("extension"),
 	}
 	v := vote.ToProto()
-	if err := vs.PrivValidator.SignVote(ctx, chainID, v); err != nil {
+	if err = vs.PrivValidator.SignVote(ctx, chainID, v); err != nil {
 		return nil, fmt.Errorf("sign vote failed: %w", err)
 	}
 
@@ -140,10 +138,12 @@ func (vs *validatorStub) signVote(
 	if signDataIsEqual(vs.lastVote, v) {
 		v.Signature = vs.lastVote.Signature
 		v.Timestamp = vs.lastVote.Timestamp
+		v.ExtensionSignature = vs.lastVote.ExtensionSignature
 	}
 
 	vote.Signature = v.Signature
 	vote.Timestamp = v.Timestamp
+	vote.ExtensionSignature = v.ExtensionSignature
 
 	return vote, err
 }
@@ -984,5 +984,6 @@ func signDataIsEqual(v1 *types.Vote, v2 *tmproto.Vote) bool {
 		v1.Height == v2.GetHeight() &&
 		v1.Round == v2.Round &&
 		bytes.Equal(v1.ValidatorAddress.Bytes(), v2.GetValidatorAddress()) &&
-		v1.ValidatorIndex == v2.GetValidatorIndex()
+		v1.ValidatorIndex == v2.GetValidatorIndex() &&
+		bytes.Equal(v1.Extension, v2.Extension)
 }
diff --git a/internal/consensus/msgs_test.go b/internal/consensus/msgs_test.go
index fd26082c1..184ae2393 100644
--- a/internal/consensus/msgs_test.go
+++ b/internal/consensus/msgs_test.go
@@ -402,7 +402,7 @@ func TestConsMsgsVectors(t *testing.T) {
 			"2a36080110011a3008011204746573741a26080110011a206164645f6d6f72655f6578636c616d6174696f6e5f6d61726b735f636f64652d"},
 		{"Vote", &tmcons.Message{Sum: &tmcons.Message_Vote{
 			Vote: &tmcons.Vote{Vote: vpb}}},
-			"3280010a7e0802100122480a206164645f6d6f72655f6578636c616d6174696f6e5f6d61726b735f636f64652d1224080112206164645f6d6f72655f6578636c616d6174696f6e5f6d61726b735f636f64652d2a0608c0b89fdc0532146164645f6d6f72655f6578636c616d6174696f6e38014a0e0a067369676e6564120461757468"},
+			"32780a760802100122480a206164645f6d6f72655f6578636c616d6174696f6e5f6d61726b735f636f64652d1224080112206164645f6d6f72655f6578636c616d6174696f6e5f6d61726b735f636f64652d2a0608c0b89fdc0532146164645f6d6f72655f6578636c616d6174696f6e38014a067369676e6564"},
 		{"HasVote", &tmcons.Message{Sum: &tmcons.Message_HasVote{
 			HasVote: &tmcons.HasVote{Height: 1, Round: 1, Type: tmproto.PrevoteType, Index: 1}}},
 			"3a080801100118012001"},
diff --git a/internal/consensus/state.go b/internal/consensus/state.go
index d81e3e8de..1f25a525b 100644
--- a/internal/consensus/state.go
+++ b/internal/consensus/state.go
@@ -1375,6 +1375,7 @@ func (cs *State) createProposalBlock(ctx context.Context) (*types.Block, error)
 	}
 
 	var commit *types.Commit
+	var votes []*types.Vote
 	switch {
 	case cs.Height == cs.state.InitialHeight:
 		// We're creating a proposal for the first block.
@@ -1384,6 +1385,7 @@ func (cs *State) createProposalBlock(ctx context.Context) (*types.Block, error)
 	case cs.LastCommit.HasTwoThirdsMajority():
 		// Make the commit from LastCommit
 		commit = cs.LastCommit.MakeCommit()
+		votes = cs.LastCommit.GetVotes()
 
 	default: // This shouldn't happen.
 		cs.logger.Error("propose step; cannot propose anything without commit for the previous block")
@@ -1399,7 +1401,7 @@ func (cs *State) createProposalBlock(ctx context.Context) (*types.Block, error)
 
 	proposerAddr := cs.privValidatorPubKey.Address()
 
-	return cs.blockExec.CreateProposalBlock(ctx, cs.Height, cs.state, commit, proposerAddr, cs.LastCommit.GetVotes())
+	return cs.blockExec.CreateProposalBlock(ctx, cs.Height, cs.state, commit, proposerAddr, votes)
 }
 
 // Enter: `timeoutPropose` after entering Propose.
@@ -2480,6 +2482,7 @@ func (cs *State) signVote(
 
 	err := cs.privValidator.SignVote(ctxto, cs.state.ChainID, v)
 	vote.Signature = v.Signature
+	vote.ExtensionSignature = v.ExtensionSignature
 	vote.Timestamp = v.Timestamp
 
 	return vote, err
diff --git a/internal/state/execution.go b/internal/state/execution.go
index fd7ae3ba9..2664df9af 100644
--- a/internal/state/execution.go
+++ b/internal/state/execution.go
@@ -437,11 +437,17 @@ func buildLastCommitInfo(block *types.Block, store Store, initialHeight int64) a
 
 func extendedCommitInfo(c abci.CommitInfo, votes []*types.Vote) abci.ExtendedCommitInfo {
 	vs := make([]abci.ExtendedVoteInfo, len(c.Votes))
+	var ext []byte
 	for i := range vs {
+		if c.Votes[i].SignedLastBlock {
+			ext = votes[i].Extension
+		} else {
+			ext = nil
+		}
 		vs[i] = abci.ExtendedVoteInfo{
 			Validator:       c.Votes[i].Validator,
 			SignedLastBlock: c.Votes[i].SignedLastBlock,
-			VoteExtension:   votes[i].Extension,
+			VoteExtension:   ext,
 		}
 	}
 	return abci.ExtendedCommitInfo{
diff --git a/internal/state/execution_test.go b/internal/state/execution_test.go
index b5c4e0975..1e245d777 100644
--- a/internal/state/execution_test.go
+++ b/internal/state/execution_test.go
@@ -645,8 +645,8 @@ func TestEmptyPrepareProposal(t *testing.T) {
 		eventBus,
 	)
 	pa, _ := state.Validators.GetByIndex(0)
-	commit := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
-	_, err = blockExec.CreateProposalBlock(ctx, height, state, commit, pa, nil)
+	commit, votes := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
+	_, err = blockExec.CreateProposalBlock(ctx, height, state, commit, pa, votes)
 	require.NoError(t, err)
 }
 
@@ -756,8 +756,8 @@ func TestPrepareProposalRemoveTxs(t *testing.T) {
 		eventBus,
 	)
 	pa, _ := state.Validators.GetByIndex(0)
-	commit := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
-	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, nil)
+	commit, votes := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
+	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, votes)
 	require.NoError(t, err)
 	require.Len(t, block.Data.Txs.ToSliceOfBytes(), len(trs)-2)
 
@@ -817,8 +817,8 @@ func TestPrepareProposalAddedTxsIncluded(t *testing.T) {
 		eventBus,
 	)
 	pa, _ := state.Validators.GetByIndex(0)
-	commit := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
-	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, nil)
+	commit, votes := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
+	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, votes)
 	require.NoError(t, err)
 
 	require.Equal(t, txs[0], block.Data.Txs[0])
@@ -875,8 +875,8 @@ func TestPrepareProposalReorderTxs(t *testing.T) {
 		eventBus,
 	)
 	pa, _ := state.Validators.GetByIndex(0)
-	commit := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
-	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, nil)
+	commit, votes := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
+	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, votes)
 	require.NoError(t, err)
 	for i, tx := range block.Data.Txs {
 		require.Equal(t, types.Tx(trs[i].Tx), tx)
@@ -938,8 +938,8 @@ func TestPrepareProposalModifiedTxFalse(t *testing.T) {
 		eventBus,
 	)
 	pa, _ := state.Validators.GetByIndex(0)
-	commit := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
-	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, nil)
+	commit, votes := makeValidCommit(ctx, t, height, types.BlockID{}, state.Validators, privVals)
+	block, err := blockExec.CreateProposalBlock(ctx, height, state, commit, pa, votes)
 	require.NoError(t, err)
 	for i, tx := range block.Data.Txs {
 		require.Equal(t, txs[i], tx)
diff --git a/internal/state/helpers_test.go b/internal/state/helpers_test.go
index dffb6f256..cfe284898 100644
--- a/internal/state/helpers_test.go
+++ b/internal/state/helpers_test.go
@@ -46,7 +46,7 @@ func makeAndCommitGoodBlock(
 	state, blockID := makeAndApplyGoodBlock(ctx, t, state, height, lastCommit, proposerAddr, blockExec, evidence)
 
 	// Simulate a lastCommit for this block from all validators for the next height
-	commit := makeValidCommit(ctx, t, height, blockID, state.Validators, privVals)
+	commit, _ := makeValidCommit(ctx, t, height, blockID, state.Validators, privVals)
 
 	return state, blockID, commit
 }
@@ -82,17 +82,19 @@ func makeValidCommit(
 	blockID types.BlockID,
 	vals *types.ValidatorSet,
 	privVals map[string]types.PrivValidator,
-) *types.Commit {
+) (*types.Commit, []*types.Vote) {
 	t.Helper()
-	sigs := make([]types.CommitSig, 0)
+	sigs := make([]types.CommitSig, vals.Size())
+	votes := make([]*types.Vote, vals.Size())
 	for i := 0; i < vals.Size(); i++ {
 		_, val := vals.GetByIndex(int32(i))
 		vote, err := factory.MakeVote(ctx, privVals[val.Address.String()], chainID, int32(i), height, 0, 2, blockID, time.Now())
 		require.NoError(t, err)
-		sigs = append(sigs, vote.CommitSig())
+		sigs[i] = vote.CommitSig()
+		votes[i] = vote
 	}
 
-	return types.NewCommit(height, 0, blockID, sigs)
+	return types.NewCommit(height, 0, blockID, sigs), votes
 }
 
 func makeState(t *testing.T, nVals, height int) (sm.State, dbm.DB, map[string]types.PrivValidator) {
diff --git a/node/node_test.go b/node/node_test.go
index fcd633d71..16832e523 100644
--- a/node/node_test.go
+++ b/node/node_test.go
@@ -531,8 +531,11 @@ func TestMaxProposalBlockSize(t *testing.T) {
 		BlockID: blockID,
 	}
 
+	votes := make([]*types.Vote, types.MaxVotesCount)
+
 	// add maximum amount of signatures to a single commit
 	for i := 0; i < types.MaxVotesCount; i++ {
+		votes[i] = &types.Vote{}
 		commit.Signatures = append(commit.Signatures, cs)
 	}
 
@@ -541,7 +544,7 @@ func TestMaxProposalBlockSize(t *testing.T) {
 		math.MaxInt64,
 		state, commit,
 		proposerAddr,
-		nil,
+		votes,
 	)
 	require.NoError(t, err)
 	partSet, err := block.MakePartSet(types.BlockPartSizeBytes)
diff --git a/privval/file.go b/privval/file.go
index 728e0dc67..2f00c9907 100644
--- a/privval/file.go
+++ b/privval/file.go
@@ -398,10 +398,19 @@ func (pv *FilePV) signVote(chainID string, vote *tmproto.Vote) error {
 	if err != nil {
 		return err
 	}
-	if err := pv.saveSigned(height, round, step, signBytes, sig); err != nil {
+	if err = pv.saveSigned(height, round, step, signBytes, sig); err != nil {
 		return err
 	}
 	vote.Signature = sig
+
+	// Sign the vote extension, if any
+	if vote.Extension != nil {
+		vote.ExtensionSignature, err = pv.Key.PrivKey.Sign(vote.Extension)
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
diff --git a/privval/msgs_test.go b/privval/msgs_test.go
index 8e3d4b2a4..31f1f61a6 100644
--- a/privval/msgs_test.go
+++ b/privval/msgs_test.go
@@ -80,8 +80,8 @@ func TestPrivvalVectors(t *testing.T) {
 		{"pubKey request", &privproto.PubKeyRequest{}, "0a00"},
 		{"pubKey response", &privproto.PubKeyResponse{PubKey: ppk, Error: nil}, "12240a220a20556a436f1218d30942efe798420f51dc9b6a311b929c578257457d05c5fcf230"},
 		{"pubKey response with error", &privproto.PubKeyResponse{PubKey: cryptoproto.PublicKey{}, Error: remoteError}, "12140a0012100801120c697427732061206572726f72"},
-		{"Vote Request", &privproto.SignVoteRequest{Vote: votepb}, "1aa8010aa501080110031802224a0a208b01023386c371778ecb6368573e539afc3cc860ec3a2f614e54fe5652f4fc80122608c0843d122072db3d959635dff1bb567bedaa70573392c5159666a3f8caf11e413aac52207a2a0608f49a8ded0532146af1f4111082efb388211bc72c55bcd61e9ac3d538d5bb034a2f0a0f6170705f646174615f7369676e6564121c6170705f646174615f73656c665f61757468656e7469636174696e67"},
-		{"Vote Response", &privproto.SignedVoteResponse{Vote: *votepb, Error: nil}, "22a8010aa501080110031802224a0a208b01023386c371778ecb6368573e539afc3cc860ec3a2f614e54fe5652f4fc80122608c0843d122072db3d959635dff1bb567bedaa70573392c5159666a3f8caf11e413aac52207a2a0608f49a8ded0532146af1f4111082efb388211bc72c55bcd61e9ac3d538d5bb034a2f0a0f6170705f646174615f7369676e6564121c6170705f646174615f73656c665f61757468656e7469636174696e67"},
+		{"Vote Request", &privproto.SignVoteRequest{Vote: votepb}, "1a88010a8501080110031802224a0a208b01023386c371778ecb6368573e539afc3cc860ec3a2f614e54fe5652f4fc80122608c0843d122072db3d959635dff1bb567bedaa70573392c5159666a3f8caf11e413aac52207a2a0608f49a8ded0532146af1f4111082efb388211bc72c55bcd61e9ac3d538d5bb034a0f6170705f646174615f7369676e6564"},
+		{"Vote Response", &privproto.SignedVoteResponse{Vote: *votepb, Error: nil}, "2288010a8501080110031802224a0a208b01023386c371778ecb6368573e539afc3cc860ec3a2f614e54fe5652f4fc80122608c0843d122072db3d959635dff1bb567bedaa70573392c5159666a3f8caf11e413aac52207a2a0608f49a8ded0532146af1f4111082efb388211bc72c55bcd61e9ac3d538d5bb034a0f6170705f646174615f7369676e6564"},
 		{"Vote Response with error", &privproto.SignedVoteResponse{Vote: tmproto.Vote{}, Error: remoteError}, "22250a11220212002a0b088092b8c398feffffff0112100801120c697427732061206572726f72"},
 		{"Proposal Request", &privproto.SignProposalRequest{Proposal: proposalpb}, "2a700a6e08011003180220022a4a0a208b01023386c371778ecb6368573e539afc3cc860ec3a2f614e54fe5652f4fc80122608c0843d122072db3d959635dff1bb567bedaa70573392c5159666a3f8caf11e413aac52207a320608f49a8ded053a10697427732061207369676e6174757265"},
 		{"Proposal Response", &privproto.SignedProposalResponse{Proposal: *proposalpb, Error: nil}, "32700a6e08011003180220022a4a0a208b01023386c371778ecb6368573e539afc3cc860ec3a2f614e54fe5652f4fc80122608c0843d122072db3d959635dff1bb567bedaa70573392c5159666a3f8caf11e413aac52207a320608f49a8ded053a10697427732061207369676e6174757265"},
diff --git a/types/priv_validator.go b/types/priv_validator.go
index 5a9b27cb6..59dbaed7a 100644
--- a/types/priv_validator.go
+++ b/types/priv_validator.go
@@ -95,6 +95,12 @@ func (pv MockPV) SignVote(ctx context.Context, chainID string, vote *tmproto.Vot
 		return err
 	}
 	vote.Signature = sig
+	if vote.Extension != nil {
+		vote.ExtensionSignature, err = pv.PrivKey.Sign(vote.Extension)
+		if err != nil {
+			return err
+		}
+	}
 	return nil
 }
 
diff --git a/types/vote.go b/types/vote.go
index 8d8331613..ea205aa3a 100644
--- a/types/vote.go
+++ b/types/vote.go
@@ -157,6 +157,9 @@ func (vote *Vote) Verify(chainID string, pubKey crypto.PubKey) error {
 	if !pubKey.VerifySignature(VoteSignBytes(chainID, v), vote.Signature) {
 		return ErrVoteInvalidSignature
 	}
+	if vote.Extension != nil && !pubKey.VerifySignature(vote.Extension, vote.ExtensionSignature) {
+		return ErrVoteInvalidSignature
+	}
 	return nil
 }
 
@@ -203,7 +206,9 @@ func (vote *Vote) ValidateBasic() error {
 		return fmt.Errorf("signature is too big (max: %d)", MaxSignatureSize)
 	}
 
-	// XXX: add length verification for vote extension?
+	if len(vote.Extension) > 0 && len(vote.ExtensionSignature) == 0 {
+		return errors.New("vote extension signature is missing")
+	}
 
 	return nil
 }
@@ -216,14 +221,16 @@ func (vote *Vote) ToProto() *tmproto.Vote {
 	}
 
 	return &tmproto.Vote{
-		Type:             vote.Type,
-		Height:           vote.Height,
-		Round:            vote.Round,
-		BlockID:          vote.BlockID.ToProto(),
-		Timestamp:        vote.Timestamp,
-		ValidatorAddress: vote.ValidatorAddress,
-		ValidatorIndex:   vote.ValidatorIndex,
-		Signature:        vote.Signature,
+		Type:               vote.Type,
+		Height:             vote.Height,
+		Round:              vote.Round,
+		BlockID:            vote.BlockID.ToProto(),
+		Timestamp:          vote.Timestamp,
+		ValidatorAddress:   vote.ValidatorAddress,
+		ValidatorIndex:     vote.ValidatorIndex,
+		Signature:          vote.Signature,
+		Extension:          vote.Extension,
+		ExtensionSignature: vote.ExtensionSignature,
 	}
 }
 
@@ -264,6 +271,8 @@ func VoteFromProto(pv *tmproto.Vote) (*Vote, error) {
 	vote.ValidatorAddress = pv.ValidatorAddress
 	vote.ValidatorIndex = pv.ValidatorIndex
 	vote.Signature = pv.Signature
+	vote.Extension = pv.Extension
+	vote.ExtensionSignature = pv.ExtensionSignature
 
 	return vote, vote.ValidateBasic()
 }