feat: Add preFinalizeBlockHook to allow VE persistence (#16898)

Co-authored-by: Aleksandr Bezobchuk <[email protected]>

CHANGELOG.md

@@ -40,6 +40,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Features
 
+* (baseapp) [#16898](https://github.com/cosmos/cosmos-sdk/pull/16898) Add `preFinalizeBlockHook` to allow vote extensions persistence.
 * (cli) [#16887](https://github.com/cosmos/cosmos-sdk/pull/16887) Add two new CLI commands: `tx simulate` for simulating a transaction; `query block-results` for querying CometBFT RPC for block results.
 * (x/gov) [#16976](https://github.com/cosmos/cosmos-sdk/pull/16976) Add `failed_reason` field to `Proposal` under `x/gov` to indicate the reason for a failed proposal. Referenced from [#238](https://github.com/bnb-chain/greenfield-cosmos-sdk/pull/238) under `bnb-chain/greenfield-cosmos-sdk`.
 

baseapp/abci.go

@@ -45,13 +45,14 @@ func (app *BaseApp) InitChain(req *abci.RequestInitChain) (*abci.ResponseInitCha
 	// On a new chain, we consider the init chain block height as 0, even though
 	// req.InitialHeight is 1 by default.
 	initHeader := cmtproto.Header{ChainID: req.ChainId, Time: req.Time}
-	app.initialHeight = req.InitialHeight
-
 	app.logger.Info("InitChain", "initialHeight", req.InitialHeight, "chainID", req.ChainId)
 
 	// Set the initial height, which will be used to determine if we are proposing
 	// or processing the first block or not.
 	app.initialHeight = req.InitialHeight
+	if app.initialHeight == 0 { // If initial height is 0, set it to 1
+		app.initialHeight = 1
+	}
 
 	// if req.InitialHeight is > 1, then we set the initial version on all stores
 	if req.InitialHeight > 1 {
@@ -675,53 +676,50 @@ func (app *BaseApp) FinalizeBlock(req *abci.RequestFinalizeBlock) (*abci.Respons
 		AppHash:            app.LastCommitID().Hash,
 	}
 
-	// Initialize the FinalizeBlock state. If this is the first block, it should
-	// already be initialized in InitChain. Otherwise app.finalizeBlockState will be
-	// nil, since it is reset on Commit.
+	// finalizeBlockState should be set on InitChain or ProcessProposal. If it is
+	// nil, it means we are replaying this block and we need to set the state here
+	// given that during block replay ProcessProposal is not executed by CometBFT.
 	if app.finalizeBlockState == nil {
 		app.setState(execModeFinalize, header)
-	} else {
-		// In the first block, app.finalizeBlockState.ctx will already be initialized
-		// by InitChain. Context is now updated with Header information.
-		app.finalizeBlockState.ctx = app.finalizeBlockState.ctx.
-			WithBlockHeader(header).
-			WithBlockHeight(req.Height).
-			WithHeaderInfo(coreheader.Info{
-				ChainID: app.chainID,
-				Height:  req.Height,
-				Time:    req.Time,
-				Hash:    req.Hash,
-				AppHash: app.LastCommitID().Hash,
-			})
 	}
 
-	gasMeter := app.getBlockGasMeter(app.finalizeBlockState.ctx)
-
+	// Context is now updated with Header information.
 	app.finalizeBlockState.ctx = app.finalizeBlockState.ctx.
-		WithBlockGasMeter(gasMeter).
+		WithBlockHeader(header).
 		WithHeaderHash(req.Hash).
-		WithConsensusParams(app.GetConsensusParams(app.finalizeBlockState.ctx)).
-		WithVoteInfos(req.DecidedLastCommit.Votes).
-		WithExecMode(sdk.ExecModeFinalize).
 		WithHeaderInfo(coreheader.Info{
 			ChainID: app.chainID,
 			Height:  req.Height,
 			Time:    req.Time,
 			Hash:    req.Hash,
 			AppHash: app.LastCommitID().Hash,
-		}).WithCometInfo(cometInfo{
-		Misbehavior:     req.Misbehavior,
-		ValidatorsHash:  req.NextValidatorsHash,
-		ProposerAddress: req.ProposerAddress,
-		LastCommit:      req.DecidedLastCommit,
-	})
+		}).
+		WithConsensusParams(app.GetConsensusParams(app.finalizeBlockState.ctx)).
+		WithVoteInfos(req.DecidedLastCommit.Votes).
+		WithExecMode(sdk.ExecModeFinalize).
+		WithCometInfo(cometInfo{
+			Misbehavior:     req.Misbehavior,
+			ValidatorsHash:  req.NextValidatorsHash,
+			ProposerAddress: req.ProposerAddress,
+			LastCommit:      req.DecidedLastCommit,
+		})
+
+	// GasMeter must be set after we get a context with updated consensus params.
+	gasMeter := app.getBlockGasMeter(app.finalizeBlockState.ctx)
+	app.finalizeBlockState.ctx = app.finalizeBlockState.ctx.WithBlockGasMeter(gasMeter)
 
 	if app.checkState != nil {
 		app.checkState.ctx = app.checkState.ctx.
 			WithBlockGasMeter(gasMeter).
 			WithHeaderHash(req.Hash)
 	}
 
+	if app.preFinalizeBlockHook != nil {
+		if err := app.preFinalizeBlockHook(app.finalizeBlockState.ctx, req); err != nil {
+			return nil, err
+		}
+	}
+
 	beginBlock := app.beginBlock(req)
 	events = append(events, beginBlock.Events...)
 

baseapp/abci_test.go

@@ -3,9 +3,11 @@ package baseapp_test
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/hex"
 	"errors"
 	"fmt"
+	"math/rand"
 	"strconv"
 	"strings"
 	"testing"
@@ -1908,3 +1910,195 @@ func TestABCI_HaltChain(t *testing.T) {
 		})
 	}
 }
+
+func TestBaseApp_PreFinalizeBlockHook(t *testing.T) {
+	db := dbm.NewMemDB()
+	name := t.Name()
+	logger := log.NewTestLogger(t)
+
+	app := baseapp.NewBaseApp(name, logger, db, nil)
+	_, err := app.InitChain(&abci.RequestInitChain{})
+	require.NoError(t, err)
+
+	wasHookCalled := false
+	app.SetPreFinalizeBlockHook(func(ctx sdk.Context, req *abci.RequestFinalizeBlock) error {
+		wasHookCalled = true
+		return nil
+	})
+	app.Seal()
+
+	_, err = app.FinalizeBlock(&abci.RequestFinalizeBlock{Height: 1})
+	require.NoError(t, err)
+	require.Equal(t, true, wasHookCalled)
+
+	// Now try erroring
+	app = baseapp.NewBaseApp(name, logger, db, nil)
+	_, err = app.InitChain(&abci.RequestInitChain{})
+	require.NoError(t, err)
+
+	app.SetPreFinalizeBlockHook(func(ctx sdk.Context, req *abci.RequestFinalizeBlock) error {
+		return errors.New("some error")
+	})
+	app.Seal()
+
+	_, err = app.FinalizeBlock(&abci.RequestFinalizeBlock{Height: 1})
+	require.Error(t, err)
+}
+
+// TestBaseApp_VoteExtensions tests vote extensions using a price as an example.
+func TestBaseApp_VoteExtensions(t *testing.T) {
+	baseappOpts := func(app *baseapp.BaseApp) {
+		app.SetExtendVoteHandler(func(sdk.Context, *abci.RequestExtendVote) (*abci.ResponseExtendVote, error) {
+			// here we would have a process to get the price from an external source
+			price := 10000000 + rand.Int63n(1000000)
+			ve := make([]byte, 8)
+			binary.BigEndian.PutUint64(ve, uint64(price))
+			return &abci.ResponseExtendVote{VoteExtension: ve}, nil
+		})
+
+		app.SetVerifyVoteExtensionHandler(func(_ sdk.Context, req *abci.RequestVerifyVoteExtension) (*abci.ResponseVerifyVoteExtension, error) {
+			vePrice := binary.BigEndian.Uint64(req.VoteExtension)
+			// here we would do some price validation, must not be 0 and not too high
+			if vePrice > 11000000 || vePrice == 0 {
+				// usually application should always return ACCEPT unless they really want to discard the entire vote
+				return &abci.ResponseVerifyVoteExtension{Status: abci.ResponseVerifyVoteExtension_REJECT}, nil
+			}
+
+			return &abci.ResponseVerifyVoteExtension{Status: abci.ResponseVerifyVoteExtension_ACCEPT}, nil
+		})
+
+		app.SetPrepareProposal(func(ctx sdk.Context, req *abci.RequestPrepareProposal) (*abci.ResponsePrepareProposal, error) {
+			txs := [][]byte{}
+
+			// add all VE as txs (in a real scenario we would need to check signatures too)
+			for _, v := range req.LocalLastCommit.Votes {
+				if len(v.VoteExtension) == 8 {
+					// pretend this is a way to check if the VE is valid
+					if binary.BigEndian.Uint64(v.VoteExtension) < 11000000 && binary.BigEndian.Uint64(v.VoteExtension) > 0 {
+						txs = append(txs, v.VoteExtension)
+					}
+				}
+			}
+
+			return &abci.ResponsePrepareProposal{Txs: txs}, nil
+		})
+
+		app.SetProcessProposal(func(ctx sdk.Context, req *abci.RequestProcessProposal) (*abci.ResponseProcessProposal, error) {
+			// here we check if the proposal is valid, mainly if the vote extensions appended to the txs are valid
+			for _, v := range req.Txs {
+				// pretend this is a way to check if the tx is actually a VE
+				if len(v) == 8 {
+					// pretend this is a way to check if the VE is valid
+					if binary.BigEndian.Uint64(v) > 11000000 || binary.BigEndian.Uint64(v) == 0 {
+						return &abci.ResponseProcessProposal{Status: abci.ResponseProcessProposal_REJECT}, nil
+					}
+				}
+			}
+
+			return &abci.ResponseProcessProposal{Status: abci.ResponseProcessProposal_ACCEPT}, nil
+		})
+
+		app.SetPreFinalizeBlockHook(func(ctx sdk.Context, req *abci.RequestFinalizeBlock) error {
+			count := uint64(0)
+			pricesSum := uint64(0)
+			for _, v := range req.Txs {
+				// pretend this is a way to check if the tx is actually a VE
+				if len(v) == 8 {
+					count++
+					pricesSum += binary.BigEndian.Uint64(v)
+				}
+			}
+
+			if count > 0 {
+				// we process the average price and store it in the context to make it available for FinalizeBlock
+				avgPrice := pricesSum / count
+				buf := make([]byte, 8)
+				binary.BigEndian.PutUint64(buf, avgPrice)
+				ctx.KVStore(capKey1).Set([]byte("avgPrice"), buf)
+			}
+
+			return nil
+		})
+	}
+
+	suite := NewBaseAppSuite(t, baseappOpts)
+
+	_, err := suite.baseApp.InitChain(&abci.RequestInitChain{
+		ConsensusParams: &cmtproto.ConsensusParams{
+			Abci: &cmtproto.ABCIParams{
+				VoteExtensionsEnableHeight: 1,
+			},
+		},
+	})
+	require.NoError(t, err)
+
+	allVEs := [][]byte{}
+	// simulate getting 10 vote extensions from 10 validators
+	for i := 0; i < 10; i++ {
+		ve, err := suite.baseApp.ExtendVote(context.TODO(), &abci.RequestExtendVote{Height: 1})
+		require.NoError(t, err)
+		allVEs = append(allVEs, ve.VoteExtension)
+	}
+
+	// add a couple of invalid vote extensions (in what regards to the check we are doing in VerifyVoteExtension/ProcessProposal)
+	// add a 0 price
+	ve := make([]byte, 8)
+	binary.BigEndian.PutUint64(ve, uint64(0))
+	allVEs = append(allVEs, ve)
+
+	// add a price too high
+	ve = make([]byte, 8)
+	binary.BigEndian.PutUint64(ve, uint64(13000000))
+	allVEs = append(allVEs, ve)
+
+	// verify all votes, only 10 should be accepted
+	successful := 0
+	for _, v := range allVEs {
+		res, err := suite.baseApp.VerifyVoteExtension(&abci.RequestVerifyVoteExtension{
+			Height:        1,
+			VoteExtension: v,
+		})
+		require.NoError(t, err)
+		if res.Status == abci.ResponseVerifyVoteExtension_ACCEPT {
+			successful++
+		}
+	}
+	require.Equal(t, 10, successful)
+
+	prepPropReq := &abci.RequestPrepareProposal{
+		Height: 1,
+		LocalLastCommit: abci.ExtendedCommitInfo{
+			Round: 0,
+			Votes: []abci.ExtendedVoteInfo{},
+		},
+	}
+
+	// add all VEs to the local last commit
+	for _, ve := range allVEs {
+		prepPropReq.LocalLastCommit.Votes = append(prepPropReq.LocalLastCommit.Votes, abci.ExtendedVoteInfo{VoteExtension: ve})
+	}
+
+	resp, err := suite.baseApp.PrepareProposal(prepPropReq)
+	require.NoError(t, err)
+	require.Len(t, resp.Txs, 10)
+
+	procPropRes, err := suite.baseApp.ProcessProposal(&abci.RequestProcessProposal{Height: 1, Txs: resp.Txs})
+	require.NoError(t, err)
+	require.Equal(t, abci.ResponseProcessProposal_ACCEPT, procPropRes.Status)
+
+	_, err = suite.baseApp.FinalizeBlock(&abci.RequestFinalizeBlock{Height: 1, Txs: resp.Txs})
+	require.NoError(t, err)
+
+	// Check if the average price was available in FinalizeBlock's context
+	avgPrice := getFinalizeBlockStateCtx(suite.baseApp).KVStore(capKey1).Get([]byte("avgPrice"))
+	require.NotNil(t, avgPrice)
+	require.GreaterOrEqual(t, binary.BigEndian.Uint64(avgPrice), uint64(10000000))
+	require.Less(t, binary.BigEndian.Uint64(avgPrice), uint64(11000000))
+
+	_, err = suite.baseApp.Commit()
+	require.NoError(t, err)
+
+	// check if avgPrice was committed
+	committedAvgPrice := suite.baseApp.NewContext(true).KVStore(capKey1).Get([]byte("avgPrice"))
+	require.Equal(t, avgPrice, committedAvgPrice)
+}

baseapp/baseapp.go

@@ -73,15 +73,16 @@ type BaseApp struct {
 	anteHandler sdk.AnteHandler // ante handler for fee and auth
 	postHandler sdk.PostHandler // post handler, optional, e.g. for tips
 
-	initChainer        sdk.InitChainer                // ABCI InitChain handler
-	beginBlocker       sdk.BeginBlocker               // (legacy ABCI) BeginBlock handler
-	endBlocker         sdk.EndBlocker                 // (legacy ABCI) EndBlock handler
-	processProposal    sdk.ProcessProposalHandler     // ABCI ProcessProposal handler
-	prepareProposal    sdk.PrepareProposalHandler     // ABCI PrepareProposal
-	extendVote         sdk.ExtendVoteHandler          // ABCI ExtendVote handler
-	verifyVoteExt      sdk.VerifyVoteExtensionHandler // ABCI VerifyVoteExtension handler
-	prepareCheckStater sdk.PrepareCheckStater         // logic to run during commit using the checkState
-	precommiter        sdk.Precommiter                // logic to run during commit using the deliverState
+	initChainer          sdk.InitChainer                // ABCI InitChain handler
+	beginBlocker         sdk.BeginBlocker               // (legacy ABCI) BeginBlock handler
+	endBlocker           sdk.EndBlocker                 // (legacy ABCI) EndBlock handler
+	processProposal      sdk.ProcessProposalHandler     // ABCI ProcessProposal handler
+	prepareProposal      sdk.PrepareProposalHandler     // ABCI PrepareProposal
+	extendVote           sdk.ExtendVoteHandler          // ABCI ExtendVote handler
+	verifyVoteExt        sdk.VerifyVoteExtensionHandler // ABCI VerifyVoteExtension handler
+	prepareCheckStater   sdk.PrepareCheckStater         // logic to run during commit using the checkState
+	precommiter          sdk.Precommiter                // logic to run during commit using the deliverState
+	preFinalizeBlockHook sdk.PreFinalizeBlockHook       // logic to run before FinalizeBlock
 
 	addrPeerFilter sdk.PeerFilter // filter peers by address and port
 	idPeerFilter   sdk.PeerFilter // filter peers by node ID
@@ -485,18 +486,6 @@ func (app *BaseApp) setState(mode execMode, header cmtproto.Header) {
 	}
 }
 
-// GetFinalizeBlockStateCtx returns the Context associated with the FinalizeBlock
-// state. This Context can be used to write data derived from processing vote
-// extensions to application state during ProcessProposal.
-//
-// NOTE:
-// - Do NOT use or write to state using this Context unless you intend for
-// that state to be committed.
-// - Do NOT use or write to state using this Context on the first block.
-func (app *BaseApp) GetFinalizeBlockStateCtx() sdk.Context {
-	return app.finalizeBlockState.ctx
-}
-
 // SetCircuitBreaker sets the circuit breaker for the BaseApp.
 // The circuit breaker is checked on every message execution to verify if a transaction should be executed or not.
 func (app *BaseApp) SetCircuitBreaker(cb CircuitBreaker) {

baseapp/options.go

@@ -180,6 +180,14 @@ func (app *BaseApp) SetPrecommiter(precommiter sdk.Precommiter) {
 	app.precommiter = precommiter
 }
 
+func (app *BaseApp) SetPreFinalizeBlockHook(hook sdk.PreFinalizeBlockHook) {
+	if app.sealed {
+		panic("SetPreFinalizeBlockHook() on sealed BaseApp")
+	}
+
+	app.preFinalizeBlockHook = hook
+}
+
 func (app *BaseApp) SetAnteHandler(ah sdk.AnteHandler) {
 	if app.sealed {
 		panic("SetAnteHandler() on sealed BaseApp")