diff --git a/protocol/app/upgrades/v8.0/migrate_accountplus_test.go b/protocol/app/upgrades/v8.0/migrate_accountplus_test.go new file mode 100644 index 0000000000..88be8a4994 --- /dev/null +++ b/protocol/app/upgrades/v8.0/migrate_accountplus_test.go @@ -0,0 +1,81 @@ +package v_8_0 + +import ( + "testing" + + "cosmossdk.io/store/prefix" + sdk "github.com/cosmos/cosmos-sdk/types" + testapp "github.com/dydxprotocol/v4-chain/protocol/testutil/app" + accountplustypes "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/types" + "github.com/stretchr/testify/suite" +) + +type UpgradeTestSuite struct { + suite.Suite + + tApp *testapp.TestApp + Ctx sdk.Context +} + +func TestMigrateAccountplusAccountState(t *testing.T) { + suite.Run(t, new(UpgradeTestSuite)) +} + +func (s *UpgradeTestSuite) SetupTest() { + s.tApp = testapp.NewTestAppBuilder(s.T()).Build() + s.Ctx = s.tApp.InitChain() +} + +func (s *UpgradeTestSuite) TestUpgrade_MigrateAccountplusAccountState() { + ctx := s.Ctx + store := ctx.KVStore(s.tApp.App.AccountPlusKeeper.GetStoreKey()) + prefixStore := prefix.NewStore(store, []byte(accountplustypes.AccountStateKeyPrefix)) + + // Create some AccountState with no prefixes + addresses := []string{"address1", "address2", "address3"} + for _, addr := range addresses { + accAddress := sdk.AccAddress([]byte(addr)) + accountState := accountplustypes.AccountState{ + Address: addr, + TimestampNonceDetails: accountplustypes.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } + bz := s.tApp.App.AccountPlusKeeper.GetCdc().MustMarshal(&accountState) + store.Set(accAddress, bz) + } + + // Verify unprefixed keys were successfully created + for _, addr := range addresses { + accAddress := sdk.AccAddress([]byte(addr)) + bz := store.Get(accAddress) + s.Require().NotNil(bz, "Unprefixed key not created for %s", addr) + } + + // Migrate + migrateAccountplusAccountState(ctx, s.tApp.App.AccountPlusKeeper) + + // Verify that unprefixed keys are deleted and prefixed keys exist + for _, addr := range addresses { + accAddress := sdk.AccAddress([]byte(addr)) + + // Check that the old key no longer exists + bzOld := store.Get(accAddress) + s.Require().Nil(bzOld, "Unprefixed AccountState should be deleted for %s", addr) + + // Check that the new prefixed key exists + bzNew := prefixStore.Get(accAddress) + var actualAccountState accountplustypes.AccountState + s.tApp.App.AccountPlusKeeper.GetCdc().MustUnmarshal(bzNew, &actualAccountState) + expectedAccountState := accountplustypes.AccountState{ + Address: addr, + TimestampNonceDetails: accountplustypes.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } + s.Require().NotNil(bzNew, "Prefixed AccountState should exist for %s", addr) + s.Require().Equal(expectedAccountState, actualAccountState, "Incorrect AccountState after migration for %s", addr) + } +} diff --git a/protocol/app/upgrades/v8.0/upgrade.go b/protocol/app/upgrades/v8.0/upgrade.go new file mode 100644 index 0000000000..12707b727b --- /dev/null +++ b/protocol/app/upgrades/v8.0/upgrade.go @@ -0,0 +1,63 @@ +package v_8_0 + +import ( + "bytes" + "fmt" + + storetypes "cosmossdk.io/store/types" + sdk "github.com/cosmos/cosmos-sdk/types" + accountpluskeeper "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/keeper" + accountplustypes "github.com/dydxprotocol/v4-chain/protocol/x/accountplus/types" +) + +// Migrate accountplus AccountState in kvstore from non-prefixed keys to prefixed keys +func migrateAccountplusAccountState(ctx sdk.Context, k accountpluskeeper.Keeper) { + ctx.Logger().Info("Migrating accountplus module AccountState in kvstore from non-prefixed keys to prefixed keys") + + store := ctx.KVStore(k.GetStoreKey()) + + // Iterate on unprefixed keys + iterator := storetypes.KVStorePrefixIterator(store, nil) + defer iterator.Close() + + var keysToDelete [][]byte + var accountStatesToSet []struct { + address sdk.AccAddress + accountState accountplustypes.AccountState + } + for ; iterator.Valid(); iterator.Next() { + key := iterator.Key() + + // Double check that key does not have prefix + if bytes.HasPrefix(key, []byte(accountplustypes.AccountStateKeyPrefix)) { + panic(fmt.Sprintf("unexpected key with prefix %X found during migration", accountplustypes.AccountStateKeyPrefix)) + } + + value := iterator.Value() + var accountState accountplustypes.AccountState + if err := k.GetCdc().Unmarshal(value, &accountState); err != nil { + panic(fmt.Sprintf("failed to unmarshal AccountState for key %X: %s", key, err)) + } + + accountStatesToSet = append(accountStatesToSet, struct { + address sdk.AccAddress + accountState accountplustypes.AccountState + }{key, accountState}) + + keysToDelete = append(keysToDelete, key) + } + + // Set prefixed keys + for _, item := range accountStatesToSet { + k.SetAccountState(ctx, item.address, item.accountState) + } + + // Delete unprefixed keys + for _, key := range keysToDelete { + store.Delete(key) + } + + ctx.Logger().Info("Successfully migrated accountplus AccountState keys") +} + +// TODO: Scaffolding for upgrade: https://linear.app/dydx/issue/OTE-886/v8-upgrade-handler-scaffold diff --git a/protocol/x/accountplus/keeper/keeper.go b/protocol/x/accountplus/keeper/keeper.go index d63297e19f..b09d435d82 100644 --- a/protocol/x/accountplus/keeper/keeper.go +++ b/protocol/x/accountplus/keeper/keeper.go @@ -1,12 +1,11 @@ package keeper import ( - "bytes" - "errors" "fmt" "strings" "cosmossdk.io/log" + "cosmossdk.io/store/prefix" storetypes "cosmossdk.io/store/types" "github.com/cosmos/cosmos-sdk/codec" sdk "github.com/cosmos/cosmos-sdk/types" @@ -38,6 +37,14 @@ func NewKeeper( } } +func (k Keeper) GetStoreKey() storetypes.StoreKey { + return k.storeKey +} + +func (k Keeper) GetCdc() codec.BinaryCodec { + return k.cdc +} + func (k Keeper) Logger(ctx sdk.Context) log.Logger { return ctx.Logger().With(log.ModuleKey, fmt.Sprintf("x/%s", types.ModuleName)) } @@ -46,26 +53,22 @@ func (k Keeper) Logger(ctx sdk.Context) log.Logger { func (k Keeper) InitializeForGenesis(ctx sdk.Context) { } -// Get all account details pairs in store +// Get all AccountStates from kvstore func (k Keeper) GetAllAccountStates(ctx sdk.Context) ([]types.AccountState, error) { - store := ctx.KVStore(k.storeKey) - - iterator := storetypes.KVStorePrefixIterator(store, nil) + iterator := storetypes.KVStorePrefixIterator( + ctx.KVStore(k.storeKey), + []byte(types.AccountStateKeyPrefix), + ) defer iterator.Close() accounts := []types.AccountState{} for ; iterator.Valid(); iterator.Next() { - key := iterator.Key() - - // Temporary workaround to exclude smart account kv pairs. - if bytes.HasPrefix(key, []byte(types.SmartAccountKeyPrefix)) { - continue + var accountState types.AccountState + err := k.cdc.Unmarshal(iterator.Value(), &accountState) + if err != nil { + return nil, fmt.Errorf("failed to unmarshal account state: %w", err) } - accountState, found := k.GetAccountState(ctx, key) - if !found { - return accounts, errors.New("Could not get account state for address: " + sdk.AccAddress(iterator.Key()).String()) - } accounts = append(accounts, accountState) } @@ -144,7 +147,7 @@ func (k Keeper) GetAllAuthenticatorData(ctx sdk.Context) ([]types.AuthenticatorD return accountAuthenticators, nil } -func GetAccountPlusStateWithTimestampNonceDetails( +func AccountStateFromTimestampNonceDetails( address sdk.AccAddress, tsNonce uint64, ) types.AccountState { @@ -162,8 +165,8 @@ func (k Keeper) GetAccountState( ctx sdk.Context, address sdk.AccAddress, ) (types.AccountState, bool) { - store := ctx.KVStore(k.storeKey) - bz := store.Get(address.Bytes()) + prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.AccountStateKeyPrefix)) + bz := prefixStore.Get(address.Bytes()) if bz == nil { return types.AccountState{}, false } @@ -185,9 +188,9 @@ func (k Keeper) SetAccountState( address sdk.AccAddress, accountState types.AccountState, ) { - store := ctx.KVStore(k.storeKey) + prefixStore := prefix.NewStore(ctx.KVStore(k.storeKey), []byte(types.AccountStateKeyPrefix)) bz := k.cdc.MustMarshal(&accountState) - store.Set(address.Bytes(), bz) + prefixStore.Set(address.Bytes(), bz) } func (k Keeper) HasAuthority(authority string) bool { diff --git a/protocol/x/accountplus/keeper/keeper_test.go b/protocol/x/accountplus/keeper/keeper_test.go index bb0d2d0bbb..58e0101c66 100644 --- a/protocol/x/accountplus/keeper/keeper_test.go +++ b/protocol/x/accountplus/keeper/keeper_test.go @@ -53,6 +53,56 @@ func (s *KeeperTestSuite) SetupTest() { ) } +func (s *KeeperTestSuite) TestKeeper_Set_Get_GetAllAccountState() { + ctx := s.Ctx + + accountState1 := types.AccountState{ + Address: "address1", + TimestampNonceDetails: types.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } + + accountState2 := types.AccountState{ + Address: "address2", + TimestampNonceDetails: types.TimestampNonceDetails{ + TimestampNonces: []uint64{1, 2, 3}, + MaxEjectedNonce: 0, + }, + } + + // SetAccountState + s.tApp.App.AccountPlusKeeper.SetAccountState( + ctx, + sdk.AccAddress([]byte(accountState1.Address)), + accountState1, + ) + + // GetAccountState + _, found := s.tApp.App.AccountPlusKeeper.GetAccountState(ctx, sdk.AccAddress([]byte(accountState1.Address))) + s.Require().True(found, "Account state not found") + + // GetAllAccountStates + accountStates, err := s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) + s.Require().NoError(err, "Should not error when getting all account states") + s.Require().Equal(len(accountStates), 1, "Incorrect number of AccountStates retrieved") + s.Require().Equal(accountStates[0], accountState1, "Incorrect AccountState retrieved") + + // Add one more AccountState and check GetAllAccountStates + s.tApp.App.AccountPlusKeeper.SetAccountState( + ctx, + sdk.AccAddress([]byte(accountState2.Address)), + accountState2, + ) + + accountStates, err = s.tApp.App.AccountPlusKeeper.GetAllAccountStates(ctx) + s.Require().NoError(err, "Should not error when getting all account states") + s.Require().Equal(len(accountStates), 2, "Incorrect number of AccountStates retrieved") + s.Require().Contains(accountStates, accountState1, "Retrieved AccountStates does not contain accountState1") + s.Require().Contains(accountStates, accountState2, "Retrieved AccountStates does not contain accountState2") +} + func (s *KeeperTestSuite) TestKeeper_AddAuthenticator() { ctx := s.Ctx diff --git a/protocol/x/accountplus/keeper/timestampnonce.go b/protocol/x/accountplus/keeper/timestampnonce.go index e23a0196a6..dad8e25a16 100644 --- a/protocol/x/accountplus/keeper/timestampnonce.go +++ b/protocol/x/accountplus/keeper/timestampnonce.go @@ -41,7 +41,7 @@ func (k Keeper) ProcessTimestampNonce(ctx sdk.Context, acc sdk.AccountI, tsNonce accountState, found := k.GetAccountState(ctx, address) if !found { // initialize accountplus state with ts nonce details - k.SetAccountState(ctx, address, GetAccountPlusStateWithTimestampNonceDetails(address, tsNonce)) + k.SetAccountState(ctx, address, AccountStateFromTimestampNonceDetails(address, tsNonce)) } else { EjectStaleTimestampNonces(&accountState, blockTs) tsNonceAccepted := AttemptTimestampNonceUpdate(tsNonce, &accountState) @@ -56,9 +56,10 @@ func (k Keeper) ProcessTimestampNonce(ctx sdk.Context, acc sdk.AccountI, tsNonce // Inplace eject all stale timestamps. func EjectStaleTimestampNonces(accountState *types.AccountState, referenceTs uint64) { tsNonceDetails := &accountState.TimestampNonceDetails + oldestAllowedTs := referenceTs - MaxTimeInPastMs var newTsNonces []uint64 for _, tsNonce := range tsNonceDetails.TimestampNonces { - if tsNonce >= referenceTs-MaxTimeInPastMs { + if tsNonce >= oldestAllowedTs { newTsNonces = append(newTsNonces, tsNonce) } else { if tsNonce > tsNonceDetails.MaxEjectedNonce { @@ -115,9 +116,5 @@ func isLargerThanSmallestValue(value uint64, values []uint64) (bool, int) { } } - if value > values[minIndex] { - return true, minIndex - } else { - return false, minIndex - } + return value > values[minIndex], minIndex } diff --git a/protocol/x/accountplus/types/keys.go b/protocol/x/accountplus/types/keys.go index 3844749bf3..2e0d0edd91 100644 --- a/protocol/x/accountplus/types/keys.go +++ b/protocol/x/accountplus/types/keys.go @@ -17,6 +17,11 @@ const ( StoreKey = ModuleName ) +// Prefix for account state. +const ( + AccountStateKeyPrefix = "AS/" +) + // Below key prefixes are for smart account implementation. const ( // SmartAccountKeyPrefix is the prefix key for all smart account store state.