Revert Capabilities on Failed Tx (cosmos#5999)

Reintroduce memKVStore to keep track of fwd and reverse mappings.

On reverse mapping, rather than store a mapping to marshalled
capability; we store the index.

capability.Keeper and all scopedKeeper have access to a capability
map that maps index to the capability pointer.

This is done to make sure that all writes to memKVStore get reverted
on a fail tx, while also allowing GetCapability to retrieve the original
memory pointer from the go map.

Go map must be accessed only by first going through the
memKVStore. SInce writes to go map cannot be automatically
reverted on tx failure, it gets cleaned up on failed GetCapability calls.

Closes: cosmos#5965

baseapp/baseapp.go

@@ -187,6 +187,14 @@ func (app *BaseApp) MountTransientStores(keys map[string]*sdk.TransientStoreKey)
 	}
 }
 
+// MountMemoryStores mounts all in-memory KVStores with the BaseApp's internal
+// commit multi-store.
+func (app *BaseApp) MountMemoryStores(keys map[string]*sdk.MemoryStoreKey) {
+	for _, memKey := range keys {
+		app.MountStore(memKey, sdk.StoreTypeMemory)
+	}
+}
+
 // MountStoreWithDB mounts a store to the provided key in the BaseApp
 // multistore, using a specified DB.
 func (app *BaseApp) MountStoreWithDB(key sdk.StoreKey, typ sdk.StoreType, db dbm.DB) {

simapp/app.go

@@ -102,8 +102,9 @@ type SimApp struct {
 	invCheckPeriod uint
 
 	// keys to access the substores
-	keys  map[string]*sdk.KVStoreKey
-	tkeys map[string]*sdk.TransientStoreKey
+	keys    map[string]*sdk.KVStoreKey
+	tkeys   map[string]*sdk.TransientStoreKey
+	memKeys map[string]*sdk.MemoryStoreKey
 
 	// subspaces
 	subspaces map[string]params.Subspace
@@ -157,13 +158,15 @@ func NewSimApp(
 		evidence.StoreKey, transfer.StoreKey, capability.StoreKey,
 	)
 	tkeys := sdk.NewTransientStoreKeys(params.TStoreKey)
+	memKeys := sdk.NewMemoryStoreKeys(capability.MemStoreKey)
 
 	app := &SimApp{
 		BaseApp:        bApp,
 		cdc:            cdc,
 		invCheckPeriod: invCheckPeriod,
 		keys:           keys,
 		tkeys:          tkeys,
+		memKeys:        memKeys,
 		subspaces:      make(map[string]params.Subspace),
 	}
 
@@ -182,7 +185,7 @@ func NewSimApp(
 	bApp.SetParamStore(app.ParamsKeeper.Subspace(baseapp.Paramspace).WithKeyTable(std.ConsensusParamsKeyTable()))
 
 	// add capability keeper and ScopeToModule for ibc module
-	app.CapabilityKeeper = capability.NewKeeper(appCodec, keys[capability.StoreKey])
+	app.CapabilityKeeper = capability.NewKeeper(appCodec, keys[capability.StoreKey], memKeys[capability.MemStoreKey])
 	scopedIBCKeeper := app.CapabilityKeeper.ScopeToModule(ibc.ModuleName)
 	scopedTransferKeeper := app.CapabilityKeeper.ScopeToModule(transfer.ModuleName)
 
@@ -321,6 +324,7 @@ func NewSimApp(
 	// initialize stores
 	app.MountKVStores(keys)
 	app.MountTransientStores(tkeys)
+	app.MountMemoryStores(memKeys)
 
 	// initialize BaseApp
 	app.SetInitChainer(app.InitChainer)
@@ -418,6 +422,13 @@ func (app *SimApp) GetTKey(storeKey string) *sdk.TransientStoreKey {
 	return app.tkeys[storeKey]
 }
 
+// GetMemKey returns the MemStoreKey for the provided mem key.
+//
+// NOTE: This is solely used for testing purposes.
+func (app *SimApp) GetMemKey(storeKey string) *sdk.MemoryStoreKey {
+	return app.memKeys[storeKey]
+}
+
 // GetSubspace returns a param subspace for a given module name.
 //
 // NOTE: This is solely to be used for testing purposes.

store/mem/mem_test.go

@@ -0,0 +1,39 @@
+package mem_test
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/cosmos/cosmos-sdk/store/mem"
+	"github.com/cosmos/cosmos-sdk/store/types"
+)
+
+func TestStore(t *testing.T) {
+	db := mem.NewStore()
+	key, value := []byte("key"), []byte("value")
+
+	require.Equal(t, types.StoreTypeMemory, db.GetStoreType())
+
+	require.Nil(t, db.Get(key))
+	db.Set(key, value)
+	require.Equal(t, value, db.Get(key))
+
+	newValue := []byte("newValue")
+	db.Set(key, newValue)
+	require.Equal(t, newValue, db.Get(key))
+
+	db.Delete(key)
+	require.Nil(t, db.Get(key))
+}
+
+func TestCommit(t *testing.T) {
+	db := mem.NewStore()
+	key, value := []byte("key"), []byte("value")
+
+	db.Set(key, value)
+	id := db.Commit()
+	require.True(t, id.IsZero())
+	require.True(t, db.LastCommitID().IsZero())
+	require.Equal(t, value, db.Get(key))
+}

store/mem/store.go

@@ -0,0 +1,53 @@
+package mem
+
+import (
+	"io"
+
+	dbm "github.com/tendermint/tm-db"
+
+	"github.com/cosmos/cosmos-sdk/store/cachekv"
+	"github.com/cosmos/cosmos-sdk/store/dbadapter"
+	"github.com/cosmos/cosmos-sdk/store/tracekv"
+	"github.com/cosmos/cosmos-sdk/store/types"
+)
+
+var (
+	_ types.KVStore   = (*Store)(nil)
+	_ types.Committer = (*Store)(nil)
+)
+
+// Store implements an in-memory only KVStore. Entries are persisted between
+// commits and thus between blocks. State in Memory store is not committed as part of app state but maintained privately by each node
+type Store struct {
+	dbadapter.Store
+}
+
+func NewStore() *Store {
+	return NewStoreWithDB(dbm.NewMemDB())
+}
+
+func NewStoreWithDB(db *dbm.MemDB) *Store { // nolint: interfacer
+	return &Store{Store: dbadapter.Store{DB: db}}
+}
+
+// GetStoreType returns the Store's type.
+func (s Store) GetStoreType() types.StoreType {
+	return types.StoreTypeMemory
+}
+
+// CacheWrap cache wraps the underlying store.
+func (s Store) CacheWrap() types.CacheWrap {
+	return cachekv.NewStore(s)
+}
+
+// CacheWrapWithTrace implements KVStore.
+func (s Store) CacheWrapWithTrace(w io.Writer, tc types.TraceContext) types.CacheWrap {
+	return cachekv.NewStore(tracekv.NewStore(s, w, tc))
+}
+
+// Commit performs a no-op as entries are persistent between commitments.
+func (s *Store) Commit() (id types.CommitID) { return }
+
+// nolint
+func (s *Store) SetPruning(pruning types.PruningOptions) {}
+func (s Store) LastCommitID() (id types.CommitID)        { return }

store/rootmulti/store.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/store/cachemulti"
 	"github.com/cosmos/cosmos-sdk/store/dbadapter"
 	"github.com/cosmos/cosmos-sdk/store/iavl"
+	"github.com/cosmos/cosmos-sdk/store/mem"
 	"github.com/cosmos/cosmos-sdk/store/tracekv"
 	"github.com/cosmos/cosmos-sdk/store/transient"
 	"github.com/cosmos/cosmos-sdk/store/types"
@@ -528,6 +529,13 @@ func (rs *Store) loadCommitStoreFromParams(key types.StoreKey, id types.CommitID
 
 		return transient.NewStore(), nil
 
+	case types.StoreTypeMemory:
+		if _, ok := key.(*types.MemoryStoreKey); !ok {
+			return nil, fmt.Errorf("unexpected key type for a MemoryStoreKey; got: %s", key.String())
+		}
+
+		return mem.NewStore(), nil
+
 	default:
 		panic(fmt.Sprintf("unrecognized store type %v", params.typ))
 	}

store/types/store.go

@@ -273,6 +273,7 @@ const (
 	StoreTypeDB
 	StoreTypeIAVL
 	StoreTypeTransient
+	StoreTypeMemory
 )
 
 func (st StoreType) String() string {
@@ -288,6 +289,9 @@ func (st StoreType) String() string {
 
 	case StoreTypeTransient:
 		return "StoreTypeTransient"
+
+	case StoreTypeMemory:
+		return "StoreTypeMemory"
 	}
 
 	return "unknown store type"
@@ -351,6 +355,25 @@ func (key *TransientStoreKey) String() string {
 	return fmt.Sprintf("TransientStoreKey{%p, %s}", key, key.name)
 }
 
+// MemoryStoreKey defines a typed key to be used with an in-memory KVStore.
+type MemoryStoreKey struct {
+	name string
+}
+
+func NewMemoryStoreKey(name string) *MemoryStoreKey {
+	return &MemoryStoreKey{name: name}
+}
+
+// Name returns the name of the MemoryStoreKey.
+func (key *MemoryStoreKey) Name() string {
+	return key.name
+}
+
+// String returns a stringified representation of the MemoryStoreKey.
+func (key *MemoryStoreKey) String() string {
+	return fmt.Sprintf("MemoryStoreKey{%p, %s}", key, key.name)
+}
+
 //----------------------------------------
 
 // key-value result for iterator queries

types/store.go

@@ -76,6 +76,7 @@ const (
 	StoreTypeDB        = types.StoreTypeDB
 	StoreTypeIAVL      = types.StoreTypeIAVL
 	StoreTypeTransient = types.StoreTypeTransient
+	StoreTypeMemory    = types.StoreTypeMemory
 )
 
 // nolint - reexport
@@ -84,6 +85,7 @@ type (
 	CapabilityKey     = types.CapabilityKey
 	KVStoreKey        = types.KVStoreKey
 	TransientStoreKey = types.TransientStoreKey
+	MemoryStoreKey    = types.MemoryStoreKey
 )
 
 // NewKVStoreKey returns a new pointer to a KVStoreKey.
@@ -120,6 +122,17 @@ func NewTransientStoreKeys(names ...string) map[string]*TransientStoreKey {
 	return keys
 }
 
+// NewMemoryStoreKeys constructs a new map matching store key names to their
+// respective MemoryStoreKey references.
+func NewMemoryStoreKeys(names ...string) map[string]*MemoryStoreKey {
+	keys := make(map[string]*MemoryStoreKey)
+	for _, name := range names {
+		keys[name] = types.NewMemoryStoreKey(name)
+	}
+
+	return keys
+}
+
 // PrefixEndBytes returns the []byte that would end a
 // range query for all []byte with a certain prefix
 // Deals with last byte of prefix being FF without overflowing

x/capability/alias.go

@@ -29,7 +29,6 @@ var (
 	ModuleCdc                = types.ModuleCdc
 	NewOwner                 = types.NewOwner
 	NewCapabilityOwners      = types.NewCapabilityOwners
-	NewCapabilityStore       = types.NewCapabilityStore
 )
 
 // nolint
@@ -38,5 +37,4 @@ type (
 	ScopedKeeper     = keeper.ScopedKeeper
 	Capability       = types.Capability
 	CapabilityOwners = types.CapabilityOwners
-	CapabilityStore  = types.CapabilityStore
 )