chore: migrate pundix chain accounts

PundiAI · Nov 4, 2024 · 6b57e32 · 6b57e32
1 parent ee7bd25
commit 6b57e32
Show file tree

Hide file tree

Showing 4 changed files with 54 additions and 5 deletions.
diff --git a/app/upgrade.go b/app/upgrade.go
@@ -35,6 +35,7 @@ func (app *App) setupUpgradeHandlers() {
 	app.UpgradeKeeper.SetUpgradeHandler(
 		plan.UpgradeName,
 		plan.CreateUpgradeHandler(
+			app.appCodec,
 			app.mm,
 			app.configurator,
 			&app.AppKeepers,

diff --git a/app/upgrades/types.go b/app/upgrades/types.go
@@ -3,6 +3,7 @@ package upgrades
 import (
 	storetypes "cosmossdk.io/store/types"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
+	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
 
@@ -18,7 +19,7 @@ type Upgrade struct {
 	UpgradeName string
 
 	// CreateUpgradeHandler defines the function that creates an upgrade handler
-	CreateUpgradeHandler func(*module.Manager, module.Configurator, *keepers.AppKeepers) upgradetypes.UpgradeHandler
+	CreateUpgradeHandler func(codec.Codec, *module.Manager, module.Configurator, *keepers.AppKeepers) upgradetypes.UpgradeHandler
 
 	// Store upgrades, should be used for any new modules introduced, new modules deleted, or store names renamed.
 	StoreUpgrades func() *storetypes.StoreUpgrades

diff --git a/app/upgrades/v8/pundix.go b/app/upgrades/v8/pundix.go
@@ -13,6 +13,7 @@ import (
 	"cosmossdk.io/store/prefix"
 	storetypes "cosmossdk.io/store/types"
 	tmjson "github.com/cometbft/cometbft/libs/json"
+	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/bech32"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -36,6 +37,7 @@ import (
 )
 
 type Pundix struct {
+	cdc               codec.Codec
 	accountKeeper     authkeeper.AccountKeeper
 	erc20Keeper       erc20keeper.Keeper
 	bankKeeper        bankkeeper.Keeper
@@ -44,8 +46,9 @@ type Pundix struct {
 	storeKey          storetypes.StoreKey
 }
 
-func NewPundix(app *keepers.AppKeepers) *Pundix {
+func NewPundix(cdc codec.Codec, app *keepers.AppKeepers) *Pundix {
 	return &Pundix{
+		cdc:               cdc,
 		accountKeeper:     app.AccountKeeper,
 		erc20Keeper:       app.Erc20Keeper,
 		bankKeeper:        app.BankKeeper,
@@ -61,6 +64,15 @@ func (m *Pundix) Migrate(ctx sdk.Context) error {
 	if err != nil {
 		return err
 	}
+
+	authRaw, ok := appState[types.ModuleName]
+	if !ok || len(authRaw) == 0 {
+		return sdkerrors.ErrNotFound.Wrap("auth genesis")
+	}
+	if err = m.migrateAuth(ctx, appState[types.ModuleName]); err != nil {
+		return err
+	}
+
 	// todo migrate other data
 	bankGenesis, ok := appState[banktypes.ModuleName]
 	if !ok || len(bankGenesis) == 0 {
@@ -69,9 +81,43 @@ func (m *Pundix) Migrate(ctx sdk.Context) error {
 	return m.migrateBank(ctx, bankGenesis)
 }
 
+func (m *Pundix) migrateAuth(ctx sdk.Context, authRaw json.RawMessage) error {
+	var authGenesis types.GenesisState
+	if err := m.cdc.UnmarshalJSON(authRaw, &authGenesis); err != nil {
+		return err
+	}
+	genesisAccounts, err := types.UnpackAccounts(authGenesis.Accounts)
+	if err != nil {
+		return err
+	}
+	for _, genAcc := range genesisAccounts {
+		pubKey := genAcc.GetPubKey()
+		if pubKey == nil {
+			continue
+		}
+		baseAcc, ok := genAcc.(*types.BaseAccount)
+		if !ok {
+			continue
+		}
+		accAddr, err := sdk.GetFromBech32(baseAcc.Address, "px")
+		if err != nil {
+			return err
+		}
+		if m.accountKeeper.HasAccount(ctx, accAddr) {
+			continue
+		}
+		newAcc := m.accountKeeper.NewAccountWithAddress(ctx, accAddr)
+		if err = newAcc.SetPubKey(pubKey); err != nil {
+			return err
+		}
+		m.accountKeeper.SetAccount(ctx, newAcc)
+	}
+	return nil
+}
+
 func (m *Pundix) migrateBank(ctx sdk.Context, bankRaw json.RawMessage) error {
 	var bankGenesis banktypes.GenesisState
-	if err := tmjson.Unmarshal(bankRaw, &bankGenesis); err != nil {
+	if err := m.cdc.UnmarshalJSON(bankRaw, &bankGenesis); err != nil {
 		return err
 	}
 	if err := m.migratePUNDIX(ctx, bankGenesis.Balances, bankGenesis.Supply); err != nil {

diff --git a/app/upgrades/v8/upgrade.go b/app/upgrades/v8/upgrade.go
@@ -6,6 +6,7 @@ import (
 
 	storetypes "cosmossdk.io/store/types"
 	upgradetypes "cosmossdk.io/x/upgrade/types"
+	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/types/module"
@@ -31,7 +32,7 @@ import (
 	fxstakingv8 "github.com/functionx/fx-core/v8/x/staking/migrations/v8"
 )
 
-func CreateUpgradeHandler(mm *module.Manager, configurator module.Configurator, app *keepers.AppKeepers) upgradetypes.UpgradeHandler {
+func CreateUpgradeHandler(cdc codec.Codec, mm *module.Manager, configurator module.Configurator, app *keepers.AppKeepers) upgradetypes.UpgradeHandler {
 	return func(ctx context.Context, plan upgradetypes.Plan, fromVM module.VersionMap) (module.VersionMap, error) {
 		cacheCtx, commit := sdk.UnwrapSDKContext(ctx).CacheContext()
 
@@ -55,7 +56,7 @@ func CreateUpgradeHandler(mm *module.Manager, configurator module.Configurator,
 			return fromVM, err
 		}
 
-		if err = NewPundix(app).Migrate(cacheCtx); err != nil {
+		if err = NewPundix(cdc, app).Migrate(cacheCtx); err != nil {
 			return fromVM, err
 		}