[OTE-882] Add prefix to accountplus keeper (#2526)

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)
+	}
+}

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

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 {

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
 

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
 }

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.