feat: add v2.0 upgrade handler

app/upgrade.go

@@ -1,25 +1,25 @@
 package app
 
 import (
-	upgrades "github.com/Lorenzo-Protocol/lorenzo/app/upgrades"
-	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	"fmt"
+
+	"github.com/Lorenzo-Protocol/lorenzo/app/upgrades"
+	v200 "github.com/Lorenzo-Protocol/lorenzo/app/upgrades/v200"
+
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
 )
 
-var plans = []upgrades.Upgrade{}
+var router = upgrades.NewUpgradeRouter()
+
+func init() {
+	// register v2.0 upgrade plan
+	router.Register(v200.Upgrade)
+}
 
 // RegisterUpgradePlans register a handler of upgrade plan
 func (app *LorenzoApp) RegisterUpgradePlans() {
-	for _, u := range plans {
-		app.registerUpgradeHandler(u.UpgradeName,
-			u.StoreUpgrades,
-			u.UpgradeHandlerConstructor(
-				app.mm,
-				app.configurator,
-				app.appKeepers(),
-			),
-		)
-	}
+	app.setupUpgradeStoreLoaders()
+	app.setupUpgradeHandlers()
 }
 
 func (app *LorenzoApp) appKeepers() upgrades.AppKeepers {
@@ -35,22 +35,35 @@
 	}
 }
 
