Add genesis changes

x/ccv/provider/keeper/genesis.go

@@ -77,6 +77,21 @@ func (k Keeper) InitGenesis(ctx sdk.Context, genState *types.GenesisState) {
 		}
 	}
 
+	if cs.KeyMap != nil {
+		for _, pkToCk := range cs.KeyMap.PkToCk {
+			k.KeyMap(ctx, cs.ChainId).Store.SetPkToCk(*pkToCk.From, *pkToCk.To)
+		}
+		for _, ckToPk := range cs.KeyMap.CkToPk {
+			k.KeyMap(ctx, cs.ChainId).Store.SetCkToPk(*ckToPk.From, *ckToPk.To)
+		}
+		for _, ckToMemo := range cs.KeyMap.CkToLastUpdateMemo {
+			k.KeyMap(ctx, cs.ChainId).Store.SetCkToMemo(*ckToMemo.Key, *ckToMemo.LastUpdateMemo)
+		}
+		for _, ccaToCk := range cs.KeyMap.CcaToCk {
+			k.KeyMap(ctx, cs.ChainId).Store.SetCcaToCk(ccaToCk.ConsAddr, *ccaToCk.Key)
+		}
+	}
+
 	k.SetParams(ctx, genState.Params)
 }
 
@@ -118,6 +133,33 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 			}
 		}
 
+		keyMap := func() *ccv.KeyMap {
+			km := &ccv.KeyMap{}
+			km.PkToCk = []ccv.KeyToKey{}
+			km.CkToPk = []ccv.KeyToKey{}
+			km.CkToLastUpdateMemo = []ccv.KeyToLastUpdateMemo{}
+			km.CcaToCk = []ccv.ConsAddrToKey{}
+			k.KeyMap(ctx, chainID).Store.IteratePkToCk(func(pk ProviderPubKey, ck ConsumerPubKey) bool {
+				km.PkToCk = append(km.PkToCk, ccv.KeyToKey{From: &pk, To: &ck})
+				return false
+			})
+			k.KeyMap(ctx, chainID).Store.IterateCkToPk(func(ck ConsumerPubKey, pk ProviderPubKey) bool {
+				km.CkToPk = append(km.CkToPk, ccv.KeyToKey{From: &ck, To: &pk})
+				return false
+			})
+			k.KeyMap(ctx, chainID).Store.IterateCkToMemo(func(ck ConsumerPubKey, m ccv.LastUpdateMemo) bool {
+				km.CkToLastUpdateMemo = append(km.CkToLastUpdateMemo, ccv.KeyToLastUpdateMemo{Key: &ck, LastUpdateMemo: &m})
+				return false
+			})
+			k.KeyMap(ctx, chainID).Store.IterateCcaToCk(func(cca ConsumerConsAddr, ck ConsumerPubKey) bool {
+				km.CcaToCk = append(km.CcaToCk, ccv.ConsAddrToKey{ConsAddr: cca, Key: &ck})
+				return false
+			})
+			return km
+		}
+
+		cs.KeyMap = keyMap()
+
 		consumerStates = append(consumerStates, cs)
 		return true
 	})

x/ccv/provider/keeper/genesis_test.go

@@ -15,6 +15,7 @@ import (
 	ccv "github.com/cosmos/interchain-security/x/ccv/types"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
+	crypto "github.com/tendermint/tendermint/proto/tendermint/crypto"
 )
 
 func TestIniAndExportGenesis(t *testing.T) {
@@ -25,6 +26,20 @@ func TestIniAndExportGenesis(t *testing.T) {
 	initHeight, vscID := uint64(5), uint64(1)
 	ubdIndex := []uint64{0, 1, 2}
 	params := providertypes.DefaultParams()
+	keyMaps := []ccv.KeyMap{
+		{
+			PkToCk:             []ccv.KeyToKey{{From: &crypto.PublicKey{}, To: &crypto.PublicKey{}}},
+			CkToPk:             []ccv.KeyToKey{},
+			CkToLastUpdateMemo: []ccv.KeyToLastUpdateMemo{},
+			CcaToCk:            []ccv.ConsAddrToKey{},
+		},
+		{
+			PkToCk:             []ccv.KeyToKey{},
+			CkToPk:             []ccv.KeyToKey{},
+			CkToLastUpdateMemo: []ccv.KeyToLastUpdateMemo{},
+			CcaToCk:            []ccv.ConsAddrToKey{},
+		},
+	}
 
 	// create genesis struct
 	pGenesis := providertypes.NewGenesisState(vscID,
@@ -42,6 +57,7 @@ func TestIniAndExportGenesis(t *testing.T) {
 				},
 				nil,
 				[]string{"slashedValidatorConsAddress"},
+				&keyMaps[0],
 			),
 			providertypes.NewConsumerStates(
 				cChainIDs[1],
@@ -53,6 +69,7 @@ func TestIniAndExportGenesis(t *testing.T) {
 				nil,
 				[]ccv.ValidatorSetChangePacketData{{ValsetUpdateId: vscID}},
 				nil,
+				&keyMaps[0],
 			),
 		},
 		[]ccv.UnbondingOp{{
@@ -104,6 +121,11 @@ func TestIniAndExportGenesis(t *testing.T) {
 	require.True(t, pk.GetPendingConsumerRemovalProp(ctx, cChainIDs[0], oneHourFromNow))
 	require.Equal(t, pGenesis.Params, pk.GetParams(ctx))
 
+	_, found = pk.KeyMap(ctx, cChainIDs[0]).Store.GetPkToCk(crypto.PublicKey{})
+	require.True(t, found)
+	_, found = pk.KeyMap(ctx, cChainIDs[1]).Store.GetPkToCk(crypto.PublicKey{})
+	require.False(t, found)
+
 	// check provider chain's consumer chain states
 	assertConsumerChainStates(ctx, t, pk, pGenesis.ConsumerStates...)