Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

initialize megavault module account in 7.0.0 upgrade handler #2394

Merged
merged 3 commits into from
Sep 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions protocol/app/upgrades.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ func (app *App) setupUpgradeHandlers() {
v7_0_0.CreateUpgradeHandler(
app.ModuleManager,
app.configurator,
app.AccountKeeper,
app.PricesKeeper,
app.VaultKeeper,
),
Expand Down
75 changes: 75 additions & 0 deletions protocol/app/upgrades/v7.0.0/upgrade.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import (
upgradetypes "cosmossdk.io/x/upgrade/types"
sdk "github.com/cosmos/cosmos-sdk/types"
"github.com/cosmos/cosmos-sdk/types/module"
authkeeper "github.com/cosmos/cosmos-sdk/x/auth/keeper"
authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
"github.com/dydxprotocol/v4-chain/protocol/lib"
"github.com/dydxprotocol/v4-chain/protocol/lib/slinky"
pricestypes "github.com/dydxprotocol/v4-chain/protocol/x/prices/types"
Expand All @@ -20,6 +22,75 @@ const (
QUOTE_QUANTUMS_PER_MEGAVAULT_SHARE = 1_000_000
)

var (
ModuleAccsToInitialize = []string{
vaulttypes.MegavaultAccountName,
}
)

// This module account initialization logic is copied from v3.0.0 upgrade handler.
func initializeModuleAccs(ctx sdk.Context, ak authkeeper.AccountKeeper) {
for _, modAccName := range ModuleAccsToInitialize {
// Get module account and relevant permissions from the accountKeeper.
addr, perms := ak.GetModuleAddressAndPermissions(modAccName)
if addr == nil {
panic(fmt.Sprintf(
"Did not find %v in `ak.GetModuleAddressAndPermissions`. This is not expected. Skipping.",
modAccName,
))
}
Comment on lines +37 to +41
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Clarify Panic Message for Missing Module Account

The panic message suggests that the module account will be skipped, but since the code is panicking, execution will halt. This could be misleading.

Consider rephrasing the panic message to accurately convey the severity of the error:

panic(fmt.Sprintf(
-	"Did not find %v in `ak.GetModuleAddressAndPermissions`. This is not expected. Skipping.",
+	"Did not find %v in `ak.GetModuleAddressAndPermissions`. This is unexpected and will halt the upgrade.",
	modAccName,
))
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
panic(fmt.Sprintf(
"Did not find %v in `ak.GetModuleAddressAndPermissions`. This is not expected. Skipping.",
modAccName,
))
}
panic(fmt.Sprintf(
"Did not find %v in `ak.GetModuleAddressAndPermissions`. This is unexpected and will halt the upgrade.",
modAccName,
))
}


// Try to get the account in state.
acc := ak.GetAccount(ctx, addr)
if acc != nil {
// Account has been initialized.
macc, isModuleAccount := acc.(sdk.ModuleAccountI)
if isModuleAccount {
// Module account was correctly initialized. Skipping
ctx.Logger().Info(fmt.Sprintf(
"module account %+v was correctly initialized. No-op",
macc,
))
continue
}
// Module account has been initialized as a BaseAccount. Change to module account.
// Note: We need to get the base account to retrieve its account number, and convert it
// in place into a module account.
baseAccount, ok := acc.(*authtypes.BaseAccount)
if !ok {
panic(fmt.Sprintf(
"cannot cast %v into a BaseAccount, acc = %+v",
modAccName,
acc,
))
}
Comment on lines +60 to +66
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Ensure Safe Type Assertion to *authtypes.BaseAccount

The type assertion baseAccount, ok := acc.(*authtypes.BaseAccount) may fail if acc is not of type *authtypes.BaseAccount, leading to a panic. Verify that this conversion is always safe in your context, or handle the case where it might fail more gracefully.

Consider adding error handling or logging to manage unexpected account types without panicking, if appropriate.

newModuleAccount := authtypes.NewModuleAccount(
baseAccount,
modAccName,
perms...,
)
ak.SetModuleAccount(ctx, newModuleAccount)
ctx.Logger().Info(fmt.Sprintf(
"Successfully converted %v to module account in state: %+v",
modAccName,
newModuleAccount,
))
continue
}

// Account has not been initialized at all. Initialize it as module.
// Implementation taken from
// https://github.com/dydxprotocol/cosmos-sdk/blob/bdf96fdd/x/auth/keeper/keeper.go#L213
newModuleAccount := authtypes.NewEmptyModuleAccount(modAccName, perms...)
maccI := (ak.NewAccount(ctx, newModuleAccount)).(sdk.ModuleAccountI) // this set the account number
ak.SetModuleAccount(ctx, maccI)
ctx.Logger().Info(fmt.Sprintf(
"Successfully initialized module account in state: %+v",
newModuleAccount,
))
}
}

