From f1b2d8da751f11c7763af49cedf045e2fb4dff18 Mon Sep 17 00:00:00 2001
From: Alex Gartner <alexg@zetachain.com>
Date: Thu, 9 May 2024 13:10:36 -0700
Subject: [PATCH] fix upgrade from version without InitChainer

---
 app/setup_handlers.go | 31 ++++++++++++++++++++++++++++---
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/app/setup_handlers.go b/app/setup_handlers.go
index 48fec6fbf6..bc6dadfe3b 100644
--- a/app/setup_handlers.go
+++ b/app/setup_handlers.go
@@ -3,13 +3,14 @@ package app
 import (
 	"os"
 
+	"golang.org/x/exp/slices"
+
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	storetypes "github.com/cosmos/cosmos-sdk/store/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
-	"github.com/cosmos/cosmos-sdk/x/consensus"
 	consensustypes "github.com/cosmos/cosmos-sdk/x/consensus/types"
 	crisistypes "github.com/cosmos/cosmos-sdk/x/crisis/types"
 	distrtypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
@@ -53,6 +54,16 @@ func SetupHandlers(app *App) {
 		}
 	}
 	baseAppLegacySS := app.ParamsKeeper.Subspace(baseapp.Paramspace).WithKeyTable(paramstypes.ConsensusParamsKeyTable())
+	needsForcedMigration := []string{
+		authtypes.ModuleName,
+		banktypes.ModuleName,
+		stakingtypes.ModuleName,
+		distrtypes.ModuleName,
+		slashingtypes.ModuleName,
+		govtypes.ModuleName,
+		crisistypes.ModuleName,
+		emissionstypes.ModuleName,
+	}
 	allUpgrades := upgradeTracker{
 		upgrades: []upgradeTrackerItem{
 			{
@@ -65,8 +76,22 @@ func SetupHandlers(app *App) {
 					// https://docs.cosmos.network/main/build/migrations/upgrading#xconsensus
 					baseapp.MigrateParams(ctx, baseAppLegacySS, &app.ConsensusParamsKeeper)
 
-					// skip consensus genesis
-					app.mm.Modules[consensustypes.ModuleName] = consensus.ConsensusVersion
+					// empty version map happens when upgrading from old versions which did not correctly call
+					// app.UpgradeKeeper.SetModuleVersionMap(ctx, app.mm.GetVersionMap()) in InitChainer.
+					// we must populate the version map if we detect this scenario
+					//
+					// this will only happen on the first upgrade. mainnet and testnet will not require this condition.
+					if len(vm) == 0 {
+						for m, mb := range app.mm.Modules {
+							if module, ok := mb.(module.HasConsensusVersion); ok {
+								if slices.Contains(needsForcedMigration, m) {
+									vm[m] = module.ConsensusVersion() - 1
+								} else {
+									vm[m] = module.ConsensusVersion()
+								}
+							}
+						}
+					}
 					return vm, nil
 				},
 			},