From 4270b69fb0d2f04bac26d92af02f4ee8cb342acf Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Thu, 27 Jul 2023 16:33:03 +0200
Subject: [PATCH 1/8] fix: ValidatorStore interface didn't match StakingKeeper

---
 baseapp/abci_test.go           |  2 +-
 baseapp/abci_utils.go          | 11 ++++++++---
 baseapp/testutil/mock/mocks.go |  8 +++++---
 3 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/baseapp/abci_test.go b/baseapp/abci_test.go
index 4e164b812be0..bac900ab4bc9 100644
--- a/baseapp/abci_test.go
+++ b/baseapp/abci_test.go
@@ -1739,7 +1739,7 @@ func TestABCI_PrepareProposal_VoteExtensions(t *testing.T) {
 
 	consAddr := sdk.ConsAddress(addr.String())
 	valStore.EXPECT().GetValidatorByConsAddr(gomock.Any(), consAddr.Bytes()).Return(val1, nil).AnyTimes()
-	valStore.EXPECT().TotalBondedTokens(gomock.Any()).Return(math.NewInt(1000)).AnyTimes()
+	valStore.EXPECT().TotalBondedTokens(gomock.Any()).Return(math.NewInt(1000), nil).AnyTimes()
 
 	// set up baseapp
 	prepareOpt := func(bapp *baseapp.BaseApp) {
diff --git a/baseapp/abci_utils.go b/baseapp/abci_utils.go
index 6dacc5407d8f..6490d51e8bde 100644
--- a/baseapp/abci_utils.go
+++ b/baseapp/abci_utils.go
@@ -2,6 +2,7 @@ package baseapp
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 
 	"github.com/cockroachdb/errors"
@@ -38,8 +39,8 @@ type (
 	// extension signatures. Typically, this will be implemented by the x/staking
 	// module, which has knowledge of the CometBFT public key.
 	ValidatorStore interface {
-		GetValidatorByConsAddr(sdk.Context, cryptotypes.Address) (Validator, error)
-		TotalBondedTokens(ctx sdk.Context) math.Int
+		GetValidatorByConsAddr(context.Context, cryptotypes.Address) (Validator, error)
+		TotalBondedTokens(ctx context.Context) (math.Int, error)
 	}
 
 	// GasTx defines the contract that a transaction with a gas limit must implement.
@@ -130,7 +131,11 @@ func ValidateVoteExtensions(
 
 	// Ensure we have at least 2/3 voting power that submitted valid vote
 	// extensions.
-	totalVP := valStore.TotalBondedTokens(ctx)
+	totalVP, err := valStore.TotalBondedTokens(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to get total bonded tokens: %w", err)
+	}
+
 	percentSubmitted := math.LegacyNewDecFromInt(sumVP).Quo(math.LegacyNewDecFromInt(totalVP))
 	if percentSubmitted.LT(VoteExtensionThreshold) {
 		return fmt.Errorf("insufficient cumulative voting power received to verify vote extensions; got: %s, expected: >=%s", percentSubmitted, VoteExtensionThreshold)
diff --git a/baseapp/testutil/mock/mocks.go b/baseapp/testutil/mock/mocks.go
index 437ed68f485b..864071d866c2 100644
--- a/baseapp/testutil/mock/mocks.go
+++ b/baseapp/testutil/mock/mocks.go
@@ -5,6 +5,7 @@
 package mock
 
 import (
+	context "context"
 	reflect "reflect"
 
 	math "cosmossdk.io/math"
@@ -91,7 +92,7 @@ func (m *MockValidatorStore) EXPECT() *MockValidatorStoreMockRecorder {
 }
 
 // GetValidatorByConsAddr mocks base method.
-func (m *MockValidatorStore) GetValidatorByConsAddr(arg0 types0.Context, arg1 types.Address) (baseapp.Validator, error) {
+func (m *MockValidatorStore) GetValidatorByConsAddr(arg0 context.Context, arg1 types.Address) (baseapp.Validator, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "GetValidatorByConsAddr", arg0, arg1)
 	ret0, _ := ret[0].(baseapp.Validator)
@@ -106,11 +107,12 @@ func (mr *MockValidatorStoreMockRecorder) GetValidatorByConsAddr(arg0, arg1 inte
 }
 
 // TotalBondedTokens mocks base method.
-func (m *MockValidatorStore) TotalBondedTokens(ctx types0.Context) math.Int {
+func (m *MockValidatorStore) TotalBondedTokens(ctx context.Context) (math.Int, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "TotalBondedTokens", ctx)
 	ret0, _ := ret[0].(math.Int)
-	return ret0
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
 }
 
 // TotalBondedTokens indicates an expected call of TotalBondedTokens.

From e0ef985c6f72966575d34756b7c5ff04d4061a6d Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Thu, 27 Jul 2023 21:27:34 +0200
Subject: [PATCH 2/8] add necessary functions to staking keeper + integration
 test

---
 baseapp/abci_utils.go                         | 34 ++-----
 baseapp/testutil/mock/mocks.go                | 91 +++++-------------
 .../staking/keeper/vote_extensions_test.go    | 96 +++++++++++++++++++
 x/staking/keeper/validator.go                 | 19 ++++
 x/staking/testutil/expected_keepers_mocks.go  | 31 ++++++
 x/staking/types/expected_keepers.go           |  7 ++
 6 files changed, 189 insertions(+), 89 deletions(-)
 create mode 100644 tests/integration/staking/keeper/vote_extensions_test.go

diff --git a/baseapp/abci_utils.go b/baseapp/abci_utils.go
index 6490d51e8bde..933f9db62798 100644
--- a/baseapp/abci_utils.go
+++ b/baseapp/abci_utils.go
@@ -7,7 +7,6 @@ import (
 
 	"github.com/cockroachdb/errors"
 	abci "github.com/cometbft/cometbft/abci/types"
-	cmtcrypto "github.com/cometbft/cometbft/crypto"
 	cryptoenc "github.com/cometbft/cometbft/crypto/encoding"
 	cmtprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
@@ -16,7 +15,6 @@ import (
 
 	"cosmossdk.io/math"
 
-	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/mempool"
 )
@@ -27,20 +25,13 @@ import (
 var VoteExtensionThreshold = math.LegacyNewDecWithPrec(667, 3)
 
 type (
-	// Validator defines the interface contract require for verifying vote extension
-	// signatures. Typically, this will be implemented by the x/staking module,
-	// which has knowledge of the CometBFT public key.
-	Validator interface {
-		CmtConsPublicKey() (cmtprotocrypto.PublicKey, error)
-		BondedTokens() math.Int
-	}
-
 	// ValidatorStore defines the interface contract require for verifying vote
 	// extension signatures. Typically, this will be implemented by the x/staking
 	// module, which has knowledge of the CometBFT public key.
 	ValidatorStore interface {
-		GetValidatorByConsAddr(context.Context, cryptotypes.Address) (Validator, error)
 		TotalBondedTokens(ctx context.Context) (math.Int, error)
+		CmtConsPublicKeyByConsAddr(context.Context, sdk.ConsAddress) (cmtprotocrypto.PublicKey, error)
+		BondedTokensByConsAddr(context.Context, sdk.ConsAddress) (math.Int, error)
 	}
 
 	// GasTx defines the contract that a transaction with a gas limit must implement.
@@ -63,7 +54,7 @@ func ValidateVoteExtensions(
 ) error {
 	cp := ctx.ConsensusParams()
 	extsEnabled := cp.Abci != nil && currentHeight >= cp.Abci.VoteExtensionsEnableHeight && cp.Abci.VoteExtensionsEnableHeight != 0
-
+	fmt.Println(currentHeight, cp)
 	marshalDelimitedFn := func(msg proto.Message) ([]byte, error) {
 		var buf bytes.Buffer
 		if err := protoio.NewDelimitedWriter(&buf).WriteMsg(msg); err != nil {
@@ -90,17 +81,8 @@ func ValidateVoteExtensions(
 			return fmt.Errorf("vote extensions enabled; received empty vote extension signature at height %d", currentHeight)
 		}
 
-		valConsAddr := cmtcrypto.Address(vote.Validator.Address)
-
-		validator, err := valStore.GetValidatorByConsAddr(ctx, valConsAddr)
-		if err != nil {
-			return fmt.Errorf("failed to get validator %X: %w", valConsAddr, err)
-		}
-		if validator == nil {
-			return fmt.Errorf("validator %X not found", valConsAddr)
-		}
-
-		cmtPubKeyProto, err := validator.CmtConsPublicKey()
+		valConsAddr := sdk.ConsAddress(vote.Validator.Address)
+		cmtPubKeyProto, err := valStore.CmtConsPublicKeyByConsAddr(ctx, valConsAddr)
 		if err != nil {
 			return fmt.Errorf("failed to get validator %X public key: %w", valConsAddr, err)
 		}
@@ -126,7 +108,11 @@ func ValidateVoteExtensions(
 			return fmt.Errorf("failed to verify validator %X vote extension signature", valConsAddr)
 		}
 
-		sumVP = sumVP.Add(validator.BondedTokens())
+		bondedTokens, err := valStore.BondedTokensByConsAddr(ctx, valConsAddr)
+		if err != nil {
+			return fmt.Errorf("failed to get validator %X bonded tokens: %w", valConsAddr, err)
+		}
+		sumVP = sumVP.Add(bondedTokens)
 	}
 
 	// Ensure we have at least 2/3 voting power that submitted valid vote
diff --git a/baseapp/testutil/mock/mocks.go b/baseapp/testutil/mock/mocks.go
index 864071d866c2..d2f34574f53c 100644
--- a/baseapp/testutil/mock/mocks.go
+++ b/baseapp/testutil/mock/mocks.go
@@ -10,64 +10,10 @@ import (
 
 	math "cosmossdk.io/math"
 	crypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
-	baseapp "github.com/cosmos/cosmos-sdk/baseapp"
-	types "github.com/cosmos/cosmos-sdk/crypto/types"
-	types0 "github.com/cosmos/cosmos-sdk/types"
+	types "github.com/cosmos/cosmos-sdk/types"
 	gomock "github.com/golang/mock/gomock"
 )
 
-// MockValidator is a mock of Validator interface.
-type MockValidator struct {
-	ctrl     *gomock.Controller
-	recorder *MockValidatorMockRecorder
-}
-
-// MockValidatorMockRecorder is the mock recorder for MockValidator.
-type MockValidatorMockRecorder struct {
-	mock *MockValidator
-}
-
-// NewMockValidator creates a new mock instance.
-func NewMockValidator(ctrl *gomock.Controller) *MockValidator {
-	mock := &MockValidator{ctrl: ctrl}
-	mock.recorder = &MockValidatorMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use.
-func (m *MockValidator) EXPECT() *MockValidatorMockRecorder {
-	return m.recorder
-}
-
-// BondedTokens mocks base method.
-func (m *MockValidator) BondedTokens() math.Int {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "BondedTokens")
-	ret0, _ := ret[0].(math.Int)
-	return ret0
-}
-
-// BondedTokens indicates an expected call of BondedTokens.
-func (mr *MockValidatorMockRecorder) BondedTokens() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokens", reflect.TypeOf((*MockValidator)(nil).BondedTokens))
-}
-
-// CmtConsPublicKey mocks base method.
-func (m *MockValidator) CmtConsPublicKey() (crypto.PublicKey, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "CmtConsPublicKey")
-	ret0, _ := ret[0].(crypto.PublicKey)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// CmtConsPublicKey indicates an expected call of CmtConsPublicKey.
-func (mr *MockValidatorMockRecorder) CmtConsPublicKey() *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CmtConsPublicKey", reflect.TypeOf((*MockValidator)(nil).CmtConsPublicKey))
-}
-
 // MockValidatorStore is a mock of ValidatorStore interface.
 type MockValidatorStore struct {
 	ctrl     *gomock.Controller
@@ -91,19 +37,34 @@ func (m *MockValidatorStore) EXPECT() *MockValidatorStoreMockRecorder {
 	return m.recorder
 }
 
-// GetValidatorByConsAddr mocks base method.
-func (m *MockValidatorStore) GetValidatorByConsAddr(arg0 context.Context, arg1 types.Address) (baseapp.Validator, error) {
+// BondedTokensByConsAddr mocks base method.
+func (m *MockValidatorStore) BondedTokensByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (math.Int, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "GetValidatorByConsAddr", arg0, arg1)
-	ret0, _ := ret[0].(baseapp.Validator)
+	ret := m.ctrl.Call(m, "BondedTokensByConsAddr", arg0, arg1)
+	ret0, _ := ret[0].(math.Int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// BondedTokensByConsAddr indicates an expected call of BondedTokensByConsAddr.
+func (mr *MockValidatorStoreMockRecorder) BondedTokensByConsAddr(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokensByConsAddr", reflect.TypeOf((*MockValidatorStore)(nil).BondedTokensByConsAddr), arg0, arg1)
+}
+
+// CmtConsPublicKeyByConsAddr mocks base method.
+func (m *MockValidatorStore) CmtConsPublicKeyByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (crypto.PublicKey, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CmtConsPublicKeyByConsAddr", arg0, arg1)
+	ret0, _ := ret[0].(crypto.PublicKey)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
 
-// GetValidatorByConsAddr indicates an expected call of GetValidatorByConsAddr.
-func (mr *MockValidatorStoreMockRecorder) GetValidatorByConsAddr(arg0, arg1 interface{}) *gomock.Call {
+// CmtConsPublicKeyByConsAddr indicates an expected call of CmtConsPublicKeyByConsAddr.
+func (mr *MockValidatorStoreMockRecorder) CmtConsPublicKeyByConsAddr(arg0, arg1 interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetValidatorByConsAddr", reflect.TypeOf((*MockValidatorStore)(nil).GetValidatorByConsAddr), arg0, arg1)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CmtConsPublicKeyByConsAddr", reflect.TypeOf((*MockValidatorStore)(nil).CmtConsPublicKeyByConsAddr), arg0, arg1)
 }
 
 // TotalBondedTokens mocks base method.
@@ -182,7 +143,7 @@ func (m *MockProposalTxVerifier) EXPECT() *MockProposalTxVerifierMockRecorder {
 }
 
 // PrepareProposalVerifyTx mocks base method.
-func (m *MockProposalTxVerifier) PrepareProposalVerifyTx(tx types0.Tx) ([]byte, error) {
+func (m *MockProposalTxVerifier) PrepareProposalVerifyTx(tx types.Tx) ([]byte, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "PrepareProposalVerifyTx", tx)
 	ret0, _ := ret[0].([]byte)
@@ -197,10 +158,10 @@ func (mr *MockProposalTxVerifierMockRecorder) PrepareProposalVerifyTx(tx interfa
 }
 
 // ProcessProposalVerifyTx mocks base method.
-func (m *MockProposalTxVerifier) ProcessProposalVerifyTx(txBz []byte) (types0.Tx, error) {
+func (m *MockProposalTxVerifier) ProcessProposalVerifyTx(txBz []byte) (types.Tx, error) {
 	m.ctrl.T.Helper()
 	ret := m.ctrl.Call(m, "ProcessProposalVerifyTx", txBz)
-	ret0, _ := ret[0].(types0.Tx)
+	ret0, _ := ret[0].(types.Tx)
 	ret1, _ := ret[1].(error)
 	return ret0, ret1
 }
diff --git a/tests/integration/staking/keeper/vote_extensions_test.go b/tests/integration/staking/keeper/vote_extensions_test.go
new file mode 100644
index 000000000000..c992c5fbfa64
--- /dev/null
+++ b/tests/integration/staking/keeper/vote_extensions_test.go
@@ -0,0 +1,96 @@
+package keeper_test
+
+import (
+	"bytes"
+	"testing"
+
+	"cosmossdk.io/math"
+	abci "github.com/cometbft/cometbft/abci/types"
+	protoio "github.com/cosmos/gogoproto/io"
+	"github.com/cosmos/gogoproto/proto"
+	"gotest.tools/v3/assert"
+
+	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
+
+	"github.com/cosmos/cosmos-sdk/baseapp"
+	ed25519 "github.com/cosmos/cosmos-sdk/crypto/keys/ed25519"
+	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
+	simtestutil "github.com/cosmos/cosmos-sdk/testutil/sims"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/x/staking/testutil"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+)
+
+func TestValidateVoteExtensions(t *testing.T) {
+	t.Parallel()
+	f := initFixture(t)
+
+	// enable vote extensions
+	cp := simtestutil.DefaultConsensusParams
+	cp.Abci = &cmtproto.ABCIParams{VoteExtensionsEnableHeight: 1}
+	f.sdkCtx = f.sdkCtx.WithConsensusParams(*cp).WithBlockHeight(2)
+
+	// setup the validators
+	numVals := 3
+	privKeys := []cryptotypes.PrivKey{}
+	for i := 0; i < numVals; i++ {
+		privKeys = append(privKeys, ed25519.GenPrivKey())
+	}
+
+	vals := []stakingtypes.Validator{}
+	for _, v := range privKeys {
+		valAddr := sdk.ValAddress(v.PubKey().Address())
+		simtestutil.AddTestAddrsFromPubKeys(f.bankKeeper, f.stakingKeeper, f.sdkCtx, []cryptotypes.PubKey{v.PubKey()}, math.NewInt(100000000000))
+		vals = append(vals, testutil.NewValidator(t, valAddr, v.PubKey()))
+	}
+
+	votes := []abci.ExtendedVoteInfo{}
+
+	for i, v := range vals {
+		v.Tokens = math.NewInt(1000000)
+		v.Status = stakingtypes.Bonded
+		assert.NilError(t, f.stakingKeeper.SetValidator(f.sdkCtx, v))
+		assert.NilError(t, f.stakingKeeper.SetValidatorByConsAddr(f.sdkCtx, v))
+		assert.NilError(t, f.stakingKeeper.SetNewValidatorByPowerIndex(f.sdkCtx, v))
+		_, err := f.stakingKeeper.Delegate(f.sdkCtx, sdk.AccAddress(privKeys[i].PubKey().Address()), v.Tokens, stakingtypes.Unbonded, v, true)
+		assert.NilError(t, err)
+
+		// each val produces a vote
+		voteExt := []byte("something" + v.OperatorAddress)
+		cve := cmtproto.CanonicalVoteExtension{
+			Extension: voteExt,
+			Height:    f.sdkCtx.BlockHeight() - 1, // the vote extension was signed in the previous height
+			Round:     0,
+			ChainId:   "chain-id-123",
+		}
+
+		extSignBytes, err := mashalVoteExt(&cve)
+		assert.NilError(t, err)
+
+		sig, err := privKeys[i].Sign(extSignBytes)
+		assert.NilError(t, err)
+
+		ve := abci.ExtendedVoteInfo{
+			Validator: abci.Validator{
+				Address: v.GetOperator(),
+				Power:   v.ConsensusPower(sdk.DefaultPowerReduction),
+			},
+			VoteExtension:      voteExt,
+			ExtensionSignature: sig,
+			BlockIdFlag:        cmtproto.BlockIDFlagCommit,
+		}
+		votes = append(votes, ve)
+	}
+
+	err := baseapp.ValidateVoteExtensions(f.sdkCtx, f.stakingKeeper, f.sdkCtx.BlockHeight(), "chain-id-123", abci.ExtendedCommitInfo{Round: 0, Votes: votes})
+	assert.NilError(t, err)
+}
+
+func mashalVoteExt(msg proto.Message) ([]byte, error) {
+	var buf bytes.Buffer
+	if err := protoio.NewDelimitedWriter(&buf).WriteMsg(msg); err != nil {
+		return nil, err
+	}
+
+	return buf.Bytes(), nil
+}
diff --git a/x/staking/keeper/validator.go b/x/staking/keeper/validator.go
index 4db228ec1ecb..468111a24407 100644
--- a/x/staking/keeper/validator.go
+++ b/x/staking/keeper/validator.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"time"
 
+	cmtprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 	gogotypes "github.com/cosmos/gogoproto/types"
 
 	corestore "cosmossdk.io/core/store"
@@ -601,3 +602,21 @@ func (k Keeper) IsValidatorJailed(ctx context.Context, addr sdk.ConsAddress) (bo
 
 	return v.Jailed, nil
 }
+
+// CmtConsPublicKeyByConsAddr returns the consensus public key by consensus address
+func (k Keeper) CmtConsPublicKeyByConsAddr(ctx context.Context, addr sdk.ConsAddress) (cmtprotocrypto.PublicKey, error) {
+	v, err := k.GetValidatorByConsAddr(ctx, addr)
+	if err != nil {
+		return cmtprotocrypto.PublicKey{}, err
+	}
+	return v.CmtConsPublicKey()
+}
+
+// ConsAddressByConsPublicKey returns the consensus address by consensus address
+func (k Keeper) BondedTokensByConsAddr(ctx context.Context, addr sdk.ConsAddress) (math.Int, error) {
+	v, err := k.GetValidatorByConsAddr(ctx, addr)
+	if err != nil {
+		return math.ZeroInt(), err
+	}
+	return v.BondedTokens(), nil
+}
diff --git a/x/staking/testutil/expected_keepers_mocks.go b/x/staking/testutil/expected_keepers_mocks.go
index e66027edd2e7..c28beaaab850 100644
--- a/x/staking/testutil/expected_keepers_mocks.go
+++ b/x/staking/testutil/expected_keepers_mocks.go
@@ -10,6 +10,7 @@ import (
 
 	address "cosmossdk.io/core/address"
 	math "cosmossdk.io/math"
+	crypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 	types "github.com/cosmos/cosmos-sdk/types"
 	types0 "github.com/cosmos/cosmos-sdk/x/staking/types"
 	gomock "github.com/golang/mock/gomock"
@@ -341,6 +342,36 @@ func (m *MockValidatorSet) EXPECT() *MockValidatorSetMockRecorder {
 	return m.recorder
 }
 
+// BondedTokensByConsAddr mocks base method.
+func (m *MockValidatorSet) BondedTokensByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (math.Int, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "BondedTokensByConsAddr", arg0, arg1)
+	ret0, _ := ret[0].(math.Int)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// BondedTokensByConsAddr indicates an expected call of BondedTokensByConsAddr.
+func (mr *MockValidatorSetMockRecorder) BondedTokensByConsAddr(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokensByConsAddr", reflect.TypeOf((*MockValidatorSet)(nil).BondedTokensByConsAddr), arg0, arg1)
+}
+
+// CmtConsPublicKeyByConsAddr mocks base method.
+func (m *MockValidatorSet) CmtConsPublicKeyByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (crypto.PublicKey, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "CmtConsPublicKeyByConsAddr", arg0, arg1)
+	ret0, _ := ret[0].(crypto.PublicKey)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// CmtConsPublicKeyByConsAddr indicates an expected call of CmtConsPublicKeyByConsAddr.
+func (mr *MockValidatorSetMockRecorder) CmtConsPublicKeyByConsAddr(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CmtConsPublicKeyByConsAddr", reflect.TypeOf((*MockValidatorSet)(nil).CmtConsPublicKeyByConsAddr), arg0, arg1)
+}
+
 // Delegation mocks base method.
 func (m *MockValidatorSet) Delegation(arg0 context.Context, arg1 types.AccAddress, arg2 types.ValAddress) (types0.DelegationI, error) {
 	m.ctrl.T.Helper()
diff --git a/x/staking/types/expected_keepers.go b/x/staking/types/expected_keepers.go
index 9ab221103179..60ac3a9dab87 100644
--- a/x/staking/types/expected_keepers.go
+++ b/x/staking/types/expected_keepers.go
@@ -5,6 +5,7 @@ import (
 
 	"cosmossdk.io/core/address"
 	"cosmossdk.io/math"
+	cmtprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
@@ -76,6 +77,12 @@ type ValidatorSet interface {
 
 	// MaxValidators returns the maximum amount of bonded validators
 	MaxValidators(context.Context) (uint32, error)
+
+	// CmtConsPublicKeyByConsAddr returns the consensus public key by consensus address
+	CmtConsPublicKeyByConsAddr(context.Context, sdk.ConsAddress) (cmtprotocrypto.PublicKey, error)
+
+	// ConsAddressByConsPublicKey returns the consensus address by consensus address
+	BondedTokensByConsAddr(context.Context, sdk.ConsAddress) (math.Int, error)
 }
 
 // DelegationSet expected properties for the set of all delegations for a particular (noalias)

From 8ea7b79e1a843c9aaea57caf9538dde68fae7084 Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Fri, 28 Jul 2023 11:00:48 +0200
Subject: [PATCH 3/8] lint and cleanup

---
 baseapp/abci_test.go                                     | 9 +++------
 baseapp/abci_utils.go                                    | 1 -
 tests/integration/staking/keeper/vote_extensions_test.go | 4 ++--
 x/staking/types/expected_keepers.go                      | 3 ++-
 4 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/baseapp/abci_test.go b/baseapp/abci_test.go
index bac900ab4bc9..71dde8cbceda 100644
--- a/baseapp/abci_test.go
+++ b/baseapp/abci_test.go
@@ -1733,12 +1733,9 @@ func TestABCI_PrepareProposal_VoteExtensions(t *testing.T) {
 		},
 	}
 
-	val1 := mock.NewMockValidator(ctrl)
-	val1.EXPECT().BondedTokens().Return(math.NewInt(667))
-	val1.EXPECT().CmtConsPublicKey().Return(tmPk, nil).AnyTimes()
-
 	consAddr := sdk.ConsAddress(addr.String())
-	valStore.EXPECT().GetValidatorByConsAddr(gomock.Any(), consAddr.Bytes()).Return(val1, nil).AnyTimes()
+	valStore.EXPECT().CmtConsPublicKeyByConsAddr(gomock.Any(), consAddr.Bytes()).Return(tmPk, nil).AnyTimes()
+	valStore.EXPECT().BondedTokensByConsAddr(gomock.Any(), consAddr.Bytes()).Return(math.NewInt(667), nil)
 	valStore.EXPECT().TotalBondedTokens(gomock.Any()).Return(math.NewInt(1000), nil).AnyTimes()
 
 	// set up baseapp
@@ -1827,7 +1824,7 @@ func TestABCI_PrepareProposal_VoteExtensions(t *testing.T) {
 	require.Equal(t, 1, len(resPrepareProposal.Txs))
 
 	// now vote extensions but our sole voter doesn't reach majority
-	val1.EXPECT().BondedTokens().Return(math.NewInt(666))
+	valStore.EXPECT().BondedTokensByConsAddr(gomock.Any(), consAddr.Bytes()).Return(math.NewInt(666), nil)
 	resPrepareProposal, err = suite.baseApp.PrepareProposal(&reqPrepareProposal)
 	require.NoError(t, err)
 	require.Equal(t, 0, len(resPrepareProposal.Txs))
diff --git a/baseapp/abci_utils.go b/baseapp/abci_utils.go
index 933f9db62798..da1f144107ba 100644
--- a/baseapp/abci_utils.go
+++ b/baseapp/abci_utils.go
@@ -54,7 +54,6 @@ func ValidateVoteExtensions(
 ) error {
 	cp := ctx.ConsensusParams()
 	extsEnabled := cp.Abci != nil && currentHeight >= cp.Abci.VoteExtensionsEnableHeight && cp.Abci.VoteExtensionsEnableHeight != 0
-	fmt.Println(currentHeight, cp)
 	marshalDelimitedFn := func(msg proto.Message) ([]byte, error) {
 		var buf bytes.Buffer
 		if err := protoio.NewDelimitedWriter(&buf).WriteMsg(msg); err != nil {
diff --git a/tests/integration/staking/keeper/vote_extensions_test.go b/tests/integration/staking/keeper/vote_extensions_test.go
index c992c5fbfa64..5947373c1c3b 100644
--- a/tests/integration/staking/keeper/vote_extensions_test.go
+++ b/tests/integration/staking/keeper/vote_extensions_test.go
@@ -4,13 +4,13 @@ import (
 	"bytes"
 	"testing"
 
-	"cosmossdk.io/math"
 	abci "github.com/cometbft/cometbft/abci/types"
+	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	protoio "github.com/cosmos/gogoproto/io"
 	"github.com/cosmos/gogoproto/proto"
 	"gotest.tools/v3/assert"
 
-	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
+	"cosmossdk.io/math"
 
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	ed25519 "github.com/cosmos/cosmos-sdk/crypto/keys/ed25519"
diff --git a/x/staking/types/expected_keepers.go b/x/staking/types/expected_keepers.go
index 60ac3a9dab87..5afc060c7b5d 100644
--- a/x/staking/types/expected_keepers.go
+++ b/x/staking/types/expected_keepers.go
@@ -3,9 +3,10 @@ package types
 import (
 	context "context"
 
+	cmtprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
+
 	"cosmossdk.io/core/address"
 	"cosmossdk.io/math"
-	cmtprotocrypto "github.com/cometbft/cometbft/proto/tendermint/crypto"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )

From c08742f228de454008fcf3f3cc103729e996634f Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Fri, 28 Jul 2023 14:04:12 +0200
Subject: [PATCH 4/8] update adr

---
 docs/architecture/adr-064-abci-2.0.md | 17 +++++------------
 1 file changed, 5 insertions(+), 12 deletions(-)

diff --git a/docs/architecture/adr-064-abci-2.0.md b/docs/architecture/adr-064-abci-2.0.md
index 98651a43abfe..2e18e3772c9c 100644
--- a/docs/architecture/adr-064-abci-2.0.md
+++ b/docs/architecture/adr-064-abci-2.0.md
@@ -218,7 +218,9 @@ a default signature verification method which applications can use:
 
 ```go
 type ValidatorStore interface {
-	GetValidatorByConsAddr(sdk.Context, cryptotypes.Address) (cryptotypes.PubKey, error)
+	TotalBondedTokens(ctx context.Context) (math.Int, error)
+	CmtConsPublicKeyByConsAddr(context.Context, sdk.ConsAddress) (cmtprotocrypto.PublicKey, error)
+	BondedTokensByConsAddr(context.Context, sdk.ConsAddress) (math.Int, error)
 }
 
 // ValidateVoteExtensions is a function that an application can execute in
@@ -229,17 +231,8 @@ func (app *BaseApp) ValidateVoteExtensions(ctx sdk.Context, currentHeight int64,
 			continue
 		}
 
-		valConsAddr := cmtcrypto.Address(vote.Validator.Address)
-
-		validator, err := app.validatorStore.GetValidatorByConsAddr(ctx, valConsAddr)
-		if err != nil {
-			return fmt.Errorf("failed to get validator %s for vote extension", valConsAddr)
-		}
-
-		cmtPubKey, err := validator.CmtConsPublicKey()
-		if err != nil {
-			return fmt.Errorf("failed to convert public key: %w", err)
-		}
+		valConsAddr := sdk.ConsAddress(vote.Validator.Address)
+		cmtPubKey, err := valStore.CmtConsPublicKeyByConsAddr(ctx, valConsAddr)
 
 		if len(vote.ExtensionSignature) == 0 {
 			return fmt.Errorf("received a non-empty vote extension with empty signature for validator %s", valConsAddr)

From 9ba8443e4d65f0e077ff5ed046eeb282076c1445 Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Fri, 28 Jul 2023 14:06:05 +0200
Subject: [PATCH 5/8] cl++

---
 CHANGELOG.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4f3b4c522d14..bc16b0001b29 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/group, x/gov) [#17109](https://github.com/cosmos/cosmos-sdk/pull/17109) Let proposal summary be 40x longer than metadata limit.
 * (all) [#16537](https://github.com/cosmos/cosmos-sdk/pull/16537) Properly propagated `fmt.Errorf` errors and using `errors.New` where appropriate.
 * (version) [#17096](https://github.com/cosmos/cosmos-sdk/pull/17096) Improve `getSDKVersion()` to handle module replacements
+* (x/staking) [#17164](https://github.com/cosmos/cosmos-sdk/pull/17164) Add `CmtConsPublicKeyByConsAddr` and `BondedTokensByConsAddr` to the keeper to enable vote extension verification.
 
 ### Bug Fixes
 

From 1fb12d744a85fab23412f2c9d0c23f38941210a7 Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Sat, 29 Jul 2023 19:41:23 +0200
Subject: [PATCH 6/8] use BondedTokensAndPubKeyByConsAddr

---
 CHANGELOG.md                                 |  2 +-
 baseapp/abci_utils.go                        | 11 ++-----
 baseapp/testutil/mock/mocks.go               | 32 ++++++--------------
 docs/architecture/adr-064-abci-2.0.md        |  5 ++-
 x/gov/testutil/expected_keepers_mocks.go     |  2 +-
 x/staking/keeper/validator.go                | 17 +++++------
 x/staking/testutil/expected_keepers_mocks.go | 32 ++++++--------------
 x/staking/types/expected_keepers.go          |  8 ++---
 8 files changed, 35 insertions(+), 74 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 4f0e54bf4378..b8a86b0eeb87 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -47,7 +47,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (x/group, x/gov) [#17109](https://github.com/cosmos/cosmos-sdk/pull/17109) Let proposal summary be 40x longer than metadata limit.
 * (all) [#16537](https://github.com/cosmos/cosmos-sdk/pull/16537) Properly propagated `fmt.Errorf` errors and using `errors.New` where appropriate.
 * (version) [#17096](https://github.com/cosmos/cosmos-sdk/pull/17096) Improve `getSDKVersion()` to handle module replacements
-* (x/staking) [#17164](https://github.com/cosmos/cosmos-sdk/pull/17164) Add `CmtConsPublicKeyByConsAddr` and `BondedTokensByConsAddr` to the keeper to enable vote extension verification.
+* (x/staking) [#17164](https://github.com/cosmos/cosmos-sdk/pull/17164) Add `BondedTokensAndPubKeyByConsAddr` to the keeper to enable vote extension verification.
 
 ### Bug Fixes
 
diff --git a/baseapp/abci_utils.go b/baseapp/abci_utils.go
index da1f144107ba..39bb4e938091 100644
--- a/baseapp/abci_utils.go
+++ b/baseapp/abci_utils.go
@@ -30,8 +30,7 @@ type (
 	// module, which has knowledge of the CometBFT public key.
 	ValidatorStore interface {
 		TotalBondedTokens(ctx context.Context) (math.Int, error)
-		CmtConsPublicKeyByConsAddr(context.Context, sdk.ConsAddress) (cmtprotocrypto.PublicKey, error)
-		BondedTokensByConsAddr(context.Context, sdk.ConsAddress) (math.Int, error)
+		BondedTokensAndPubKeyByConsAddr(context.Context, sdk.ConsAddress) (math.Int, cmtprotocrypto.PublicKey, error)
 	}
 
 	// GasTx defines the contract that a transaction with a gas limit must implement.
@@ -81,9 +80,9 @@ func ValidateVoteExtensions(
 		}
 
 		valConsAddr := sdk.ConsAddress(vote.Validator.Address)
-		cmtPubKeyProto, err := valStore.CmtConsPublicKeyByConsAddr(ctx, valConsAddr)
+		bondedTokens, cmtPubKeyProto, err := valStore.BondedTokensAndPubKeyByConsAddr(ctx, valConsAddr)
 		if err != nil {
-			return fmt.Errorf("failed to get validator %X public key: %w", valConsAddr, err)
+			return fmt.Errorf("failed to get validator %X info (bonded tokens and public key): %w", valConsAddr, err)
 		}
 
 		cmtPubKey, err := cryptoenc.PubKeyFromProto(cmtPubKeyProto)
@@ -107,10 +106,6 @@ func ValidateVoteExtensions(
 			return fmt.Errorf("failed to verify validator %X vote extension signature", valConsAddr)
 		}
 
-		bondedTokens, err := valStore.BondedTokensByConsAddr(ctx, valConsAddr)
-		if err != nil {
-			return fmt.Errorf("failed to get validator %X bonded tokens: %w", valConsAddr, err)
-		}
 		sumVP = sumVP.Add(bondedTokens)
 	}
 
diff --git a/baseapp/testutil/mock/mocks.go b/baseapp/testutil/mock/mocks.go
index d2f34574f53c..482e8f59dd17 100644
--- a/baseapp/testutil/mock/mocks.go
+++ b/baseapp/testutil/mock/mocks.go
@@ -37,34 +37,20 @@ func (m *MockValidatorStore) EXPECT() *MockValidatorStoreMockRecorder {
 	return m.recorder
 }
 
-// BondedTokensByConsAddr mocks base method.
-func (m *MockValidatorStore) BondedTokensByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (math.Int, error) {
+// BondedTokensAndPubKeyByConsAddr mocks base method.
+func (m *MockValidatorStore) BondedTokensAndPubKeyByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (math.Int, crypto.PublicKey, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "BondedTokensByConsAddr", arg0, arg1)
+	ret := m.ctrl.Call(m, "BondedTokensAndPubKeyByConsAddr", arg0, arg1)
 	ret0, _ := ret[0].(math.Int)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// BondedTokensByConsAddr indicates an expected call of BondedTokensByConsAddr.
-func (mr *MockValidatorStoreMockRecorder) BondedTokensByConsAddr(arg0, arg1 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokensByConsAddr", reflect.TypeOf((*MockValidatorStore)(nil).BondedTokensByConsAddr), arg0, arg1)
-}
-
-// CmtConsPublicKeyByConsAddr mocks base method.
-func (m *MockValidatorStore) CmtConsPublicKeyByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (crypto.PublicKey, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "CmtConsPublicKeyByConsAddr", arg0, arg1)
-	ret0, _ := ret[0].(crypto.PublicKey)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret1, _ := ret[1].(crypto.PublicKey)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
 }
 
-// CmtConsPublicKeyByConsAddr indicates an expected call of CmtConsPublicKeyByConsAddr.
-func (mr *MockValidatorStoreMockRecorder) CmtConsPublicKeyByConsAddr(arg0, arg1 interface{}) *gomock.Call {
+// BondedTokensAndPubKeyByConsAddr indicates an expected call of BondedTokensAndPubKeyByConsAddr.
+func (mr *MockValidatorStoreMockRecorder) BondedTokensAndPubKeyByConsAddr(arg0, arg1 interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CmtConsPublicKeyByConsAddr", reflect.TypeOf((*MockValidatorStore)(nil).CmtConsPublicKeyByConsAddr), arg0, arg1)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokensAndPubKeyByConsAddr", reflect.TypeOf((*MockValidatorStore)(nil).BondedTokensAndPubKeyByConsAddr), arg0, arg1)
 }
 
 // TotalBondedTokens mocks base method.
diff --git a/docs/architecture/adr-064-abci-2.0.md b/docs/architecture/adr-064-abci-2.0.md
index 2e18e3772c9c..fc1278ecf379 100644
--- a/docs/architecture/adr-064-abci-2.0.md
+++ b/docs/architecture/adr-064-abci-2.0.md
@@ -219,8 +219,7 @@ a default signature verification method which applications can use:
 ```go
 type ValidatorStore interface {
 	TotalBondedTokens(ctx context.Context) (math.Int, error)
-	CmtConsPublicKeyByConsAddr(context.Context, sdk.ConsAddress) (cmtprotocrypto.PublicKey, error)
-	BondedTokensByConsAddr(context.Context, sdk.ConsAddress) (math.Int, error)
+	BondedTokensAndPubKeyByConsAddr(context.Context, sdk.ConsAddress) (math.Int, cmtprotocrypto.PublicKey, error)
 }
 
 // ValidateVoteExtensions is a function that an application can execute in
@@ -232,7 +231,7 @@ func (app *BaseApp) ValidateVoteExtensions(ctx sdk.Context, currentHeight int64,
 		}
 
 		valConsAddr := sdk.ConsAddress(vote.Validator.Address)
-		cmtPubKey, err := valStore.CmtConsPublicKeyByConsAddr(ctx, valConsAddr)
+		bondedTokens, cmtPubKey, err := valStore.BondedTokensAndPubKeyByConsAddr(ctx, valConsAddr)
 
 		if len(vote.ExtensionSignature) == 0 {
 			return fmt.Errorf("received a non-empty vote extension with empty signature for validator %s", valConsAddr)
diff --git a/x/gov/testutil/expected_keepers_mocks.go b/x/gov/testutil/expected_keepers_mocks.go
index 1a7c78d61df6..0dc0ac0b6365 100644
--- a/x/gov/testutil/expected_keepers_mocks.go
+++ b/x/gov/testutil/expected_keepers_mocks.go
@@ -1,7 +1,7 @@
 // Code generated by MockGen. DO NOT EDIT.
 // Source: x/gov/testutil/expected_keepers.go
 
-// Package mock_testutil is a generated GoMock package.
+// Package testutil is a generated GoMock package.
 package testutil
 
 import (
diff --git a/x/staking/keeper/validator.go b/x/staking/keeper/validator.go
index 468111a24407..cd80cb6eb410 100644
--- a/x/staking/keeper/validator.go
+++ b/x/staking/keeper/validator.go
@@ -603,20 +603,17 @@ func (k Keeper) IsValidatorJailed(ctx context.Context, addr sdk.ConsAddress) (bo
 	return v.Jailed, nil
 }
 
-// CmtConsPublicKeyByConsAddr returns the consensus public key by consensus address
-func (k Keeper) CmtConsPublicKeyByConsAddr(ctx context.Context, addr sdk.ConsAddress) (cmtprotocrypto.PublicKey, error) {
+// CmtConsPublicKeyByConsAddr returns the consensus public key and bonded tokens by consensus address
+func (k Keeper) BondedTokensAndPubKeyByConsAddr(ctx context.Context, addr sdk.ConsAddress) (math.Int, cmtprotocrypto.PublicKey, error) {
 	v, err := k.GetValidatorByConsAddr(ctx, addr)
 	if err != nil {
-		return cmtprotocrypto.PublicKey{}, err
+		return math.ZeroInt(), cmtprotocrypto.PublicKey{}, err
 	}
-	return v.CmtConsPublicKey()
-}
 
-// ConsAddressByConsPublicKey returns the consensus address by consensus address
-func (k Keeper) BondedTokensByConsAddr(ctx context.Context, addr sdk.ConsAddress) (math.Int, error) {
-	v, err := k.GetValidatorByConsAddr(ctx, addr)
+	pubkey, err := v.CmtConsPublicKey()
 	if err != nil {
-		return math.ZeroInt(), err
+		return math.ZeroInt(), cmtprotocrypto.PublicKey{}, err
 	}
-	return v.BondedTokens(), nil
+
+	return v.BondedTokens(), pubkey, nil
 }
diff --git a/x/staking/testutil/expected_keepers_mocks.go b/x/staking/testutil/expected_keepers_mocks.go
index c28beaaab850..85071728020e 100644
--- a/x/staking/testutil/expected_keepers_mocks.go
+++ b/x/staking/testutil/expected_keepers_mocks.go
@@ -342,34 +342,20 @@ func (m *MockValidatorSet) EXPECT() *MockValidatorSetMockRecorder {
 	return m.recorder
 }
 
-// BondedTokensByConsAddr mocks base method.
-func (m *MockValidatorSet) BondedTokensByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (math.Int, error) {
+// BondedTokensAndPubKeyByConsAddr mocks base method.
+func (m *MockValidatorSet) BondedTokensAndPubKeyByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (math.Int, crypto.PublicKey, error) {
 	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "BondedTokensByConsAddr", arg0, arg1)
+	ret := m.ctrl.Call(m, "BondedTokensAndPubKeyByConsAddr", arg0, arg1)
 	ret0, _ := ret[0].(math.Int)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
-}
-
-// BondedTokensByConsAddr indicates an expected call of BondedTokensByConsAddr.
-func (mr *MockValidatorSetMockRecorder) BondedTokensByConsAddr(arg0, arg1 interface{}) *gomock.Call {
-	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokensByConsAddr", reflect.TypeOf((*MockValidatorSet)(nil).BondedTokensByConsAddr), arg0, arg1)
-}
-
-// CmtConsPublicKeyByConsAddr mocks base method.
-func (m *MockValidatorSet) CmtConsPublicKeyByConsAddr(arg0 context.Context, arg1 types.ConsAddress) (crypto.PublicKey, error) {
-	m.ctrl.T.Helper()
-	ret := m.ctrl.Call(m, "CmtConsPublicKeyByConsAddr", arg0, arg1)
-	ret0, _ := ret[0].(crypto.PublicKey)
-	ret1, _ := ret[1].(error)
-	return ret0, ret1
+	ret1, _ := ret[1].(crypto.PublicKey)
+	ret2, _ := ret[2].(error)
+	return ret0, ret1, ret2
 }
 
-// CmtConsPublicKeyByConsAddr indicates an expected call of CmtConsPublicKeyByConsAddr.
-func (mr *MockValidatorSetMockRecorder) CmtConsPublicKeyByConsAddr(arg0, arg1 interface{}) *gomock.Call {
+// BondedTokensAndPubKeyByConsAddr indicates an expected call of BondedTokensAndPubKeyByConsAddr.
+func (mr *MockValidatorSetMockRecorder) BondedTokensAndPubKeyByConsAddr(arg0, arg1 interface{}) *gomock.Call {
 	mr.mock.ctrl.T.Helper()
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CmtConsPublicKeyByConsAddr", reflect.TypeOf((*MockValidatorSet)(nil).CmtConsPublicKeyByConsAddr), arg0, arg1)
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "BondedTokensAndPubKeyByConsAddr", reflect.TypeOf((*MockValidatorSet)(nil).BondedTokensAndPubKeyByConsAddr), arg0, arg1)
 }
 
 // Delegation mocks base method.
diff --git a/x/staking/types/expected_keepers.go b/x/staking/types/expected_keepers.go
index 5afc060c7b5d..6488a9d3f68b 100644
--- a/x/staking/types/expected_keepers.go
+++ b/x/staking/types/expected_keepers.go
@@ -79,11 +79,9 @@ type ValidatorSet interface {
 	// MaxValidators returns the maximum amount of bonded validators
 	MaxValidators(context.Context) (uint32, error)
 
-	// CmtConsPublicKeyByConsAddr returns the consensus public key by consensus address
-	CmtConsPublicKeyByConsAddr(context.Context, sdk.ConsAddress) (cmtprotocrypto.PublicKey, error)
-
-	// ConsAddressByConsPublicKey returns the consensus address by consensus address
-	BondedTokensByConsAddr(context.Context, sdk.ConsAddress) (math.Int, error)
+	// BondedTokensAndPubKeyByConsAddr returns the bonded tokens and consensus public key for a validator.
+	// Used in vote extension validation.
+	BondedTokensAndPubKeyByConsAddr(context.Context, sdk.ConsAddress) (math.Int, cmtprotocrypto.PublicKey, error)
 }
 
 // DelegationSet expected properties for the set of all delegations for a particular (noalias)

From f5ff37496323d76b37242778493f72da743c9a34 Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Sat, 29 Jul 2023 19:51:57 +0200
Subject: [PATCH 7/8] fix test

---
 baseapp/abci_test.go | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/baseapp/abci_test.go b/baseapp/abci_test.go
index 71dde8cbceda..4ea39b81a188 100644
--- a/baseapp/abci_test.go
+++ b/baseapp/abci_test.go
@@ -1734,8 +1734,7 @@ func TestABCI_PrepareProposal_VoteExtensions(t *testing.T) {
 	}
 
 	consAddr := sdk.ConsAddress(addr.String())
-	valStore.EXPECT().CmtConsPublicKeyByConsAddr(gomock.Any(), consAddr.Bytes()).Return(tmPk, nil).AnyTimes()
-	valStore.EXPECT().BondedTokensByConsAddr(gomock.Any(), consAddr.Bytes()).Return(math.NewInt(667), nil)
+	valStore.EXPECT().BondedTokensAndPubKeyByConsAddr(gomock.Any(), consAddr.Bytes()).Return(math.NewInt(667), tmPk, nil)
 	valStore.EXPECT().TotalBondedTokens(gomock.Any()).Return(math.NewInt(1000), nil).AnyTimes()
 
 	// set up baseapp
@@ -1824,7 +1823,7 @@ func TestABCI_PrepareProposal_VoteExtensions(t *testing.T) {
 	require.Equal(t, 1, len(resPrepareProposal.Txs))
 
 	// now vote extensions but our sole voter doesn't reach majority
-	valStore.EXPECT().BondedTokensByConsAddr(gomock.Any(), consAddr.Bytes()).Return(math.NewInt(666), nil)
+	valStore.EXPECT().BondedTokensAndPubKeyByConsAddr(gomock.Any(), consAddr.Bytes()).Return(math.NewInt(666), tmPk, nil)
 	resPrepareProposal, err = suite.baseApp.PrepareProposal(&reqPrepareProposal)
 	require.NoError(t, err)
 	require.Equal(t, 0, len(resPrepareProposal.Txs))

From 72bbdfd4dc5b530317cc6fd147cabbadf2090d03 Mon Sep 17 00:00:00 2001
From: Facundo Medica <facundomedica@gmail.com>
Date: Mon, 31 Jul 2023 10:09:51 +0200
Subject: [PATCH 8/8] fixes

---
 docs/architecture/adr-064-abci-2.0.md | 3 +++
 x/staking/keeper/validator.go         | 2 +-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/docs/architecture/adr-064-abci-2.0.md b/docs/architecture/adr-064-abci-2.0.md
index fc1278ecf379..1d6dd1df522b 100644
--- a/docs/architecture/adr-064-abci-2.0.md
+++ b/docs/architecture/adr-064-abci-2.0.md
@@ -232,6 +232,9 @@ func (app *BaseApp) ValidateVoteExtensions(ctx sdk.Context, currentHeight int64,
 
 		valConsAddr := sdk.ConsAddress(vote.Validator.Address)
 		bondedTokens, cmtPubKey, err := valStore.BondedTokensAndPubKeyByConsAddr(ctx, valConsAddr)
+		if err != nil {
+			return fmt.Errorf("failed to get bonded tokens and public key for validator %s: %w", valConsAddr, err)
+		}
 
 		if len(vote.ExtensionSignature) == 0 {
 			return fmt.Errorf("received a non-empty vote extension with empty signature for validator %s", valConsAddr)
diff --git a/x/staking/keeper/validator.go b/x/staking/keeper/validator.go
index cd80cb6eb410..e54ad6e55ac1 100644
--- a/x/staking/keeper/validator.go
+++ b/x/staking/keeper/validator.go
@@ -603,7 +603,7 @@ func (k Keeper) IsValidatorJailed(ctx context.Context, addr sdk.ConsAddress) (bo
 	return v.Jailed, nil
 }
 
-// CmtConsPublicKeyByConsAddr returns the consensus public key and bonded tokens by consensus address
+// BondedTokensAndPubKeyByConsAddr returns the consensus public key and bonded tokens by consensus address
 func (k Keeper) BondedTokensAndPubKeyByConsAddr(ctx context.Context, addr sdk.ConsAddress) (math.Int, cmtprotocrypto.PublicKey, error) {
 	v, err := k.GetValidatorByConsAddr(ctx, addr)
 	if err != nil {