func initCurrencyPairIDCache(ctx sdk.Context, k pricestypes.PricesKeeper) {
marketParams := k.GetAllMarketParams(ctx)
for _, mp := range marketParams {
Expand Down Expand Up @@ -126,13 +197,17 @@ func migrateVaultSharesToMegavaultShares(ctx sdk.Context, k vaultkeeper.Keeper)
func CreateUpgradeHandler(
mm *module.Manager,
configurator module.Configurator,
accountKeeper authkeeper.AccountKeeper,
pricesKeeper pricestypes.PricesKeeper,
vaultKeeper vaultkeeper.Keeper,
) upgradetypes.UpgradeHandler {
return func(ctx context.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
sdkCtx := lib.UnwrapSDKContext(ctx, "app/upgrades")
sdkCtx.Logger().Info(fmt.Sprintf("Running %s Upgrade...", UpgradeName))

// Initialize module accounts.
initializeModuleAccs(sdkCtx, accountKeeper)

// Initialize the currency pair ID cache for all existing market params.
initCurrencyPairIDCache(sdkCtx, pricesKeeper)

Expand Down
23 changes: 22 additions & 1 deletion protocol/app/upgrades/v7.0.0/upgrade_container_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import (

"github.com/cosmos/gogoproto/proto"

authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
v_7_0_0 "github.com/dydxprotocol/v4-chain/protocol/app/upgrades/v7.0.0"
"github.com/dydxprotocol/v4-chain/protocol/dtypes"
"github.com/dydxprotocol/v4-chain/protocol/testing/containertest"
Expand Down Expand Up @@ -48,9 +49,12 @@ func preUpgradeChecks(node *containertest.Node, t *testing.T) {
}

func postUpgradeChecks(node *containertest.Node, t *testing.T) {
// Add test for your upgrade handler logic below
// Check that vault quoting params are successfully migrated to vault params.
postUpgradeVaultParamsCheck(node, t)
// Check that vault shares are successfully migrated to megavault shares.
postUpgradeMegavaultSharesCheck(node, t)
// Check that megavault module account is successfully initialized.
postUpgradeMegavaultModuleAccCheck(node, t)

// Check that the affiliates module has been initialized with the default tiers.
postUpgradeAffiliatesModuleTiersCheck(node, t)
Expand Down Expand Up @@ -182,3 +186,20 @@ func postUpgradeAffiliatesModuleTiersCheck(node *containertest.Node, t *testing.
require.NoError(t, err)
require.Equal(t, affiliatestypes.DefaultAffiliateTiers, affiliateTiersResp.Tiers)
}

func postUpgradeMegavaultModuleAccCheck(node *containertest.Node, t *testing.T) {
resp, err := containertest.Query(
node,
authtypes.NewQueryClient,
authtypes.QueryClient.ModuleAccountByName,
&authtypes.QueryModuleAccountByNameRequest{
Name: vaulttypes.MegavaultAccountName,
},
)
require.NoError(t, err)
require.NotNil(t, resp)

moduleAccResp := authtypes.QueryModuleAccountByNameResponse{}
err = proto.UnmarshalText(resp.String(), &moduleAccResp)
require.NoError(t, err)
}
Comment on lines +190 to +205
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Enhance verification of the module account's properties

After querying and unmarshalling the module account, it's important to verify that the account has the expected attributes. This ensures the module account is not only initialized but also correctly configured.

Consider adding assertions like:

require.NotNil(t, moduleAccResp.Account)
require.Equal(t, vaulttypes.MegavaultAccountName, moduleAccResp.Account.Name)
// Add additional checks as necessary

Simplify response handling by directly using the response object

Instead of unmarshalling the response from text, you can directly use the response object returned by the query. This simplifies the code and avoids unnecessary parsing.

Apply this diff to simplify the code:

-moduleAccResp := authtypes.QueryModuleAccountByNameResponse{}
-err = proto.UnmarshalText(resp.String(), &moduleAccResp)
-require.NoError(t, err)
+moduleAccResp := resp.(*authtypes.QueryModuleAccountByNameResponse)

Loading