From f82bc19b9929f5e7c84e55d3ab08d2d504290c49 Mon Sep 17 00:00:00 2001
From: Jack Zampolin <jack.zampolin@gmail.com>
Date: Wed, 22 Apr 2020 10:49:09 -0700
Subject: [PATCH] Fix capabilities issue (#6057)

* Fix capabilities issue

* Update x/capability/module.go

quick doc fix

* Address PR comments

Co-authored-by: Aditya <adityasripal@gmail.com>
---
 x/capability/keeper/keeper.go | 12 ++++--------
 x/capability/module.go        | 19 +++++++++++++------
 x/capability/types/genesis.go | 11 +++++++++++
 3 files changed, 28 insertions(+), 14 deletions(-)
 create mode 100644 x/capability/types/genesis.go

diff --git a/x/capability/keeper/keeper.go b/x/capability/keeper/keeper.go
index 48c07ac77..138f500e4 100644
--- a/x/capability/keeper/keeper.go
+++ b/x/capability/keeper/keeper.go
@@ -130,16 +130,12 @@ func (k *Keeper) InitializeAndSeal(ctx sdk.Context) {
 	k.sealed = true
 }
 
-// InitializeIndex sets the index to one in InitChain
+// SetIndex sets the index to one in InitChain
 // Since it is an exported function, we check that index is indeed unset, before initializing
-func (k Keeper) InitializeIndex(ctx sdk.Context) {
-	// set the global index to start at 1 if it is unset
-	index := k.GetLatestIndex(ctx)
-	if index != 0 {
-		return
-	}
+func (k Keeper) SetIndex(ctx sdk.Context, index uint64) {
+	// set the global index to the passed index
 	store := ctx.KVStore(k.storeKey)
-	store.Set(types.KeyIndex, types.IndexToKey(1))
+	store.Set(types.KeyIndex, types.IndexToKey(index))
 }
 
 // GetLatestIndex returns the latest index of the CapabilityKeeper
diff --git a/x/capability/module.go b/x/capability/module.go
index c2a3fb18e..3c8f2117a 100644
--- a/x/capability/module.go
+++ b/x/capability/module.go
@@ -7,6 +7,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	"github.com/cosmos/cosmos-sdk/x/capability/types"
 
 	"github.com/gorilla/mux"
 	"github.com/spf13/cobra"
@@ -44,7 +45,9 @@ func (AppModuleBasic) RegisterCodec(cdc *codec.Codec) {
 }
 
 // DefaultGenesis returns the capability module's default genesis state.
-func (AppModuleBasic) DefaultGenesis(_ codec.JSONMarshaler) json.RawMessage { return []byte("{}") }
+func (AppModuleBasic) DefaultGenesis(cdc codec.JSONMarshaler) json.RawMessage {
+	return cdc.MustMarshalJSON(types.DefaultGenesis())
+}
 
 // ValidateGenesis performs genesis state validation for the capability module.
 func (AppModuleBasic) ValidateGenesis(_ codec.JSONMarshaler, _ json.RawMessage) error { return nil }
@@ -98,16 +101,20 @@ func (am AppModule) RegisterInvariants(_ sdk.InvariantRegistry) {}
 
 // InitGenesis performs the capability module's genesis initialization It returns
 // no validator updates.
-func (am AppModule) InitGenesis(ctx sdk.Context, _ codec.JSONMarshaler, _ json.RawMessage) []abci.ValidatorUpdate {
-	// Initialize global index to 1
-	am.keeper.InitializeIndex(ctx)
+func (am AppModule) InitGenesis(ctx sdk.Context, cdc codec.JSONMarshaler, gs json.RawMessage) []abci.ValidatorUpdate {
+	var genState types.GenesisState
+	// Initialize global index to index in genesis state
+	cdc.MustUnmarshalJSON(gs, &genState)
+
+	am.keeper.SetIndex(ctx, genState.Index)
 
 	return []abci.ValidatorUpdate{}
 }
 
 // ExportGenesis returns the capability module's exported genesis state as raw JSON bytes.
-func (am AppModule) ExportGenesis(_ sdk.Context, cdc codec.JSONMarshaler) json.RawMessage {
-	return am.DefaultGenesis(cdc)
+func (am AppModule) ExportGenesis(ctx sdk.Context, cdc codec.JSONMarshaler) json.RawMessage {
+	index := am.keeper.GetLatestIndex(ctx)
+	return cdc.MustMarshalJSON(types.GenesisState{index})
 }
 
 // BeginBlock executes all ABCI BeginBlock logic respective to the capability module.
diff --git a/x/capability/types/genesis.go b/x/capability/types/genesis.go
new file mode 100644
index 000000000..85c7f4162
--- /dev/null
+++ b/x/capability/types/genesis.go
@@ -0,0 +1,11 @@
+package types
+
+// GenesisState represents the Capability module genesis state
+type GenesisState struct {
+	Index uint64 `json:"index" yaml:"index"`
+}
+
+// DefaultGenesis returns the default Capability genesis state
+func DefaultGenesis() GenesisState {
+	return GenesisState{Index: 1}
+}