refactor(auth): use collections for GlobalAccountNumber (#15830)

Co-authored-by: testinginprod <testinginprod@somewhere.idk>

x/auth/keeper/keeper.go

@@ -7,12 +7,10 @@ import (
 
 	"cosmossdk.io/collections"
 
-	"cosmossdk.io/log"
-	gogotypes "github.com/cosmos/gogoproto/types"
-
 	"cosmossdk.io/core/address"
 	"cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
+	"cosmossdk.io/log"
 
 	"github.com/cosmos/cosmos-sdk/codec"
 	cryptotypes "github.com/cosmos/cosmos-sdk/crypto/types"
@@ -74,7 +72,8 @@ type AccountKeeper struct {
 
 	// State
 
-	ParamsState collections.Item[types.Params] // NOTE: name is this because it conflicts with the Params gRPC method impl
+	ParamsState   collections.Item[types.Params] // NOTE: name is this because it conflicts with the Params gRPC method impl
+	AccountNumber collections.Sequence
 }
 
 var _ AccountKeeperI = &AccountKeeper{}
@@ -99,13 +98,14 @@ func NewAccountKeeper(
 	sb := collections.NewSchemaBuilder(storeService)
 
 	return AccountKeeper{
-		storeService: storeService,
-		proto:        proto,
-		cdc:          cdc,
-		permAddrs:    permAddrs,
-		addressCdc:   bech32Codec,
-		authority:    authority,
-		ParamsState:  collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
+		storeService:  storeService,
+		proto:         proto,
+		cdc:           cdc,
+		permAddrs:     permAddrs,
+		addressCdc:    bech32Codec,
+		authority:     authority,
+		ParamsState:   collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
+		AccountNumber: collections.NewSequence(sb, types.GlobalAccountNumberKey, "account_number"),
 	}
 }
 
@@ -148,32 +148,11 @@ func (ak AccountKeeper) GetSequence(ctx context.Context, addr sdk.AccAddress) (u
 // NextAccountNumber returns and increments the global account number counter.
 // If the global account number is not set, it initializes it with value 0.
 func (ak AccountKeeper) NextAccountNumber(ctx context.Context) uint64 {
-	var accNumber uint64
-	store := ak.storeService.OpenKVStore(ctx)
-
-	bz, err := store.Get(types.GlobalAccountNumberKey)
+	n, err := ak.AccountNumber.Next(ctx)
 	if err != nil {
-		// panics only on nil key, which should not be possible
 		panic(err)
 	}
-	if bz == nil {
-		// initialize the account numbers
-		accNumber = 0
-	} else {
-		val := gogotypes.UInt64Value{}
-
-		err := ak.cdc.Unmarshal(bz, &val)
-		if err != nil {
-			panic(err)
-		}
-
-		accNumber = val.GetValue()
-	}
-
-	bz = ak.cdc.MustMarshal(&gogotypes.UInt64Value{Value: accNumber + 1})
-	store.Set(types.GlobalAccountNumberKey, bz)
-
-	return accNumber
+	return n
 }
 
 // GetModulePermissions fetches per-module account permissions.

x/auth/keeper/migrations.go

@@ -8,6 +8,7 @@ import (
 	v2 "github.com/cosmos/cosmos-sdk/x/auth/migrations/v2"
 	v3 "github.com/cosmos/cosmos-sdk/x/auth/migrations/v3"
 	v4 "github.com/cosmos/cosmos-sdk/x/auth/migrations/v4"
+	v5 "github.com/cosmos/cosmos-sdk/x/auth/migrations/v5"
 	"github.com/cosmos/cosmos-sdk/x/auth/types"
 )
 
@@ -59,6 +60,13 @@ func (m Migrator) Migrate3to4(ctx sdk.Context) error {
 	return v4.Migrate(ctx, m.keeper.storeService, m.legacySubspace, m.keeper.cdc)
 }
 
+// Migrate4To5 migrates the x/auth module state from the consensus version 4 to 5.
+// It migrates the GlobalAccountNumber from being a protobuf defined value to a
+// big-endian encoded uint64, it also migrates it to use a more canonical prefix.
+func (m Migrator) Migrate4To5(ctx sdk.Context) error {
+	return v5.Migrate(ctx, m.keeper.storeService, m.keeper.AccountNumber)
+}
+
 // V45_SetAccount implements V45_SetAccount
 // set the account without map to accAddr to accNumber.
 //

x/auth/migrations/v5/doc.go

@@ -1,4 +0,0 @@
-// v5 is an empty package that exists because of the group module.
-// the group module v2 migration actually migrates the auth module state (replace group policies accounts from module accounts to base accounts).
-// the auth state does not migrate if the group module is not enabled.
-package v5

x/auth/migrations/v5/migrate.go

@@ -0,0 +1,45 @@
+package v5
+
+import (
+	"context"
+
+	"cosmossdk.io/collections"
+	storetypes "cosmossdk.io/core/store"
+	"github.com/cosmos/gogoproto/types"
+)
+
+var LegacyGlobalAccountNumberKey = []byte("globalAccountNumber")
+
+func Migrate(ctx context.Context, storeService storetypes.KVStoreService, sequence collections.Sequence) error {
+	store := storeService.OpenKVStore(ctx)
+	b, err := store.Get(LegacyGlobalAccountNumberKey)
+	if err != nil {
+		return err
+	}
+	if b == nil {
+		// this would mean no account was ever created in this chain which is being migrated?
+		// we're doing nothing as the collections.Sequence already handles the non-existing value.
+		return nil
+	}
+
+	// get old value
+	v := new(types.UInt64Value)
+	err = v.Unmarshal(b)
+	if err != nil {
+		return err
+	}
+
+	// set the old value in the collection
+	err = sequence.Set(ctx, v.Value)
+	if err != nil {
+		return err
+	}
+
+	// remove the value from the old prefix.
+	err = store.Delete(LegacyGlobalAccountNumberKey)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}

x/auth/migrations/v5/migrate_test.go

@@ -0,0 +1,44 @@
+package v5
+
+import (
+	"testing"
+
+	"cosmossdk.io/collections"
+	"cosmossdk.io/collections/colltest"
+	"github.com/cosmos/gogoproto/types"
+	"github.com/stretchr/testify/require"
+)
+
+func TestMigrate(t *testing.T) {
+	kv, ctx := colltest.MockStore()
+	sb := collections.NewSchemaBuilder(kv)
+	seq := collections.NewSequence(sb, collections.NewPrefix(0), "seq")
+
+	wantValue := uint64(100)
+
+	// set old sequence to wanted value
+	legacySeqBytes, err := (&types.UInt64Value{Value: wantValue}).Marshal()
+	require.NoError(t, err)
+
+	err = kv.OpenKVStore(ctx).Set(LegacyGlobalAccountNumberKey, legacySeqBytes)
+	require.NoError(t, err)
+
+	err = Migrate(ctx, kv, seq)
+	require.NoError(t, err)
+
+	// check that after migration the sequence is what we want it to be
+	gotValue, err := seq.Peek(ctx)
+	require.NoError(t, err)
+	require.Equal(t, wantValue, gotValue)
+
+	// case the global account number was not set
+	ctx = kv.NewStoreContext() // this resets the store to zero
+	wantValue = collections.DefaultSequenceStart
+
+	err = Migrate(ctx, kv, seq)
+	require.NoError(t, err)
+
+	gotValue, err = seq.Next(ctx)
+	require.NoError(t, err)
+	require.Equal(t, wantValue, gotValue)
+}

x/auth/module.go

@@ -33,7 +33,7 @@ import (
 )
 
 // ConsensusVersion defines the current x/auth module consensus version.
-const ConsensusVersion = 4
+const ConsensusVersion = 5
 
 var (
 	_ module.AppModule           = AppModule{}
@@ -147,9 +147,8 @@ func (am AppModule) RegisterServices(cfg module.Configurator) {
 		panic(fmt.Sprintf("failed to migrate x/%s from version 3 to 4: %v", types.ModuleName, err))
 	}
 
-	// see migrations/v5/doc.go
-	if err := cfg.RegisterMigration(types.ModuleName, 4, func(ctx sdk.Context) error { return nil }); err != nil {
-		panic(fmt.Sprintf("failed to migrate x/%s from version 4 to 5: %v", types.ModuleName, err))
+	if err := cfg.RegisterMigration(types.ModuleName, 4, m.Migrate4To5); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/%s from version 4 to 5", types.ModuleName))
 	}
 }
 

x/auth/types/keys.go

@@ -23,8 +23,9 @@ var (
 	// AddressStoreKeyPrefix prefix for account-by-address store
 	AddressStoreKeyPrefix = []byte{0x01}
 
-	// param key for global account number
-	GlobalAccountNumberKey = []byte("globalAccountNumber")
+	// GlobalAccountNumberKey identifies the prefix where the monotonically increasing
+	// account number is stored.
+	GlobalAccountNumberKey = collections.NewPrefix(2)
 
 	// AccountNumberStoreKeyPrefix prefix for account-by-id store
 	AccountNumberStoreKeyPrefix = []byte("accountNumber")