fix: ValidatorStore interface doesn't match StakingKeeper (#17164)

Co-authored-by: Aleksandr Bezobchuk <alexanderbez@users.noreply.github.com>

CHANGELOG.md

@@ -48,6 +48,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 `BondedTokensAndPubKeyByConsAddr` to the keeper to enable vote extension verification.
 
 ### Bug Fixes
 

baseapp/abci_test.go

@@ -1790,13 +1790,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().TotalBondedTokens(gomock.Any()).Return(math.NewInt(1000)).AnyTimes()
+	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
 	prepareOpt := func(bapp *baseapp.BaseApp) {
@@ -1884,7 +1880,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().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))

baseapp/abci_utils.go

@@ -2,11 +2,11 @@ package baseapp
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 
 	"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"
@@ -15,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"
 )
@@ -26,20 +25,12 @@ 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(sdk.Context, cryptotypes.Address) (Validator, error)
-		TotalBondedTokens(ctx sdk.Context) math.Int
+		TotalBondedTokens(ctx context.Context) (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.
@@ -62,7 +53,6 @@ func ValidateVoteExtensions(
 ) error {
 	cp := ctx.ConsensusParams()
 	extsEnabled := cp.Abci != nil && currentHeight >= cp.Abci.VoteExtensionsEnableHeight && cp.Abci.VoteExtensionsEnableHeight != 0
-
 	marshalDelimitedFn := func(msg proto.Message) ([]byte, error) {
 		var buf bytes.Buffer
 		if err := protoio.NewDelimitedWriter(&buf).WriteMsg(msg); err != nil {
@@ -89,19 +79,10 @@ 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)
+		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)
@@ -125,12 +106,16 @@ func ValidateVoteExtensions(
 			return fmt.Errorf("failed to verify validator %X vote extension signature", valConsAddr)
 		}
 
-		sumVP = sumVP.Add(validator.BondedTokens())
+		sumVP = sumVP.Add(bondedTokens)
 	}
 
 	// 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)

docs/architecture/adr-064-abci-2.0.md

@@ -219,7 +219,8 @@ 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)
+	BondedTokensAndPubKeyByConsAddr(context.Context, sdk.ConsAddress) (math.Int, cmtprotocrypto.PublicKey, error)
 }
 
 // ValidateVoteExtensions is a function that an application can execute in
@@ -230,16 +231,10 @@ 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()
+		valConsAddr := sdk.ConsAddress(vote.Validator.Address)
+		bondedTokens, cmtPubKey, err := valStore.BondedTokensAndPubKeyByConsAddr(ctx, valConsAddr)
 		if err != nil {
-			return fmt.Errorf("failed to convert public key: %w", err)
+			return fmt.Errorf("failed to get bonded tokens and public key for validator %s: %w", valConsAddr, err)
 		}
 
 		if len(vote.ExtensionSignature) == 0 {

tests/integration/staking/keeper/vote_extensions_test.go

@@ -0,0 +1,96 @@
+package keeper_test
+
+import (
+	"bytes"
+	"testing"
+
+	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"
+
+	"cosmossdk.io/math"
+
+	"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
+}