v2.0.0 hard fork for initializing module accounts

protocol/app/upgrades.go

@@ -2,31 +2,35 @@ package app
 
 import (
 	"fmt"
+
 	"github.com/dydxprotocol/v4-chain/protocol/app/upgrades"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	v2_0_0 "github.com/dydxprotocol/v4-chain/protocol/app/upgrades/v2.0.0"
 )
 
 var (
 	// `Upgrades` defines the upgrade handlers and store loaders for the application.
 	// New upgrades should be added to this slice after they are implemented.
-	Upgrades = []upgrades.Upgrade{}
-	Forks    = []upgrades.Fork{}
+	Upgrades = []upgrades.Upgrade{v2_0_0.Upgrade}
+	Forks    = []upgrades.Fork{v2_0_0.Fork}
 )
 
 // setupUpgradeHandlers registers the upgrade handlers to perform custom upgrade
 // logic and state migrations for software upgrades.
 func (app *App) setupUpgradeHandlers() {
-	for _, upgrade := range Upgrades {
-		if app.UpgradeKeeper.HasHandler(upgrade.UpgradeName) {
-			panic(fmt.Sprintf("Cannot register duplicate upgrade handler '%s'", upgrade.UpgradeName))
-		}
-		app.UpgradeKeeper.SetUpgradeHandler(
-			upgrade.UpgradeName,
-			upgrade.CreateUpgradeHandler(app.ModuleManager, app.configurator),
-		)
+	if app.UpgradeKeeper.HasHandler(v2_0_0.UpgradeName) {
+		panic(fmt.Sprintf("Cannot register duplicate upgrade handler '%s'", v2_0_0.UpgradeName))
 	}
+	app.UpgradeKeeper.SetUpgradeHandler(
+		v2_0_0.UpgradeName,
+		v2_0_0.CreateUpgradeHandler(
+			app.ModuleManager,
+			app.configurator,
+			app.AccountKeeper,
+		),
+	)
 }
 
 // setUpgradeStoreLoaders sets custom store loaders to customize the rootMultiStore
@@ -65,6 +69,7 @@ func (app *App) scheduleForkUpgrade(ctx sdk.Context) {
 				Info:   fork.UpgradeInfo,
 			}
 
+			ctx.Logger().Info(fmt.Sprintf("Scheduling fork upgrade: %+v", fork))
 			// schedule the upgrade plan to the current block height, effectively performing
 			// a hard fork that uses the upgrade handler to manage the migration.
 			if err := app.UpgradeKeeper.ScheduleUpgrade(ctx, upgradePlan); err != nil {

protocol/app/upgrades/v2.0.0/constants.go

@@ -0,0 +1,25 @@
+package v_2_0_0
+
+import (
+	store "github.com/cosmos/cosmos-sdk/store/types"
+	"github.com/dydxprotocol/v4-chain/protocol/app/upgrades"
+)
+
+const (
+	UpgradeName = "v2.0.0"
+
+	UpgradeHeight = 1810000 // Estimated to be 5:50 PM ET Nov 23, 2023
+)
+
+var (
+	Fork = upgrades.Fork{
+		UpgradeName:   UpgradeName,
+		UpgradeHeight: UpgradeHeight,
+		UpgradeInfo:   "",
+	}
+
+	Upgrade = upgrades.Upgrade{
+		UpgradeName:   UpgradeName,
+		StoreUpgrades: store.StoreUpgrades{},
+	}
+)

protocol/app/upgrades/v2.0.0/upgrade.go

@@ -0,0 +1,102 @@
+package v_2_0_0
+
+import (
+	"fmt"
+
+	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"
+	upgradetypes "github.com/cosmos/cosmos-sdk/x/upgrade/types"
+	bridgemoduletypes "github.com/dydxprotocol/v4-chain/protocol/x/bridge/types"
+	clobmoduletypes "github.com/dydxprotocol/v4-chain/protocol/x/clob/types"
+	rewardsmoduletypes "github.com/dydxprotocol/v4-chain/protocol/x/rewards/types"
+	satypes "github.com/dydxprotocol/v4-chain/protocol/x/subaccounts/types"
+	vestmoduletypes "github.com/dydxprotocol/v4-chain/protocol/x/vest/types"
+)
+
+var (
+	// List of module accounts to check in state.
+	// These include all dYdX custom module accounts.
+	ModuleAccsToInitialize = []string{
+		bridgemoduletypes.ModuleName,
+		satypes.ModuleName,
+		clobmoduletypes.InsuranceFundName,
+		rewardsmoduletypes.TreasuryAccountName,
+		rewardsmoduletypes.VesterAccountName,
+		vestmoduletypes.CommunityTreasuryAccountName,
+		vestmoduletypes.CommunityVesterAccountName,
+	}
+)
+
+func CreateUpgradeHandler(
+	mm *module.Manager,
+	configurator module.Configurator,
+	ak authkeeper.AccountKeeper,
+) upgradetypes.UpgradeHandler {
+	return func(ctx sdk.Context, plan upgradetypes.Plan, vm module.VersionMap) (module.VersionMap, error) {
+		ctx.Logger().Info("Running v2.0.0 Upgrade...")
+
+		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,
+				))
+			}
+
+			// Try to get the account in state.
+			acc := ak.GetAccount(ctx, addr)
+			if acc != nil {
+				// Account has been initialized.
+				macc, isModuleAccount := acc.(authtypes.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,
+					))
+				}
+				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)).(authtypes.ModuleAccountI) // this set the account number
+			ak.SetModuleAccount(ctx, maccI)
+			ctx.Logger().Info(fmt.Sprintf(
+				"Successfully initialized module account in state: %+v",
+				newModuleAccount,
+			))
+		}
+
+		return mm.RunMigrations(ctx, configurator, vm)
+	}
+}

protocol/app/upgrades_test.go

@@ -75,6 +75,6 @@ func TestSetupUpgradeHandlers(t *testing.T) {
 }
 
 func TestDefaultUpgradesAndForks(t *testing.T) {
-	require.Empty(t, app.Upgrades, "Expected empty upgrades list")
-	require.Empty(t, app.Forks, "Expected empty forks list")
+	require.Len(t, app.Upgrades, 1)
+	require.Len(t, app.Forks, 1)
 }