-// registerUpgradeHandler implements the upgrade execution logic of the upgrade module
-func (app *LorenzoApp) registerUpgradeHandler(
-	planName string,
-	upgrades *storetypes.StoreUpgrades,
-	upgradeHandler upgradetypes.UpgradeHandler,
-) {
+// configure store loader that checks if version == upgradeHeight and applies store upgrades
+func (app *LorenzoApp) setupUpgradeStoreLoaders() {
 	upgradeInfo, err := app.UpgradeKeeper.ReadUpgradeInfoFromDisk()
 	if err != nil {
-		app.Logger().Info("not found upgrade plan", "planName", planName, "err", err.Error())
+		panic(fmt.Sprintf("failed to read upgrade info from disk %s", err))
+	}
+
+	if app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
 		return
 	}
 
-	if upgradeInfo.Name == planName && !app.UpgradeKeeper.IsSkipHeight(upgradeInfo.Height) {
-		// this configures a no-op upgrade handler for the planName upgrade
-		app.UpgradeKeeper.SetUpgradeHandler(planName, upgradeHandler)
-		// configure store loader that checks if version+1 == upgradeHeight and applies store upgrades
-		app.SetStoreLoader(upgradetypes.UpgradeStoreLoader(upgradeInfo.Height, upgrades))
+	app.SetStoreLoader(
+		upgradetypes.UpgradeStoreLoader(
+			upgradeInfo.Height,
+			router.UpgradeInfo(upgradeInfo.Name).StoreUpgrades,
+		),
+	)
+}
+
+func (app *LorenzoApp) setupUpgradeHandlers() {
+	for upgradeName, upgrade := range router.Routers() {
+		// SAFE: upgrade handlers are registered in the init function
+		app.UpgradeKeeper.SetUpgradeHandler(
+			upgradeName,
+			upgrade.UpgradeHandlerConstructor(
+				app.mm,
+				app.configurator,
+				app.appKeepers(),
+			),
+		)
 	}
 }

app/upgrades/types.go

@@ -1,6 +1,10 @@
 package upgrades
 
 import (
+	agentkeeper "github.com/Lorenzo-Protocol/lorenzo/x/agent/keeper"
+	btcstakingkeeper "github.com/Lorenzo-Protocol/lorenzo/x/btcstaking/keeper"
+	plankeeper "github.com/Lorenzo-Protocol/lorenzo/x/plan/keeper"
+	tokenkeeper "github.com/Lorenzo-Protocol/lorenzo/x/token/keeper"
 	tmproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/cosmos/cosmos-sdk/codec"
 	store "github.com/cosmos/cosmos-sdk/store/types"
@@ -10,6 +14,7 @@ import (
 	authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
 	bankkeeper "github.com/cosmos/cosmos-sdk/x/bank/keeper"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	ibckeeper "github.com/cosmos/ibc-go/v7/modules/core/keeper"
 	evmkeeper "github.com/evmos/ethermint/x/evm/keeper"
 	feemarketkeeper "github.com/evmos/ethermint/x/feemarket/keeper"
 )
@@ -35,13 +40,42 @@ type ConsensusParamsReaderWriter interface {
 }
 
 type AppKeepers struct {
-	AppCodec        codec.Codec
-	AccountKeeper   authkeeper.AccountKeeper
-	BankKeeper      bankkeeper.Keeper
-	GetKey          func(moduleName string) *storetypes.KVStoreKey
-	ModuleManager   *module.Manager
-	EvmKeeper       *evmkeeper.Keeper
-	FeeMarketKeeper feemarketkeeper.Keeper
+	AppCodec         codec.Codec
+	AccountKeeper    authkeeper.AccountKeeper
+	BankKeeper       bankkeeper.Keeper
+	GetKey           func(moduleName string) *storetypes.KVStoreKey
+	ModuleManager    *module.Manager
+	IBCKeeper        *ibckeeper.Keeper
+	EvmKeeper        *evmkeeper.Keeper
+	FeeMarketKeeper  feemarketkeeper.Keeper
+	AgentKeeper      agentkeeper.Keeper
+	PlanKeeper       *plankeeper.Keeper
+	BTCStakingKeeper btcstakingkeeper.Keeper
+	TokenKeeper      *tokenkeeper.Keeper
 
 	ReaderWriter ConsensusParamsReaderWriter
 }
+
+type UpgradeRouter struct {
+	mu map[string]Upgrade
+}
+
+func NewUpgradeRouter() *UpgradeRouter {
+	return &UpgradeRouter{make(map[string]Upgrade)}
+}
+
+func (r *UpgradeRouter) Register(u Upgrade) *UpgradeRouter {
+	if _, has := r.mu[u.UpgradeName]; has {
+		panic(u.UpgradeName + " already registered")
+	}
+	r.mu[u.UpgradeName] = u
+	return r
+}
+
+func (r *UpgradeRouter) Routers() map[string]Upgrade {
+	return r.mu
+}
+
+func (r *UpgradeRouter) UpgradeInfo(planName string) Upgrade {
+	return r.mu[planName]
+}

app/upgrades/v200/migarations.go

@@ -0,0 +1,29 @@
+package v200
+
+import (
+	"github.com/Lorenzo-Protocol/lorenzo/app/upgrades"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+)
+
+func migrateAgentFromBTCStakingToAgent(
+	ctx sdk.Context,
+	app upgrades.AppKeepers,
+) error {
+	btcStakingParams := app.BTCStakingKeeper.GetParams(ctx)
+	for _, receiver := range btcStakingParams.Receivers {
+		app.AgentKeeper.AddAgent(
+			ctx,
+			receiver.Name,
+			receiver.Addr, receiver.EthAddr,
+			"",
+			"",
+		)
+	}
+	// TODO: Is params.Receiver of btcstaking module removed?
+	btcStakingParams.Receivers = nil
+	if err := app.BTCStakingKeeper.SetParams(ctx, btcStakingParams); err != nil {
+		return err
+	}
+	return nil
+}

app/upgrades/v200/upgrades.go

@@ -0,0 +1,72 @@
+package v200
+
+import (
+	"github.com/Lorenzo-Protocol/lorenzo/app/upgrades"
+	"github.com/Lorenzo-Protocol/lorenzo/x/agent"
+	agenttypes "github.com/Lorenzo-Protocol/lorenzo/x/agent/types"
+	"github.com/Lorenzo-Protocol/lorenzo/x/plan"
+	plantypes "github.com/Lorenzo-Protocol/lorenzo/x/plan/types"
+	"github.com/Lorenzo-Protocol/lorenzo/x/token"
+	tokentypes "github.com/Lorenzo-Protocol/lorenzo/x/token/types"
+
+	storetypes "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/module"
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+)
+
+// Upgrade defines a struct containing necessary fields that a SoftwareUpgradeProposal
+var Upgrade = upgrades.Upgrade{
+	UpgradeName:               "v2.0",
+	UpgradeHandlerConstructor: upgradeHandlerConstructor,
+	StoreUpgrades: &storetypes.StoreUpgrades{
+		Added: []string{
+			agenttypes.StoreKey,
+			plantypes.StoreKey,
+			tokentypes.StoreKey,
+		},
+	},
+}
+
+func upgradeHandlerConstructor(
+	m *module.Manager,
+	c module.Configurator,
+	app upgrades.AppKeepers,
+) upgradetypes.UpgradeHandler {
+	return func(ctx sdk.Context, _ upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
+		fromVM[agenttypes.ModuleName] = agent.AppModule{}.ConsensusVersion()
+		fromVM[plantypes.ModuleName] = plan.AppModule{}.ConsensusVersion()
+		fromVM[tokentypes.ModuleName] = token.AppModule{}.ConsensusVersion()
+
+		// agent module init
+		// 1. set admin
+		admin, err := sdk.AccAddressFromBech32("")
+		if err != nil {
+			return nil, err
+		}
+
+		if err := app.AgentKeeper.SetAdmin(ctx, admin); err != nil {
+			return nil, err
+		}
+
+		// 2. set agents
+		if err := migrateAgentFromBTCStakingToAgent(ctx, app); err != nil {
+			return nil, err
+		}
+
+		// plan module init
+		planParams := plantypes.Params{
+			AllowList: []string{"*"},
+		}
+
+		if err := app.PlanKeeper.SetParams(ctx, planParams); err != nil {
+			return nil, err
+		}
+
+		// 3. set token params
+		tokenParams := tokentypes.DefaultParams()
+		app.TokenKeeper.SetParams(ctx, tokenParams)
+
+		return app.ModuleManager.RunMigrations(ctx, c, fromVM)
+	}
+}

x/agent/keeper/agent.go

@@ -44,6 +44,15 @@ func (k Keeper) GetNextNumber(ctx sdk.Context) uint64 {
 	return sdk.BigEndianToUint64(bz)
 }
 
+func (k Keeper) SetAdmin(ctx sdk.Context, admin sdk.AccAddress) error {
+	// check if the sender is the current admin
+	if !k.GetAdmin(ctx).Equals(admin) {
+		return types.ErrAdminExists
+	}
+	k.setAdmin(ctx, admin)
+	return nil
+}
+
 // GetAdmin retrieves the admin address from the Keeper's store.
 //
 // Parameters:

x/agent/types/errors.go

@@ -12,4 +12,5 @@ var (
 	ErrUnAuthorized             = errorsmod.Register(ModuleName, 6, "unauthorized")
 	ErrInvalidBtcAddress        = errorsmod.Register(ModuleName, 7, "invalid btc address")
 	ErrInvalidEthAddress        = errorsmod.Register(ModuleName, 8, "invalid eth address")
+	ErrAdminExists              = errorsmod.Register(ModuleName, 9, "admin already exists")
 )