feat(accounts): add Genesis handling (#17802)

Co-authored-by: unknown unknown <unknown@unknown>

proto/cosmos/accounts/v1/genesis.proto

@@ -0,0 +1,31 @@
+syntax = "proto3";
+
+package cosmos.accounts.v1;
+
+option go_package = "cosmossdk.io/x/accounts/v1";
+
+// GenesisState defines the accounts' module's genesis state.
+message GenesisState {
+  // account_number is the latest account number.
+  uint64 account_number = 1;
+  // accounts are the genesis accounts.
+  repeated GenesisAccount accounts = 2;
+}
+
+// GenesisAccount defines an account to be initialized in the genesis state.
+message GenesisAccount {
+  // address is the address of the account.
+  string address = 1;
+  // account_type is the account type of the account.
+  string account_type = 2;
+  // state is the account state represented as a slice of raw key value byte pairs.
+  repeated KVPair state = 3;
+}
+
+// KVPair defines a key value pair.
+message KVPair {
+  // key is the key of the pair.
+  bytes key = 1;
+  // value is the value of the pair.
+  bytes value = 2;
+}

x/accounts/account_test.go

@@ -9,12 +9,21 @@ import (
 
 	bankv1beta1 "cosmossdk.io/api/cosmos/bank/v1beta1"
 	basev1beta1 "cosmossdk.io/api/cosmos/base/v1beta1"
+	"cosmossdk.io/collections"
 	"cosmossdk.io/x/accounts/internal/implementation"
 )
 
 var _ implementation.Account = (*TestAccount)(nil)
 
-type TestAccount struct{}
+func NewTestAccount(sb *collections.SchemaBuilder) *TestAccount {
+	return &TestAccount{
+		Counter: collections.NewSequence(sb, collections.NewPrefix(0), "counter"),
+	}
+}
+
+type TestAccount struct {
+	Counter collections.Sequence
+}
 
 func (t TestAccount) RegisterInitHandler(builder *implementation.InitBuilder) {
 	implementation.RegisterInitHandler(builder, func(ctx context.Context, _ *emptypb.Empty) (*emptypb.Empty, error) {
@@ -62,6 +71,11 @@ func (t TestAccount) RegisterExecuteHandlers(builder *implementation.ExecuteBuil
 
 		return &emptypb.Empty{}, nil
 	})
+
+	// genesis testing
+	implementation.RegisterExecuteHandler(builder, func(ctx context.Context, req *wrapperspb.UInt64Value) (*emptypb.Empty, error) {
+		return new(emptypb.Empty), t.Counter.Set(ctx, req.Value)
+	})
 }
 
 func (t TestAccount) RegisterQueryHandlers(builder *implementation.QueryBuilder) {
@@ -90,4 +104,14 @@ func (t TestAccount) RegisterQueryHandlers(builder *implementation.QueryBuilder)
 		}
 		return wrapperspb.Int64(amt), nil
 	})
+
+	// genesis testing; DoubleValue does not make sense as a request type for this query, but empty is already taken
+	// and this is only used for testing.
+	implementation.RegisterQueryHandler(builder, func(ctx context.Context, _ *wrapperspb.DoubleValue) (*wrapperspb.UInt64Value, error) {
+		v, err := t.Counter.Peek(ctx)
+		if err != nil {
+			return nil, err
+		}
+		return &wrapperspb.UInt64Value{Value: v}, nil
+	})
 }

x/accounts/genesis.go

