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

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

collections/colltest/store.go

@@ -7,21 +7,32 @@ import (
 	db "github.com/cosmos/cosmos-db"
 )
 
+type contextStoreKey struct{}
+
 // MockStore returns a mock store.KVStoreService and a mock context.Context.
-// They can be used to test collections.
-func MockStore() (store.KVStoreService, context.Context) {
+// They can be used to test collections. The StoreService.NewStoreContext
+// can be used to instantiate a new empty KVStore.
+func MockStore() (*StoreService, context.Context) {
 	kv := db.NewMemDB()
-	return &testStore{kv}, context.Background()
+	ctx := context.WithValue(context.Background(), contextStoreKey{}, &testStore{kv})
+	return &StoreService{}, ctx
+}
+
+type StoreService struct{}
+
+func (s StoreService) OpenKVStore(ctx context.Context) store.KVStore {
+	return ctx.Value(contextStoreKey{}).(store.KVStore)
+}
+
+func (s StoreService) NewStoreContext() context.Context {
+	kv := db.NewMemDB()
+	return context.WithValue(context.Background(), contextStoreKey{}, &testStore{kv})
 }
 
 type testStore struct {
 	db db.DB
 }
 
-func (t testStore) OpenKVStore(ctx context.Context) store.KVStore {
-	return t
-}
-
 func (t testStore) Get(key []byte) ([]byte, error) {
 	return t.db.Get(key)
 }

collections/sequence.go

@@ -6,7 +6,7 @@ import (
 )
 
 // DefaultSequenceStart defines the default starting number of a sequence.
-const DefaultSequenceStart uint64 = 1
+const DefaultSequenceStart uint64 = 0
 
 // Sequence builds on top of an Item, and represents a monotonically increasing number.
 type Sequence Item[uint64]

go.mod

@@ -4,7 +4,7 @@ module github.com/cosmos/cosmos-sdk
 
 require (
 	cosmossdk.io/api v0.4.0
-	cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4
+	cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e
 	cosmossdk.io/core v0.6.1
 	cosmossdk.io/depinject v1.0.0-alpha.3
 	cosmossdk.io/errors v1.0.0-beta.7

go.sum

@@ -37,8 +37,8 @@ cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9
 cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3fOKtUw0Xmo=
 cosmossdk.io/api v0.4.0 h1:x90DmdidP6EhzktAa/6/IofSHidDnPjahdlrUvyQZQw=
 cosmossdk.io/api v0.4.0/go.mod h1:TWDzBhUBhI1LhSf2XSYpfIBf6D4mbLu/fvzvDfhcaYM=
-cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4 h1:QQZ0Qz8Gy/EmUNMRiHkUPG3BMA6OqEBp67IsfKETXIU=
-cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4/go.mod h1:/vS4ugR7ad3IciUd5TQuP2Ldz3NukHK2u/l5xTxXbbE=
+cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e h1:3gV/RSkCybtvU7unzH/Tt7yI7W8O6q22Lg7FSwSP1K8=
+cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e/go.mod h1:/vS4ugR7ad3IciUd5TQuP2Ldz3NukHK2u/l5xTxXbbE=
 cosmossdk.io/core v0.6.1 h1:OBy7TI2W+/gyn2z40vVvruK3di+cAluinA6cybFbE7s=
 cosmossdk.io/core v0.6.1/go.mod h1:g3MMBCBXtxbDWBURDVnJE7XML4BG5qENhs0gzkcpuFA=
 cosmossdk.io/depinject v1.0.0-alpha.3 h1:6evFIgj//Y3w09bqOUOzEpFj5tsxBqdc5CfkO7z+zfw=

simapp/go.mod

@@ -36,7 +36,7 @@ require (
 	cloud.google.com/go/compute/metadata v0.2.3 // indirect
 	cloud.google.com/go/iam v0.13.0 // indirect
 	cloud.google.com/go/storage v1.30.0 // indirect
-	cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4 // indirect
+	cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e // indirect
 	cosmossdk.io/errors v1.0.0-beta.7 // indirect
 	cosmossdk.io/x/tx v0.5.1-0.20230407182919-057d2e09bd63 // indirect
 	filippo.io/edwards25519 v1.0.0 // indirect

simapp/go.sum

@@ -190,8 +190,8 @@ cloud.google.com/go/workflows v1.6.0/go.mod h1:6t9F5h/unJz41YqfBmqSASJSXccBLtD1V
 cloud.google.com/go/workflows v1.7.0/go.mod h1:JhSrZuVZWuiDfKEFxU0/F1PQjmpnpcoISEXH2bcHC3M=
 cosmossdk.io/client/v2 v2.0.0-20230309163709-87da587416ba h1:LuPHCncU2KLMNPItFECs709uo46I9wSu2fAWYVCx+/U=
 cosmossdk.io/client/v2 v2.0.0-20230309163709-87da587416ba/go.mod h1:SXdwqO7cN5htalh/lhXWP8V4zKtBrhhcSTU+ytuEtmM=
-cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4 h1:QQZ0Qz8Gy/EmUNMRiHkUPG3BMA6OqEBp67IsfKETXIU=
-cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4/go.mod h1:/vS4ugR7ad3IciUd5TQuP2Ldz3NukHK2u/l5xTxXbbE=
+cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e h1:3gV/RSkCybtvU7unzH/Tt7yI7W8O6q22Lg7FSwSP1K8=
+cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e/go.mod h1:/vS4ugR7ad3IciUd5TQuP2Ldz3NukHK2u/l5xTxXbbE=
 cosmossdk.io/core v0.6.1 h1:OBy7TI2W+/gyn2z40vVvruK3di+cAluinA6cybFbE7s=
 cosmossdk.io/core v0.6.1/go.mod h1:g3MMBCBXtxbDWBURDVnJE7XML4BG5qENhs0gzkcpuFA=
 cosmossdk.io/depinject v1.0.0-alpha.3 h1:6evFIgj//Y3w09bqOUOzEpFj5tsxBqdc5CfkO7z+zfw=

tests/go.mod

@@ -37,7 +37,7 @@ require (
 	cloud.google.com/go/iam v0.13.0 // indirect
 	cloud.google.com/go/storage v1.30.0 // indirect
 	cosmossdk.io/client/v2 v2.0.0-20230309163709-87da587416ba // indirect
-	cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4 // indirect
+	cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e // indirect
 	cosmossdk.io/core v0.6.1 // indirect
 	filippo.io/edwards25519 v1.0.0 // indirect
 	github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect

tests/go.sum

@@ -190,8 +190,8 @@ cloud.google.com/go/workflows v1.6.0/go.mod h1:6t9F5h/unJz41YqfBmqSASJSXccBLtD1V
 cloud.google.com/go/workflows v1.7.0/go.mod h1:JhSrZuVZWuiDfKEFxU0/F1PQjmpnpcoISEXH2bcHC3M=
 cosmossdk.io/client/v2 v2.0.0-20230309163709-87da587416ba h1:LuPHCncU2KLMNPItFECs709uo46I9wSu2fAWYVCx+/U=
 cosmossdk.io/client/v2 v2.0.0-20230309163709-87da587416ba/go.mod h1:SXdwqO7cN5htalh/lhXWP8V4zKtBrhhcSTU+ytuEtmM=
-cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4 h1:QQZ0Qz8Gy/EmUNMRiHkUPG3BMA6OqEBp67IsfKETXIU=
-cosmossdk.io/collections v0.0.0-20230411101845-3d1a0b8840e4/go.mod h1:/vS4ugR7ad3IciUd5TQuP2Ldz3NukHK2u/l5xTxXbbE=
+cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e h1:3gV/RSkCybtvU7unzH/Tt7yI7W8O6q22Lg7FSwSP1K8=
+cosmossdk.io/collections v0.0.0-20230413101615-b81cd0ebb86e/go.mod h1:/vS4ugR7ad3IciUd5TQuP2Ldz3NukHK2u/l5xTxXbbE=
 cosmossdk.io/core v0.6.1 h1:OBy7TI2W+/gyn2z40vVvruK3di+cAluinA6cybFbE7s=
 cosmossdk.io/core v0.6.1/go.mod h1:g3MMBCBXtxbDWBURDVnJE7XML4BG5qENhs0gzkcpuFA=
 cosmossdk.io/depinject v1.0.0-alpha.3 h1:6evFIgj//Y3w09bqOUOzEpFj5tsxBqdc5CfkO7z+zfw=

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")