fix: execute PreBlocker for upgrade module (#707)

app/app.go

@@ -264,7 +264,6 @@ func NewBabylonApp(
 		wasmOpts,
 		BlockedAddresses(),
 	)
-	app.setupUpgradeStoreLoaders()
 
 	/****  Module Options ****/
 
@@ -456,9 +455,6 @@ func NewBabylonApp(
 	// add test gRPC service for testing gRPC queries in isolation
 	testdata.RegisterQueryServer(app.GRPCQueryRouter(), testdata.QueryImpl{})
 
-	// set upgrade handler
-	app.setupUpgradeHandlers()
-
 	// create the simulation manager and define the order of the modules for deterministic simulations
 	//
 	// NOTE: this is not required apps that don't use the simulator for fuzz testing
@@ -509,7 +505,30 @@ func NewBabylonApp(
 		NewBtcValidationDecorator(btcConfig, &app.BtcCheckpointKeeper),
 	)
 
+	// set proposal extension
+	proposalHandler := checkpointing.NewProposalHandler(
+		logger, &app.CheckpointingKeeper, bApp.Mempool(), bApp)
+	proposalHandler.SetHandlers(bApp)
+
+	// set vote extension
+	voteExtHandler := checkpointing.NewVoteExtensionHandler(logger, &app.CheckpointingKeeper)
+	voteExtHandler.SetHandlers(bApp)
+
 	app.SetInitChainer(app.InitChainer)
+	app.SetPreBlocker(func(ctx sdk.Context, req *abci.RequestFinalizeBlock) (*sdk.ResponsePreBlock, error) {
+		// execute the existing PreBlocker
+		res, err := app.PreBlocker(ctx, req)
+		if err != nil {
+			return res, err
+		}
+		// execute checkpointing module's PreBlocker
+		// NOTE: this does not change the consensus parameter in `res`
+		ckptPreBlocker := proposalHandler.PreBlocker()
+		if _, err := ckptPreBlocker(ctx, req); err != nil {
+			return res, err
+		}
+		return res, nil
+	})
 	app.SetBeginBlocker(app.BeginBlocker)
 	app.SetEndBlocker(app.EndBlocker)
 	app.SetAnteHandler(anteHandler)
@@ -545,6 +564,10 @@ func NewBabylonApp(
 		_, _ = fmt.Fprintln(os.Stderr, err.Error())
 	}
 
+	// set upgrade handler and store loader for supporting software upgrade
+	app.setupUpgradeHandlers()
+	app.setupUpgradeStoreLoaders()
+
 	if loadLatest {
 		if err := app.LoadLatestVersion(); err != nil {
 			cmtos.Exit(err.Error())

app/keepers/keepers.go

@@ -66,7 +66,6 @@ import (
 	btclightclienttypes "github.com/babylonchain/babylon/x/btclightclient/types"
 	btcstakingkeeper "github.com/babylonchain/babylon/x/btcstaking/keeper"
 	btcstakingtypes "github.com/babylonchain/babylon/x/btcstaking/types"
-	"github.com/babylonchain/babylon/x/checkpointing"
 	checkpointingkeeper "github.com/babylonchain/babylon/x/checkpointing/keeper"
 	checkpointingtypes "github.com/babylonchain/babylon/x/checkpointing/types"
 	epochingkeeper "github.com/babylonchain/babylon/x/epoching/keeper"
@@ -244,15 +243,6 @@ func (ak *AppKeepers) InitKeepers(
 		epochingKeeper,
 	)
 
-	// set proposal extension
-	proposalHandler := checkpointing.NewProposalHandler(
-		logger, &checkpointingKeeper, bApp.Mempool(), bApp)
-	proposalHandler.SetHandlers(bApp)
-
-	// set vote extension
-	voteExtHandler := checkpointing.NewVoteExtensionHandler(logger, &checkpointingKeeper)
-	voteExtHandler.SetHandlers(bApp)
-
 	// register streaming services
 	if err := bApp.RegisterStreamingServices(appOpts, keys); err != nil {
 		panic(err)

app/upgrades/vanilla/upgrades.go

@@ -0,0 +1,61 @@
+package vanilla
+
+import (
+	"context"
+
+	store "cosmossdk.io/store/types"
+	upgradetypes "cosmossdk.io/x/upgrade/types"
+	"github.com/babylonchain/babylon/app/keepers"
+	"github.com/babylonchain/babylon/app/upgrades"
+	bbn "github.com/babylonchain/babylon/types"
+	btcstakingkeeper "github.com/babylonchain/babylon/x/btcstaking/keeper"
+	bstypes "github.com/babylonchain/babylon/x/btcstaking/types"
+	"github.com/btcsuite/btcd/btcec/v2"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
+)
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:          "vanilla",
+	CreateUpgradeHandler: CreateUpgradeHandler,
+	StoreUpgrades:        store.StoreUpgrades{},
+}
+
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	cfg module.Configurator,
+	_ upgrades.BaseAppParamManager,
+	keepers *keepers.AppKeepers,
+) upgradetypes.UpgradeHandler {
+	return func(context context.Context, _plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+
+		ctx := sdk.UnwrapSDKContext(context)
+
+		propVanilla(ctx, &keepers.AccountKeeper, &keepers.BTCStakingKeeper)
+
+		return mm.RunMigrations(ctx, cfg, fromVM)
+	}
+}
+
+func propVanilla(
+	ctx sdk.Context,
+	accountKeeper *authkeeper.AccountKeeper,
+	bsKeeper *btcstakingkeeper.Keeper,
+) {
+	// remove an account
+	allAccounts := accountKeeper.GetAllAccounts(ctx)
+	accountKeeper.RemoveAccount(ctx, allAccounts[len(allAccounts)-1])
+
+	// insert a FP
+	sk, err := btcec.NewPrivateKey()
+	if err != nil {
+		panic(err)
+	}
+	btcPK := bbn.NewBIP340PubKeyFromBTCPK(sk.PubKey())
+	fp := &bstypes.FinalityProvider{
+		Addr:  allAccounts[0].GetAddress().String(),
+		BtcPk: btcPK,
+	}
+	bsKeeper.SetFinalityProvider(ctx, fp)
+}

app/upgrades/types_test.go → app/upgrades/vanilla/upgrades_test.go

@@ -1,74 +1,26 @@
-package upgrades_test
+package vanilla_test
 
 import (
-	"context"
 	"fmt"
-	"math/rand"
 	"testing"
 	"time"
 
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/core/header"
-	store "cosmossdk.io/store/types"
 	"cosmossdk.io/x/upgrade"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
 	"github.com/babylonchain/babylon/app"
-	"github.com/babylonchain/babylon/app/keepers"
-	"github.com/babylonchain/babylon/app/upgrades"
-	"github.com/babylonchain/babylon/testutil/datagen"
-	btcstakingkeeper "github.com/babylonchain/babylon/x/btcstaking/keeper"
+	v1 "github.com/babylonchain/babylon/app/upgrades/vanilla"
 	bstypes "github.com/babylonchain/babylon/x/btcstaking/types"
 	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
-	"github.com/cosmos/cosmos-sdk/types/module"
-	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	"github.com/stretchr/testify/suite"
 )
 
 const (
 	DummyUpgradeHeight = 5
 )
 
-var Upgrade = upgrades.Upgrade{
-	UpgradeName:          "vanilla",
-	CreateUpgradeHandler: CreateUpgradeHandler,
-	StoreUpgrades:        store.StoreUpgrades{},
-}
-
-func CreateUpgradeHandler(
-	mm *module.Manager,
-	_ module.Configurator,
-	_ upgrades.BaseAppParamManager,
-	keepers *keepers.AppKeepers,
-) upgradetypes.UpgradeHandler {
-	return func(context context.Context, _plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
-		ctx := sdk.UnwrapSDKContext(context)
-
-		propVanilla(ctx, &keepers.AccountKeeper, &keepers.BTCStakingKeeper)
-
-		return vm, nil
-	}
-}
-
-func propVanilla(
-	ctx sdk.Context,
-	accountKeeper *authkeeper.AccountKeeper,
-	bsKeeper *btcstakingkeeper.Keeper,
-) {
-	r := rand.New(rand.NewSource(time.Now().Unix()))
-
-	// remove an account
-	allAccounts := accountKeeper.GetAllAccounts(ctx)
-	accountKeeper.RemoveAccount(ctx, allAccounts[len(allAccounts)-1])
-
-	// insert a FP
-	fp, err := datagen.GenRandomFinalityProvider(r)
-	if err != nil {
-		panic(err)
-	}
-	bsKeeper.SetFinalityProvider(ctx, fp)
-}
-
 type UpgradeTestSuite struct {
 	suite.Suite
 
@@ -78,6 +30,10 @@ type UpgradeTestSuite struct {
 }
 
 func (s *UpgradeTestSuite) SetupTest() {
+	// add the upgrade plan
+	app.Upgrades = append(app.Upgrades, v1.Upgrade)
+
+	// set up app
 	s.app = app.Setup(s.T(), false)
 	s.ctx = s.app.BaseApp.NewContextLegacy(false, tmproto.Header{Height: 1, ChainID: "babylon-1", Time: time.Now().UTC()})
 	s.preModule = upgrade.NewAppModule(s.app.UpgradeKeeper, s.app.AccountKeeper.AddressCodec())
@@ -106,23 +62,16 @@ func (s *UpgradeTestSuite) TestUpgradePayments() {
 			func() {
 				// inject upgrade plan
 				s.ctx = s.ctx.WithBlockHeight(DummyUpgradeHeight - 1)
-				plan := upgradetypes.Plan{Name: Upgrade.UpgradeName, Height: DummyUpgradeHeight}
-				s.app.UpgradeKeeper.SetUpgradeHandler(
-					Upgrade.UpgradeName,
-					Upgrade.CreateUpgradeHandler(
-						s.app.ModuleManager,
-						nil,
-						nil,
-						s.app.AppKeepers,
-					),
-				)
-
-				// run upgrade
+				plan := upgradetypes.Plan{Name: v1.Upgrade.UpgradeName, Height: DummyUpgradeHeight}
 				err := s.app.UpgradeKeeper.ScheduleUpgrade(s.ctx, plan)
 				s.NoError(err)
-				_, err = s.app.UpgradeKeeper.GetUpgradePlan(s.ctx)
+
+				// ensure upgrade plan exists
+				actualPlan, err := s.app.UpgradeKeeper.GetUpgradePlan(s.ctx)
 				s.NoError(err)
+				s.Equal(plan, actualPlan)
 
+				// execute upgrade
 				s.ctx = s.ctx.WithHeaderInfo(header.Info{Height: DummyUpgradeHeight, Time: s.ctx.BlockTime().Add(time.Second)}).WithBlockHeight(DummyUpgradeHeight)
 				s.NotPanics(func() {
 					_, err := s.preModule.PreBlock(s.ctx)

x/checkpointing/module.go

@@ -2,10 +2,11 @@ package checkpointing
 
 import (
 	"context"
-	"cosmossdk.io/core/appmodule"
 	"encoding/json"
 	"fmt"
 
+	"cosmossdk.io/core/appmodule"
+
 	"github.com/gorilla/mux"
 	"github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/spf13/cobra"

x/checkpointing/proposal.go

@@ -45,7 +45,6 @@ func NewProposalHandler(
 func (h *ProposalHandler) SetHandlers(bApp *baseapp.BaseApp) {
 	bApp.SetPrepareProposal(h.PrepareProposal())
 	bApp.SetProcessProposal(h.ProcessProposal())
-	bApp.SetPreBlocker(h.PreBlocker())
 }
 
 // PrepareProposal examines the vote extensions from the previous block, accumulates
@@ -329,14 +328,14 @@ func (h *ProposalHandler) ProcessProposal() sdk.ProcessProposalHandler {
 	}
 }
 
-// PreBlocker extracts the checkpoint from the injected tx and stores it in
-// the application
+// PreBlocker extracts the checkpoint from the injected tx and stores it in the application
 // no more validation is needed as it is already done in ProcessProposal
+// NOTE: this is appended to the existing PreBlocker in BabylonApp at app.go
 func (h *ProposalHandler) PreBlocker() sdk.PreBlocker {
 	return func(ctx sdk.Context, req *abci.RequestFinalizeBlock) (*sdk.ResponsePreBlock, error) {
-		k := h.ckptKeeper
 		res := &sdk.ResponsePreBlock{}
 
+		k := h.ckptKeeper
 		epoch := k.GetEpoch(ctx)
 		// BLS signatures are sent in the last block of the previous epoch,
 		// so they should be aggregated in the first block of the new epoch