@@ -0,0 +1,94 @@
+package accounts
+
+import (
+	"context"
+	"fmt"
+
+	"cosmossdk.io/collections"
+	v1 "cosmossdk.io/x/accounts/v1"
+)
+
+func (k Keeper) ExportState(ctx context.Context) (*v1.GenesisState, error) {
+	genState := &v1.GenesisState{}
+
+	// get account number
+	accountNumber, err := k.AccountNumber.Peek(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	genState.AccountNumber = accountNumber
+
+	err = k.AccountsByType.Walk(ctx, nil, func(key []byte, value string) (stop bool, err error) {
+		accState, err := k.exportAccount(ctx, key, value)
+		if err != nil {
+			return true, err
+		}
+		genState.Accounts = append(genState.Accounts, accState)
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return genState, nil
+}
+
+func (k Keeper) exportAccount(ctx context.Context, addr []byte, accType string) (*v1.GenesisAccount, error) {
+	addrString, err := k.addressCodec.BytesToString(addr)
+	if err != nil {
+		return nil, err
+	}
+	account := &v1.GenesisAccount{
+		Address:     addrString,
+		AccountType: accType,
+	}
+	rng := new(collections.Range[[]byte]).
+		Prefix(addr)
+	err = k.AccountsState.Walk(ctx, rng, func(key, value []byte) (stop bool, err error) {
+		account.State = append(account.State, &v1.KVPair{
+			Key:   key,
+			Value: value,
+		})
+		return false, nil
+	})
+	if err != nil {
+		return nil, err
+	}
+	return account, nil
+}
+
+func (k Keeper) ImportState(ctx context.Context, genState *v1.GenesisState) error {
+	err := k.AccountNumber.Set(ctx, genState.AccountNumber)
+	if err != nil {
+		return err
+	}
+
+	// import accounts
+	for _, acc := range genState.Accounts {
+		err = k.importAccount(ctx, acc)
+		if err != nil {
+			return fmt.Errorf("%w: %s", err, acc.Address)
+		}
+	}
+	return nil
+}
+
+func (k Keeper) importAccount(ctx context.Context, acc *v1.GenesisAccount) error {
+	// TODO: maybe check if impl exists?
+	addrBytes, err := k.addressCodec.StringToBytes(acc.Address)
+	if err != nil {
+		return err
+	}
+	err = k.AccountsByType.Set(ctx, addrBytes, acc.AccountType)
+	if err != nil {
+		return err
+	}
+	for _, kv := range acc.State {
+		err = k.AccountsState.Set(ctx, kv.Key, kv.Value)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}

x/accounts/genesis_test.go

@@ -0,0 +1,61 @@
+package accounts
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/emptypb"
+	"google.golang.org/protobuf/types/known/wrapperspb"
+
+	bankv1beta1 "cosmossdk.io/api/cosmos/bank/v1beta1"
+	"cosmossdk.io/collections"
+	"cosmossdk.io/collections/colltest"
+	"cosmossdk.io/x/accounts/internal/implementation"
+)
+
+func TestGenesis(t *testing.T) {
+	sb := collections.NewSchemaBuilderFromAccessor(implementation.OpenKVStore)
+	acc := NewTestAccount(sb)
+
+	k, ctx := newKeeper(t, map[string]implementation.Account{
+		"test": acc,
+	})
+	k.queryModuleFunc = func(ctx context.Context, _ proto.Message) (proto.Message, error) {
+		return &bankv1beta1.QueryBalanceResponse{}, nil
+	}
+
+	// we init two accounts of the same type
+
+	// we set counter to 10
+	_, addr1, err := k.Init(ctx, "test", []byte("sender"), &emptypb.Empty{})
+	require.NoError(t, err)
+	_, err = k.Execute(ctx, addr1, []byte("sender"), &wrapperspb.UInt64Value{Value: 10})
+	require.NoError(t, err)
+
+	// we set counter to 20
+	_, addr2, err := k.Init(ctx, "test", []byte("sender"), &emptypb.Empty{})
+	require.NoError(t, err)
+	_, err = k.Execute(ctx, addr2, []byte("sender"), &wrapperspb.UInt64Value{Value: 20})
+	require.NoError(t, err)
+
+	// export state
+	state, err := k.ExportState(ctx)
+	require.NoError(t, err)
+
+	// reset state
+	_, ctx = colltest.MockStore()
+	err = k.ImportState(ctx, state)
+	require.NoError(t, err)
+
+	// if genesis import went fine, we should be able to query the accounts
+	// and get the expected values.
+	resp, err := k.Query(ctx, addr1, &wrapperspb.DoubleValue{})
+	require.NoError(t, err)
+	require.Equal(t, &wrapperspb.UInt64Value{Value: 10}, resp)
+
+	resp, err = k.Query(ctx, addr2, &wrapperspb.DoubleValue{})
+	require.NoError(t, err)
+	require.Equal(t, &wrapperspb.UInt64Value{Value: 20}, resp)
+}

x/accounts/internal/implementation/account_test.go

@@ -46,7 +46,7 @@ func (t TestAccount) RegisterExecuteHandlers(builder *ExecuteBuilder) {
 	})
 }
 
-func (TestAccount) RegisterQueryHandlers(builder *QueryBuilder) {
+func (t TestAccount) RegisterQueryHandlers(builder *QueryBuilder) {
 	RegisterQueryHandler(builder, func(_ context.Context, req *wrapperspb.StringValue) (*wrapperspb.StringValue, error) {
 		return &wrapperspb.StringValue{Value: req.Value + "query-echo"}, nil
 	})

x/accounts/internal/implementation/context.go

@@ -8,29 +8,35 @@ import (
 
 	"google.golang.org/protobuf/proto"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/core/store"
 	"cosmossdk.io/x/accounts/internal/prefixstore"
 )
 
-var errUnauthorized = errors.New("unauthorized")
+var (
+	errUnauthorized    = errors.New("unauthorized")
+	AccountStatePrefix = collections.NewPrefix(255)
+)
 
 type contextKey struct{}
 
 type contextValue struct {
-	store             store.KVStore   // store is the prefixed store for the account.
-	sender            []byte          // sender is the address of the entity invoking the account action.
-	whoami            []byte          // whoami is the address of the account being invoked.
-	originalContext   context.Context // originalContext that was used to build the account context.
-	getExpectedSender func(msg proto.Message) ([]byte, error)
-	moduleExec        func(ctx context.Context, msg proto.Message) (proto.Message, error)
-	moduleQuery       func(ctx context.Context, msg proto.Message) (proto.Message, error)
+	store             store.KVStore                                                       // store is the prefixed store for the account.
+	sender            []byte                                                              // sender is the address of the entity invoking the account action.
+	whoami            []byte                                                              // whoami is the address of the account being invoked.
+	originalContext   context.Context                                                     // originalContext that was used to build the account context.
+	getExpectedSender func(msg proto.Message) ([]byte, error)                             // getExpectedSender is a function that returns the expected sender for a given message.
+	moduleExec        func(ctx context.Context, msg proto.Message) (proto.Message, error) // moduleExec is a function that executes a module message.
+	moduleQuery       func(ctx context.Context, msg proto.Message) (proto.Message, error) // moduleQuery is a function that queries a module.
 }
 
 // MakeAccountContext creates a new account execution context given:
 // storeSvc: which fetches the x/accounts module store.
 // accountAddr: the address of the account being invoked, which is used to give the
 // account a prefixed storage.
 // sender: the address of entity invoking the account action.
+// moduleExec: a function that executes a module message.
+// moduleQuery: a function that queries a module.
 func MakeAccountContext(
 	ctx context.Context,
 	storeSvc store.KVStoreService,
@@ -41,7 +47,7 @@ func MakeAccountContext(
 	moduleQuery func(ctx context.Context, msg proto.Message) (proto.Message, error),
 ) context.Context {
 	return context.WithValue(ctx, contextKey{}, contextValue{
-		store:             prefixstore.New(storeSvc.OpenKVStore(ctx), accountAddr),
+		store:             prefixstore.New(storeSvc.OpenKVStore(ctx), append(AccountStatePrefix, accountAddr...)),
 		sender:            sender,
 		whoami:            accountAddr,
 		originalContext:   ctx,

x/accounts/internal/implementation/context_test.go

@@ -38,8 +38,8 @@ func TestMakeAccountContext(t *testing.T) {
 	// this store is the global x/accounts module store.
 	store := storeService.OpenKVStore(originalContext)
 
-	// now we want the value to be store in the following accounts prefix (accountAddr + itemPrefix)
-	value, err := store.Get(append(accountAddr, itemPrefix...))
+	// now we want the value to be store in the following accounts prefix (AccountsStatePrefix + accountAddr + itemPrefix)
+	value, err := store.Get(append(AccountStatePrefix, append(accountAddr, itemPrefix...)...))
 	require.NoError(t, err)
 	require.Equal(t, []byte{0, 0, 0, 0, 0, 0, 3, 232}, value)
 

x/accounts/keeper.go

@@ -43,6 +43,7 @@ func NewKeeper(
 		Schema:          collections.Schema{},
 		AccountNumber:   collections.NewSequence(sb, AccountNumberKey, "account_number"),
 		AccountsByType:  collections.NewMap(sb, AccountTypeKeyPrefix, "accounts_by_type", collections.BytesKey, collections.StringValue),
+		AccountsState:   collections.NewMap(sb, implementation.AccountStatePrefix, "accounts_state", collections.BytesKey, collections.BytesValue),
 	}
 
 	// make accounts implementation
@@ -77,6 +78,12 @@ type Keeper struct {
 	AccountNumber collections.Sequence
 	// AccountsByType maps account address to their implementation.
 	AccountsByType collections.Map[[]byte, string]
+
+	// AccountsState keeps track of the state of each account.
+	// NOTE: this is only used for genesis import and export.
+	// Contracts set and get their own state but this helps providing a nice mapping
+	// between: (account address, account state key) => account state value.
+	AccountsState collections.Map[[]byte, []byte]
 }
 
 // Init creates a new account of the given type.