From 47b901a114717953ab6493f36d2ff8aacbb37748 Mon Sep 17 00:00:00 2001
From: catShaark <gundamaster5@gmail.com>
Date: Wed, 19 Jul 2023 21:29:12 +0700
Subject: [PATCH] fix: upgrade handler and store loader

---
 app/app.go                       | 38 ++++++++++++++++++++------------
 app/upgrades/v2_2_5/constants.go | 15 +++++++++++++
 app/upgrades/v2_2_5/upgrades.go  | 19 ++++++++++++++++
 3 files changed, 58 insertions(+), 14 deletions(-)
 create mode 100644 app/upgrades/v2_2_5/constants.go
 create mode 100644 app/upgrades/v2_2_5/upgrades.go

diff --git a/app/app.go b/app/app.go
index b3f552b9..c510417c 100644
--- a/app/app.go
+++ b/app/app.go
@@ -139,6 +139,7 @@ import (
 	// Upgrade Handler
 	upgrades "github.com/White-Whale-Defi-Platform/migaloo-chain/v2/app/upgrades"
 	v2 "github.com/White-Whale-Defi-Platform/migaloo-chain/v2/app/upgrades/v2"
+	v2_2_5 "github.com/White-Whale-Defi-Platform/migaloo-chain/v2/app/upgrades/v2_2_5"
 )
 
 const (
@@ -157,7 +158,7 @@ var (
 	// https://github.com/CosmWasm/wasmd/blob/02a54d33ff2c064f3539ae12d75d027d9c665f05/x/wasm/internal/types/proposal.go#L28-L34
 	EnableSpecificProposals = ""
 
-	Upgrades = []upgrades.Upgrade{v2.Upgrade}
+	Upgrades = []upgrades.Upgrade{v2.Upgrade, v2_2_5.Upgrade}
 )
 
 // GetEnabledProposals parses the ProposalsEnabled / EnableSpecificProposals values to
@@ -1042,21 +1043,30 @@ func RegisterSwaggerAPI(rtr *mux.Router) {
 
 // Setup Upgrade Handler
 func (app *MigalooApp) setupUpgradeHandlers(cfg module.Configurator) {
-	for _, upgrade := range Upgrades {
-		upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
-		if err != nil {
-			panic(fmt.Sprintf("failed to read upgrade info from disk %s", err))
-		}
+	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
+	if err != nil {
+		panic(fmt.Sprintf("failed to read upgrade info from disk %s", err))
+	}
 
-		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &upgrade.StoreUpgrades))
+	if app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
+		return
+	}
 
-		app.UpgradeKeeper.SetUpgradeHandler(
-			upgrade.UpgradeName,
-			upgrade.CreateUpgradeHandler(
-				app.mm,
-				cfg,
-			),
-		)
+	currentHeight := app.CommitMultiStore().LastCommitID().Version
+
+	for _, upgrade := range Upgrades {
+		if upgradeInfo.Name == upgrade.UpgradeName {
+			app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, &upgrade.StoreUpgrades))
+		}
+		if upgradeInfo.Height == currentHeight+1 {
+			app.UpgradeKeeper.SetUpgradeHandler(
+				upgrade.UpgradeName,
+				upgrade.CreateUpgradeHandler(
+					app.mm,
+					cfg,
+				),
+			)
+		}
 	}
 }
 
diff --git a/app/upgrades/v2_2_5/constants.go b/app/upgrades/v2_2_5/constants.go
new file mode 100644
index 00000000..dedf8b2d
--- /dev/null
+++ b/app/upgrades/v2_2_5/constants.go
@@ -0,0 +1,15 @@
+package v2_2_5
+
+import (
+	"github.com/White-Whale-Defi-Platform/migaloo-chain/v2/app/upgrades"
+	store "github.com/cosmos/cosmos-sdk/store/types"
+)
+
+// UpgradeName defines the on-chain upgrade name for the Migaloo v2 upgrade.
+const UpgradeName = "v2.2.5"
+
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:          UpgradeName,
+	CreateUpgradeHandler: CreateUpgradeHandler,
+	StoreUpgrades:        store.StoreUpgrades{},
+}
diff --git a/app/upgrades/v2_2_5/upgrades.go b/app/upgrades/v2_2_5/upgrades.go
new file mode 100644
index 00000000..1093fd38
--- /dev/null
+++ b/app/upgrades/v2_2_5/upgrades.go
@@ -0,0 +1,19 @@
+package v2_2_5
+
+import (
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+)
+
+// We set the app version to pre-upgrade because it will be incremented by one
+// after the upgrade is applied by the handler.
+
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configurator module.Configurator,
+) upgradetypes.UpgradeHandler {
+	return func(ctx sdk.Context, _plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		return mm.RunMigrations(ctx, configurator, vm)
+	}
+}