From c749cebfd4d709a938f33f272d1c8988caacbffe Mon Sep 17 00:00:00 2001
From: zakir <80246097+zakir-code@users.noreply.github.com>
Date: Thu, 16 Jan 2025 18:24:38 +0800
Subject: [PATCH] feat: update token handling to improve security (#914)

Co-authored-by: nulnut <151493716+nulnut@users.noreply.github.com>
---
 app/upgrade_test.go              | 31 ++++++++++++++
 app/upgrades/v8/upgrade.go       | 69 ++++++++++++++++++++++++++++++--
 scripts/linter.sh                |  2 +-
 x/erc20/migrations/v8/keys.go    |  3 +-
 x/erc20/migrations/v8/migrate.go |  8 ++--
 5 files changed, 105 insertions(+), 8 deletions(-)

diff --git a/app/upgrade_test.go b/app/upgrade_test.go
index 8946ef51..1f491634 100644
--- a/app/upgrade_test.go
+++ b/app/upgrade_test.go
@@ -5,6 +5,7 @@ import (
 	"math/big"
 	"os"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"cosmossdk.io/collections"
@@ -20,6 +21,7 @@ import (
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
 	distributiontypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
 	govtypes "github.com/cosmos/cosmos-sdk/x/gov/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -92,6 +94,9 @@ func Test_UpgradeTestnet(t *testing.T) {
 
 	// check bridge fee contract status
 	checkBridgeFeeContract(t, ctx, myApp)
+
+	checkMetadataValidate(t, ctx, myApp)
+	checkPundiAIFXERC20Token(t, ctx, myApp)
 }
 
 func buildApp(t *testing.T) *app.App {
@@ -174,6 +179,8 @@ func checkAppUpgrade(t *testing.T, ctx sdk.Context, myApp *app.App, bdd BeforeUp
 	checkPundixPurse(t, ctx, myApp)
 	checkTotalSupply(t, ctx, myApp)
 	checkLayer2OracleIsOnline(t, ctx, myApp.Layer2Keeper)
+	checkMetadataValidate(t, ctx, myApp)
+	checkPundiAIFXERC20Token(t, ctx, myApp)
 }
 
 func checkLayer2OracleIsOnline(t *testing.T, ctx sdk.Context, layer2Keeper crosschainkeeper.Keeper) {
@@ -565,3 +572,27 @@ func checkBridgeFeeContract(t *testing.T, ctx sdk.Context, myApp *app.App) {
 		require.Empty(t, oracles)
 	}
 }
+
+func checkMetadataValidate(t *testing.T, ctx sdk.Context, myApp *app.App) {
+	t.Helper()
+
+	myApp.BankKeeper.IterateAllDenomMetaData(ctx, func(metadata banktypes.Metadata) bool {
+		if len(metadata.DenomUnits) <= 1 {
+			return false
+		}
+		require.NoError(t, metadata.Validate())
+		require.NotEqual(t, metadata.Display, metadata.Base)
+		return false
+	})
+}
+
+func checkPundiAIFXERC20Token(t *testing.T, ctx sdk.Context, myApp *app.App) {
+	t.Helper()
+
+	has, err := myApp.Erc20Keeper.ERC20Token.Has(ctx, fxtypes.DefaultDenom)
+	require.NoError(t, err)
+	require.True(t, has)
+	has, err = myApp.Erc20Keeper.ERC20Token.Has(ctx, strings.ToUpper(fxtypes.FXDenom))
+	require.NoError(t, err)
+	require.False(t, has)
+}
diff --git a/app/upgrades/v8/upgrade.go b/app/upgrades/v8/upgrade.go
index c20ec1dc..ae130a34 100644
--- a/app/upgrades/v8/upgrade.go
+++ b/app/upgrades/v8/upgrade.go
@@ -85,13 +85,22 @@ func upgradeTestnet(ctx sdk.Context, app *keepers.AppKeepers) error {
 	if err := migrateFeemarketGasPrice(ctx, app.FeeMarketKeeper); err != nil {
 		return err
 	}
+	if err := redeployTestnetContract(ctx, app.AccountKeeper, app.EvmKeeper, app.Erc20Keeper, app.EthKeeper); err != nil {
+		return err
+	}
 	if err := migrateCrosschainParams(ctx, app.CrosschainKeepers); err != nil {
 		return err
 	}
-
-	return redeployTestnetContract(ctx, app.AccountKeeper, app.EvmKeeper, app.Erc20Keeper, app.EthKeeper)
+	if err := migrateMetadataDisplay(ctx, app.BankKeeper); err != nil {
+		return err
+	}
+	if err := migrateErc20FXToPundiAI(ctx, app.Erc20Keeper); err != nil {
+		return err
+	}
+	return migrateMetadataFXToPundiAI(ctx, app.BankKeeper)
 }
 
+//nolint:gocyclo // mainnet
 func upgradeMainnet(
 	ctx sdk.Context,
 	mm *module.Manager,
@@ -171,7 +180,15 @@ func upgradeMainnet(
 	if err = migrateFeemarketGasPrice(ctx, app.FeeMarketKeeper); err != nil {
 		return toVM, err
 	}
-
+	if err = migrateMetadataDisplay(ctx, app.BankKeeper); err != nil {
+		return toVM, err
+	}
+	if err = migrateErc20FXToPundiAI(ctx, app.Erc20Keeper); err != nil {
+		return toVM, err
+	}
+	if err = migrateMetadataFXToPundiAI(ctx, app.BankKeeper); err != nil {
+		return toVM, err
+	}
 	return toVM, nil
 }
 
@@ -569,3 +586,49 @@ func migrateCrosschainParams(ctx sdk.Context, keepers keepers.CrosschainKeepers)
 	}
 	return nil
 }
+
+func migrateMetadataDisplay(ctx sdk.Context, bankKeeper bankkeeper.Keeper) error {
+	mds := bankKeeper.GetAllDenomMetaData(ctx)
+	for _, md := range mds {
+		if md.Display != md.Base || len(md.DenomUnits) <= 1 {
+			continue
+		}
+		for _, dus := range md.DenomUnits {
+			if dus.Denom != md.Base {
+				md.Display = dus.Denom
+				break
+			}
+		}
+		if err := md.Validate(); err != nil {
+			return err
+		}
+		bankKeeper.SetDenomMetaData(ctx, md)
+	}
+	return nil
+}
+
+func migrateErc20FXToPundiAI(ctx sdk.Context, keeper erc20keeper.Keeper) error {
+	fxDenom := strings.ToUpper(fxtypes.FXDenom)
+	erc20Token, err := keeper.GetERC20Token(ctx, fxDenom)
+	if err != nil {
+		return err
+	}
+	erc20Token.Denom = fxtypes.DefaultDenom
+	if err = keeper.ERC20Token.Set(ctx, erc20Token.Denom, erc20Token); err != nil {
+		return err
+	}
+	return keeper.ERC20Token.Remove(ctx, fxDenom)
+}
+
+func migrateMetadataFXToPundiAI(ctx sdk.Context, keeper bankkeeper.Keeper) error {
+	// add pundiai metadata
+	metadata := fxtypes.NewDefaultMetadata()
+	keeper.SetDenomMetaData(ctx, metadata)
+
+	// remove FX metadata
+	bk, ok := keeper.(bankkeeper.BaseKeeper)
+	if !ok {
+		return errors.New("bank keeper not implement bank.BaseKeeper")
+	}
+	return bk.BaseViewKeeper.DenomMetadata.Remove(ctx, strings.ToUpper(fxtypes.FXDenom))
+}
diff --git a/scripts/linter.sh b/scripts/linter.sh
index 0f9acb23..e2573a1b 100755
--- a/scripts/linter.sh
+++ b/scripts/linter.sh
@@ -3,7 +3,7 @@
 set -eo pipefail
 
 patternLimits=(
-  "nolint:22"
+  "nolint:23"
   "#nosec:5"
   "CrossChain:4"
   "cross chain:0"
diff --git a/x/erc20/migrations/v8/keys.go b/x/erc20/migrations/v8/keys.go
index 26b61718..c252620e 100644
--- a/x/erc20/migrations/v8/keys.go
+++ b/x/erc20/migrations/v8/keys.go
@@ -64,6 +64,7 @@ func (m Migrator) migrateParams(ctx sdk.Context, store storetypes.KVStore) error
 }
 
 func (m Migrator) migrateTokenPair(ctx sdk.Context, store storetypes.KVStore) error {
+	fxDenom := strings.ToUpper(fxtypes.FXDenom)
 	iterator := storetypes.KVStorePrefixIterator(store, KeyPrefixTokenPair)
 	defer iterator.Close()
 	for ; iterator.Valid(); iterator.Next() {
@@ -73,7 +74,7 @@ func (m Migrator) migrateTokenPair(ctx sdk.Context, store storetypes.KVStore) er
 		if !found {
 			return sdkerrors.ErrKeyNotFound.Wrapf("metadata not found: %s", tokenPair.GetDenom())
 		}
-		if md.Base == fxtypes.FXDenom || md.Base == strings.ToLower(md.Symbol) {
+		if md.Base == fxDenom || md.Base == strings.ToLower(md.Symbol) {
 			if err := m.keeper.ERC20Token.Set(ctx, md.Base, tokenPair); err != nil {
 				return err
 			}
diff --git a/x/erc20/migrations/v8/migrate.go b/x/erc20/migrations/v8/migrate.go
index e2bbfb0a..45690162 100644
--- a/x/erc20/migrations/v8/migrate.go
+++ b/x/erc20/migrations/v8/migrate.go
@@ -16,15 +16,16 @@ import (
 )
 
 func (m Migrator) MigrateToken(ctx sdk.Context) error {
+	fxDenom := strings.ToUpper(fxtypes.FXDenom)
 	// add FX bridge token
-	if err := m.addToken(ctx, fxtypes.FXDenom, ""); err != nil {
+	if err := m.addToken(ctx, fxDenom, ""); err != nil {
 		return err
 	}
 
 	mds := m.bankKeeper.GetAllDenomMetaData(ctx)
 	for _, md := range mds {
 		// exclude FX and alias empty, except PUNDIX
-		if md.Base == fxtypes.FXDenom || (len(md.DenomUnits) == 0 || len(md.DenomUnits[0].Aliases) == 0) && md.Symbol != "PUNDIX" {
+		if md.Base == fxDenom || (len(md.DenomUnits) == 0 || len(md.DenomUnits[0].Aliases) == 0) && md.Symbol != "PUNDIX" {
 			continue
 		}
 
@@ -73,11 +74,12 @@ func (m Migrator) addIBCToken(ctx sdk.Context, base, alias string) error {
 }
 
 func (m Migrator) addBridgeToken(ctx sdk.Context, base, alias string) error {
+	fxDenom := strings.ToUpper(fxtypes.FXDenom)
 	if getExcludeBridgeToken(ctx, alias) {
 		return nil
 	}
 	for _, ck := range m.crosschainKeepers {
-		canAddFxBridgeToken := base == fxtypes.FXDenom && ck.ModuleName() == ethtypes.ModuleName
+		canAddFxBridgeToken := base == fxDenom && ck.ModuleName() == ethtypes.ModuleName
 
 		canAddBridgeToken := strings.HasPrefix(alias, ck.ModuleName())
 		excludeModule := ck.ModuleName() != arbitrumtypes.ModuleName && ck.ModuleName() != optimismtypes.ModuleName