refactor: improve `x/upgrade` app wiring (part of upgrade audit) (backport #14216) (#14306)

Co-authored-by: Julien Robert <julien@rbrt.fr>

CHANGELOG.md

@@ -37,8 +37,6 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ## [Unreleased]
 
-## [v0.47.0-alpha3](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.47.0-alpha3) - 2022-12-13
-
 ### Features
 
 * (client) [#14051](https://github.com/cosmos/cosmos-sdk/pull/14051) Add `--grpc` client option.
@@ -48,6 +46,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 * (types) [#14163](https://github.com/cosmos/cosmos-sdk/pull/14163) Refactor `(coins Coins) Validate()` to avoid unnecessary map.
 
+### API Breaking Changes
+
+* (x/upgrade) [#14216](https://github.com/cosmos/cosmos-sdk/pull/14216) Change upgrade keeper receiver to upgrade keeper pointers.
+
 ### State Machine Breaking
 
 * (x/gov) [#14214](https://github.com/cosmos/cosmos-sdk/pull/14214) Fix gov v0.46 migration to v1 votes.

RELEASE_NOTES.md

@@ -1,4 +1,4 @@
-# Cosmos SDK v0.47.0-alpha3 Release Notes
+# Cosmos SDK v0.47.0-beta1 Release Notes
 
 There is no release notes for pre-releases. 
 
@@ -6,4 +6,4 @@ Please see the [CHANGELOG](https://github.com/cosmos/cosmos-sdk/blob/release/v0.
 Refer to the [UPGRADING.md](https://github.com/cosmos/cosmos-sdk/blob/release/v0.47.x/UPGRADING.md) for upgrading your application.
 
 Full Commit History (previous version): https://github.com/cosmos/cosmos-sdk/compare/release/v0.46.x...release/v0.47.x
-Full Commit History (`alpha2..alpha3`): https://github.com/cosmos/cosmos-sdk/compare/v0.47.0-alpha2...v0.47.0-alpha3
+Full Commit History (`alpha2..beta1`): https://github.com/cosmos/cosmos-sdk/compare/v0.47.0-alpha2...v0.47.0-beta1

core/sonar-project.properties

@@ -1,7 +1,7 @@
 sonar.projectKey=cosmos-sdk-core
 sonar.organization=cosmos
 
-sonar.projectName=Cosmos SDK Core
+sonar.projectName=Cosmos SDK - Core
 sonar.project.monorepo.enabled=true
 
 sonar.sources=.

runtime/app.go

@@ -48,6 +48,9 @@ type App struct {
 	baseAppOptions    []BaseAppOption
 	msgServiceRouter  *baseapp.MsgServiceRouter
 	appConfig         *appv1alpha1.Config
+	// initChainer is the init chainer function defined by the app config.
+	// this is only required if the chain wants to add special InitChainer logic.
+	initChainer sdk.InitChainer
 }
 
 // RegisterModules registers the provided modules with the module manager and
@@ -82,7 +85,9 @@ func (a *App) Load(loadLatest bool) error {
 
 	if len(a.config.InitGenesis) != 0 {
 		a.ModuleManager.SetOrderInitGenesis(a.config.InitGenesis...)
-		a.SetInitChainer(a.InitChainer)
+		if a.initChainer == nil {
+			a.SetInitChainer(a.InitChainer)
+		}
 	}
 
 	if len(a.config.ExportGenesis) != 0 {
@@ -165,10 +170,12 @@ func (a *App) RegisterTendermintService(clientCtx client.Context) {
 	)
 }
 
+// RegisterNodeService registers the node gRPC service on the app gRPC router.
 func (a *App) RegisterNodeService(clientCtx client.Context) {
 	nodeservice.RegisterNodeService(clientCtx, a.GRPCQueryRouter())
 }
 
+// Configurator returns the app's configurator.
 func (a *App) Configurator() module.Configurator {
 	return a.configurator
 }
@@ -183,11 +190,18 @@ func (a *App) DefaultGenesis() map[string]json.RawMessage {
 	return a.basicManager.DefaultGenesis(a.cdc)
 }
 
-// GetStoreKeys returns all the keys stored store keys.
+// GetStoreKeys returns all the stored store keys.
 func (a *App) GetStoreKeys() []storetypes.StoreKey {
 	return a.storeKeys
 }
 
+// SetInitChainer sets the init chainer function
+// It wraps `BaseApp.SetInitChainer` to allow setting a custom init chainer from an app.
+func (a *App) SetInitChainer(initChainer sdk.InitChainer) {
+	a.initChainer = initChainer
+	a.BaseApp.SetInitChainer(initChainer)
+}
+
 // UnsafeFindStoreKey fetches a registered StoreKey from the App in linear time.
 //
 // NOTE: This should only be used in testing.

simapp/app.go

@@ -4,7 +4,6 @@ package simapp
 
 import (
 	"encoding/json"
-	"fmt"
 	"io"
 	"os"
 	"path/filepath"
@@ -184,7 +183,7 @@ type SimApp struct {
 	DistrKeeper           distrkeeper.Keeper
 	GovKeeper             govkeeper.Keeper
 	CrisisKeeper          *crisiskeeper.Keeper
-	UpgradeKeeper         upgradekeeper.Keeper
+	UpgradeKeeper         *upgradekeeper.Keeper
 	ParamsKeeper          paramskeeper.Keeper
 	AuthzKeeper           authzkeeper.Keeper
 	EvidenceKeeper        evidencekeeper.Keeper
@@ -265,7 +264,7 @@ func NewSimApp(
 
 	// load state streaming if enabled
 	if _, _, err := streaming.LoadStreamingServices(bApp, appOpts, appCodec, logger, keys); err != nil {
-		fmt.Printf("failed to load state streaming: %s", err)
+		logger.Error("failed to load state streaming", "err", err)
 		os.Exit(1)
 	}
 
@@ -360,15 +359,15 @@ func NewSimApp(
 		app.StakingKeeper, app.MsgServiceRouter(), govConfig, authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
+	// Set legacy router for backwards compatibility with gov v1beta1
+	govKeeper.SetLegacyRouter(govRouter)
+
 	app.GovKeeper = *govKeeper.SetHooks(
 		govtypes.NewMultiGovHooks(
 		// register the governance hooks
 		),
 	)
 
-	// Set legacy router for backwards compatibility with gov v1beta1
-	govKeeper.SetLegacyRouter(govRouter)
-
 	app.NFTKeeper = nftkeeper.NewKeeper(keys[nftkeeper.StoreKey], appCodec, app.AccountKeeper, app.BankKeeper)
 
 	// create evidence keeper with router
@@ -514,7 +513,7 @@ func NewSimApp(
 
 	if loadLatest {
 		if err := app.LoadLatestVersion(); err != nil {
-			fmt.Println(err.Error())
+			logger.Error("error on loading last version", "err", err)
 			os.Exit(1)
 		}
 	}

simapp/app_config.go

@@ -148,7 +148,7 @@ var (
 							KvStoreKey: "acc",
 						},
 					},
-					// InitGenesis: genesisModuleOrder,
+					InitGenesis: genesisModuleOrder,
 					// When ExportGenesis is not specified, the export genesis module order
 					// is equal to the init genesis order
 					// ExportGenesis: genesisModuleOrder,

simapp/app_v2.go

@@ -4,12 +4,10 @@ package simapp
 
 import (
 	_ "embed"
-	"fmt"
 	"io"
 	"os"
 	"path/filepath"
 
-	abci "github.com/tendermint/tendermint/abci/types"
 	"github.com/tendermint/tendermint/libs/log"
 	dbm "github.com/tendermint/tm-db"
 
@@ -27,7 +25,6 @@ import (
 	"github.com/cosmos/cosmos-sdk/store/streaming"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	"github.com/cosmos/cosmos-sdk/testutil/testdata_pulsar"
-	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	"github.com/cosmos/cosmos-sdk/x/auth"
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
@@ -136,7 +133,7 @@ type SimApp struct {
 	DistrKeeper           distrkeeper.Keeper
 	GovKeeper             *govkeeper.Keeper
 	CrisisKeeper          *crisiskeeper.Keeper
-	UpgradeKeeper         upgradekeeper.Keeper
+	UpgradeKeeper         *upgradekeeper.Keeper
 	ParamsKeeper          paramskeeper.Keeper
 	AuthzKeeper           authzkeeper.Keeper
 	EvidenceKeeper        evidencekeeper.Keeper
@@ -250,18 +247,12 @@ func NewSimApp(
 
 	// load state streaming if enabled
 	if _, _, err := streaming.LoadStreamingServices(app.App.BaseApp, appOpts, app.appCodec, logger, app.kvStoreKeys()); err != nil {
-		fmt.Printf("failed to load state streaming: %s", err)
+		logger.Error("failed to load state streaming", "err", err)
 		os.Exit(1)
 	}
 
 	/****  Module Options ****/
 
-	// Set upgrade module options
-	app.UpgradeKeeper.SetVersionSetter(app.BaseApp)
-
-	app.ModuleManager.SetOrderInitGenesis(genesisModuleOrder...)
-	app.ModuleManager.SetOrderExportGenesis(genesisModuleOrder...)
-
 	app.ModuleManager.RegisterInvariants(app.CrisisKeeper)
 
 	// RegisterUpgradeHandlers is used for registering any on-chain upgrades.
@@ -281,8 +272,16 @@ func NewSimApp(
 
 	app.sm.RegisterStoreDecoders()
 
-	// initialize BaseApp
+	// A custom InitChainer can be set if extra pre-init-genesis logic is required.
-	app.SetInitChainer(app.InitChainer)
+	// By default, when using app wiring enabled module, this is not required.
+	// For instance, the upgrade module will set automatically the module version map in its init genesis thanks to app wiring.
+	// However, when registering a module manually (i.e. that does not support app wiring), the module version map
+	// must be set manually as follow. The upgrade module will de-duplicate the module version map.
+	//
+	// app.SetInitChainer(func(ctx sdk.Context, req abci.RequestInitChain) abci.ResponseInitChain {
+	// 	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.ModuleManager.GetVersionMap())
+	// 	return app.App.InitChainer(ctx, req)
+	// })
 
 	if err := app.Load(loadLatest); err != nil {
 		panic(err)
@@ -294,12 +293,6 @@ func NewSimApp(
 // Name returns the name of the App
 func (app *SimApp) Name() string { return app.BaseApp.Name() }
 
-// InitChainer application update at chain initialization
-func (app *SimApp) InitChainer(ctx sdk.Context, req abci.RequestInitChain) abci.ResponseInitChain {
-	app.UpgradeKeeper.SetModuleVersionMap(ctx, app.ModuleManager.GetVersionMap())
-	return app.App.InitChainer(ctx, req)
-}
-
 // LegacyAmino returns SimApp's amino codec.
 //
 // NOTE: This is solely to be used for testing purposes as it may be desirable

simapp/state.go

@@ -144,7 +144,6 @@ func AppStateRandomizedFn(
 	accs []simtypes.Account, genesisTimestamp time.Time, appParams simtypes.AppParams,
 ) (json.RawMessage, []simtypes.Account) {
 	numAccs := int64(len(accs))
-	// TODO - in case runtime.RegisterModules(...) is used, the genesis state of the module won't be reflected here
 	genesisState := ModuleBasics.DefaultGenesis(cdc)
 
 	// generate a random amount of initial stake coins and a random initial

tests/e2e/upgrade/suite.go

@@ -13,7 +13,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
-func NewE2ETestSuite(cfg network.Config, keeper keeper.Keeper, ctx sdk.Context) *E2ETestSuite {
+func NewE2ETestSuite(cfg network.Config, keeper *keeper.Keeper, ctx sdk.Context) *E2ETestSuite {
 	return &E2ETestSuite{
 		cfg:           cfg,
 		upgradeKeeper: keeper,
@@ -24,7 +24,7 @@ func NewE2ETestSuite(cfg network.Config, keeper keeper.Keeper, ctx sdk.Context)
 type E2ETestSuite struct {
 	suite.Suite
 
-	upgradeKeeper keeper.Keeper
+	upgradeKeeper *keeper.Keeper
 	cfg           network.Config
 	network       *network.Network
 	ctx           sdk.Context

x/upgrade/abci.go

@@ -20,7 +20,7 @@ import (
 // The purpose is to ensure the binary is switched EXACTLY at the desired block, and to allow
 // a migration to be executed if needed upon this switch (migration defined in the new binary)
 // skipUpgradeHeightArray is a set of block heights for which the upgrade must be skipped
-func BeginBlocker(k keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
+func BeginBlocker(k *keeper.Keeper, ctx sdk.Context, _ abci.RequestBeginBlock) {
 	defer telemetry.ModuleMeasureSince(types.ModuleName, time.Now(), telemetry.MetricKeyBeginBlocker)
 
 	plan, found := k.GetUpgradePlan(ctx)

x/upgrade/abci_test.go

@@ -32,7 +32,7 @@ type TestSuite struct {
 	suite.Suite
 
 	module  module.BeginBlockAppModule
-	keeper  keeper.Keeper
+	keeper  *keeper.Keeper
 	handler govtypesv1beta1.Handler
 	ctx     sdk.Context
 	baseApp *baseapp.BaseApp
@@ -516,19 +516,19 @@ func TestDowngradeVerification(t *testing.T) {
 	})
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
-	testCases := map[string]struct{
+	testCases := map[string]struct {
-		preRun      func(keeper.Keeper, sdk.Context, string)
+		preRun      func(*keeper.Keeper, sdk.Context, string)
 		expectPanic bool
 	}{
 		"valid binary": {
-			preRun: func(k keeper.Keeper, ctx sdk.Context, name string) {
+			preRun: func(k *keeper.Keeper, ctx sdk.Context, name string) {
 				k.SetUpgradeHandler(planName, func(ctx sdk.Context, plan types.Plan, vm module.VersionMap) (module.VersionMap, error) {
 					return vm, nil
 				})
 			},
 		},
 		"downgrade with an active plan": {
-			preRun: func(k keeper.Keeper, ctx sdk.Context, name string) {
+			preRun: func(k *keeper.Keeper, ctx sdk.Context, name string) {
 				handler := upgrade.NewSoftwareUpgradeProposalHandler(k)
 				err := handler(ctx, &types.SoftwareUpgradeProposal{Title: "test", Plan: types.Plan{Name: "another" + planName, Height: ctx.BlockHeight() + 1}})
 				require.NoError(t, err, name)

x/upgrade/handler.go

@@ -13,7 +13,7 @@ import (
 // to abort a previously voted upgrade.
 //
 //nolint:staticcheck // we are intentionally using a deprecated proposal here.
-func NewSoftwareUpgradeProposalHandler(k keeper.Keeper) govtypes.Handler {
+func NewSoftwareUpgradeProposalHandler(k *keeper.Keeper) govtypes.Handler {
 	return func(ctx sdk.Context, content govtypes.Content) error {
 		switch c := content.(type) {
 		case *types.SoftwareUpgradeProposal:
@@ -29,12 +29,12 @@ func NewSoftwareUpgradeProposalHandler(k keeper.Keeper) govtypes.Handler {
 }
 
 //nolint:staticcheck // we are intentionally using a deprecated proposal here.
-func handleSoftwareUpgradeProposal(ctx sdk.Context, k keeper.Keeper, p *types.SoftwareUpgradeProposal) error {
+func handleSoftwareUpgradeProposal(ctx sdk.Context, k *keeper.Keeper, p *types.SoftwareUpgradeProposal) error {
 	return k.ScheduleUpgrade(ctx, p.Plan)
 }
 
 //nolint:staticcheck // we are intentionally using a deprecated proposal here.
-func handleCancelSoftwareUpgradeProposal(ctx sdk.Context, k keeper.Keeper, _ *types.CancelSoftwareUpgradeProposal) error {
+func handleCancelSoftwareUpgradeProposal(ctx sdk.Context, k *keeper.Keeper, _ *types.CancelSoftwareUpgradeProposal) error {
 	k.ClearUpgradePlan(ctx)
 	return nil
 }

x/upgrade/keeper/grpc_query_test.go

@@ -23,7 +23,7 @@ import (
 type UpgradeTestSuite struct {
 	suite.Suite
 
-	upgradeKeeper keeper.Keeper
+	upgradeKeeper *keeper.Keeper
 	ctx           sdk.Context
 	queryClient   types.QueryClient
 	encCfg        moduletestutil.TestEncodingConfig

x/upgrade/keeper/keeper.go

@@ -35,6 +35,7 @@ type Keeper struct {
 	versionSetter      xp.ProtocolVersionSetter        // implements setting the protocol version field on BaseApp
 	downgradeVerified  bool                            // tells if we've already sanity checked that this binary version isn't being used against an old state.
 	authority          string                          // the address capable of executing and cancelling an upgrade. Usually the gov module account
+	initVersionMap     module.VersionMap               // the module version map at init genesis
 }
 
 // NewKeeper constructs an upgrade Keeper which requires the following arguments:
@@ -43,8 +44,8 @@ type Keeper struct {
 // cdc - the app-wide binary codec
 // homePath - root directory of the application's config
 // vs - the interface implemented by baseapp which allows setting baseapp's protocol version field
-func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, vs xp.ProtocolVersionSetter, authority string) Keeper {
+func NewKeeper(skipUpgradeHeights map[int64]bool, storeKey storetypes.StoreKey, cdc codec.BinaryCodec, homePath string, vs xp.ProtocolVersionSetter, authority string) *Keeper {
-	return Keeper{
+	return &Keeper{
 		homePath:           homePath,
 		skipUpgradeHeights: skipUpgradeHeights,
 		storeKey:           storeKey,
@@ -65,6 +66,18 @@ func (k *Keeper) GetVersionSetter() xp.ProtocolVersionSetter {
 	return k.versionSetter
 }
 
+// SetInitVersionMap sets the initial version map.
+// This is only used in app wiring and should not be used in any other context.
+func (k *Keeper) SetInitVersionMap(vm module.VersionMap) {
+	k.initVersionMap = vm
+}
+
+// GetInitVersionMap gets the initial version map
+// This is only used in upgrade InitGenesis and should not be used in any other context.
+func (k *Keeper) GetInitVersionMap() module.VersionMap {
+	return k.initVersionMap
+}
+
 // SetUpgradeHandler sets an UpgradeHandler for the upgrade specified by name. This handler will be called when the upgrade
 // with this name is applied. In order for an upgrade with the given name to proceed, a handler for this upgrade
 // must be set even if it is a no-op function.

x/upgrade/keeper/keeper_test.go

@@ -28,7 +28,7 @@ type KeeperTestSuite struct {
 
 	key           *storetypes.KVStoreKey
 	baseApp       *baseapp.BaseApp
-	upgradeKeeper keeper.Keeper
+	upgradeKeeper *keeper.Keeper
 	homeDir       string
 	ctx           sdk.Context
 	msgSrvr       types.MsgServer

x/upgrade/keeper/migrations.go

@@ -11,11 +11,11 @@ import (
 
 // Migrator is a struct for handling in-place store migrations.
 type Migrator struct {
-	keeper Keeper
+	keeper *Keeper
 }
 
 // NewMigrator returns a new Migrator.
-func NewMigrator(keeper Keeper) Migrator {
+func NewMigrator(keeper *Keeper) Migrator {
 	return Migrator{keeper: keeper}
 }
 

x/upgrade/keeper/msg_server.go

@@ -10,12 +10,12 @@ import (
 )
 
 type msgServer struct {
-	Keeper
+	*Keeper
 }
 
 // NewMsgServerImpl returns an implementation of the upgrade MsgServer interface
 // for the provided Keeper.
-func NewMsgServerImpl(k Keeper) types.MsgServer {
+func NewMsgServerImpl(k *Keeper) types.MsgServer {
 	return &msgServer{
 		Keeper: k,
 	}

x/upgrade/module.go

@@ -5,20 +5,21 @@ import (
 	"encoding/json"
 	"fmt"
 
-	govv1beta1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1beta1"
-
 	gwruntime "github.com/grpc-ecosystem/grpc-gateway/runtime"
 	"github.com/spf13/cast"
 	"github.com/spf13/cobra"
 	abci "github.com/tendermint/tendermint/abci/types"
 
+	modulev1 "cosmossdk.io/api/cosmos/upgrade/module/v1"
 	"cosmossdk.io/core/appmodule"
 	"cosmossdk.io/depinject"
 
+	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/client"
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/codec"
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
+	"github.com/cosmos/cosmos-sdk/runtime"
 	"github.com/cosmos/cosmos-sdk/server"
 	servertypes "github.com/cosmos/cosmos-sdk/server/types"
 	store "github.com/cosmos/cosmos-sdk/store/types"
@@ -26,9 +27,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/types/module"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
-
+	govv1beta1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1beta1"
-	modulev1 "cosmossdk.io/api/cosmos/upgrade/module/v1"
-
 	"github.com/cosmos/cosmos-sdk/x/upgrade/client/cli"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/keeper"
 	"github.com/cosmos/cosmos-sdk/x/upgrade/types"
@@ -85,11 +84,11 @@ func (b AppModuleBasic) RegisterInterfaces(registry codectypes.InterfaceRegistry
 // AppModule implements the sdk.AppModule interface
 type AppModule struct {
 	AppModuleBasic
-	keeper keeper.Keeper
+	keeper *keeper.Keeper
 }
 
 // NewAppModule creates a new AppModule object
-func NewAppModule(keeper keeper.Keeper) AppModule {
+func NewAppModule(keeper *keeper.Keeper) AppModule {
 	return AppModule{
 		AppModuleBasic: AppModuleBasic{},
 		keeper:         keeper,
@@ -120,7 +119,20 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 }
 
 // InitGenesis is ignored, no sense in serializing future upgrades
-func (am AppModule) InitGenesis(_ sdk.Context, _ codec.JSONCodec, _ json.RawMessage) []abci.ValidatorUpdate {
+func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONCodec, _ json.RawMessage) []abci.ValidatorUpdate {
+	// set version map automatically if available
+	if versionMap := am.keeper.GetInitVersionMap(); versionMap != nil {
+		// chains can still use a custom init chainer for setting the version map
+		// this means that we need to combine the manually wired modules version map with app wiring enabled modules version map
+		for name, version := range am.keeper.GetModuleVersionMap(ctx) {
+			if _, ok := versionMap[name]; !ok {
+				versionMap[name] = version
+			}
+		}
+
+		am.keeper.SetModuleVersionMap(ctx, versionMap)
+	}
+
 	return []abci.ValidatorUpdate{}
 }
 
@@ -156,6 +168,7 @@ func (am AppModule) BeginBlock(ctx sdk.Context, req abci.RequestBeginBlock) {
 func init() {
 	appmodule.Register(&modulev1.Module{},
 		appmodule.Provide(ProvideModule),
+		appmodule.Invoke(PopulateVersionMap),
 	)
 }
 
@@ -172,9 +185,10 @@ type UpgradeInputs struct {
 type UpgradeOutputs struct {
 	depinject.Out
 
-	UpgradeKeeper keeper.Keeper
+	UpgradeKeeper *keeper.Keeper
 	Module        appmodule.AppModule
 	GovHandler    govv1beta1.HandlerRoute
+	BaseAppOption runtime.BaseAppOption
 }
 
 func ProvideModule(in UpgradeInputs) UpgradeOutputs {
@@ -199,8 +213,19 @@ func ProvideModule(in UpgradeInputs) UpgradeOutputs {
 
 	// set the governance module account as the authority for conducting upgrades
 	k := keeper.NewKeeper(skipUpgradeHeights, in.Key, in.Cdc, homePath, nil, authority.String())
+	baseappOpt := func(app *baseapp.BaseApp) {
+		k.SetVersionSetter(app)
+	}
 	m := NewAppModule(k)
 	gh := govv1beta1.HandlerRoute{RouteKey: types.RouterKey, Handler: NewSoftwareUpgradeProposalHandler(k)}
 
-	return UpgradeOutputs{UpgradeKeeper: k, Module: m, GovHandler: gh}
+	return UpgradeOutputs{UpgradeKeeper: k, Module: m, GovHandler: gh, BaseAppOption: baseappOpt}
+}
+
+func PopulateVersionMap(upgradeKeeper *keeper.Keeper, modules map[string]appmodule.AppModule) {
+	if upgradeKeeper == nil {
+		return
+	}
+
+	upgradeKeeper.SetInitVersionMap(module.NewManagerFromMap(modules).GetVersionMap())
 }