refactor(x/staking): refactor keeper and msg server structs and add p…

…ubkey invariant

app/app.go

@@ -135,6 +135,7 @@ import (
 	"github.com/sedaprotocol/seda-chain/app/keepers"
 	appparams "github.com/sedaprotocol/seda-chain/app/params"
 	"github.com/sedaprotocol/seda-chain/app/utils"
+
 	// Used in cosmos-sdk when registering the route for swagger docs.
 	_ "github.com/sedaprotocol/seda-chain/client/docs/statik"
 	"github.com/sedaprotocol/seda-chain/cmd/sedad/gentx"
@@ -426,7 +427,11 @@ func NewApp(
 		authcodec.NewBech32Codec(sdk.GetConfig().GetBech32ValidatorAddrPrefix()),
 		authcodec.NewBech32Codec(sdk.GetConfig().GetBech32ConsensusAddrPrefix()),
 	)
-	app.StakingKeeper = stakingkeeper.NewKeeper(sdkStakingKeeper)
+	app.StakingKeeper = stakingkeeper.NewKeeper(
+		sdkStakingKeeper,
+		app.PubKeyKeeper,
+		app.AccountKeeper.AddressCodec(),
+	)
 
 	app.FeeGrantKeeper = feegrantkeeper.NewKeeper(
 		appCodec,

cmd/sedad/gentx/gentx.go

@@ -94,10 +94,6 @@ $ %s gentx my-key-name 1000000seda --home=/path/to/home/dir --keyring-backend=os
 			if err != nil {
 				return errors.Wrap(err, "failed to initialize node validator files")
 			}
-			// vrfPubKey, err := utils.LoadOrGenVRFKey(serverCtx.Config, "") // TODO (#314)
-			// if err != nil {
-			// 	return errors.Wrap(err, "failed to initialize VRF key")
-			// }
 
 			// read --nodeID, if empty take it from priv_validator.json
 			if nodeIDString, _ := cmd.Flags().GetString(cli.FlagNodeID); nodeIDString != "" {

x/batching/keeper/integration_test.go

@@ -140,6 +140,7 @@ func initFixture(tb testing.TB) *fixture {
 		log.NewNopLogger(),
 	)
 
+	var pubKeyKeeper *pubkeykeeper.Keeper
 	sdkStakingKeeper := sdkstakingkeeper.NewKeeper(
 		cdc,
 		runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]),
@@ -149,7 +150,11 @@ func initFixture(tb testing.TB) *fixture {
 		addresscodec.NewBech32Codec(params.Bech32PrefixValAddr),
 		addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr),
 	)
-	stakingKeeper := stakingkeeper.NewKeeper(sdkStakingKeeper)
+	stakingKeeper := stakingkeeper.NewKeeper(
+		sdkStakingKeeper,
+		pubKeyKeeper,
+		addresscodec.NewBech32Codec(params.Bech32PrefixValAddr),
+	)
 
 	stakingParams := sdkstakingtypes.DefaultParams()
 	stakingParams.BondDenom = bondDenom
@@ -196,7 +201,7 @@ func initFixture(tb testing.TB) *fixture {
 		viewKeeper,
 	)
 
-	pubKeyKeeper := pubkeykeeper.NewKeeper(
+	pubKeyKeeper = pubkeykeeper.NewKeeper(
 		cdc,
 		runtime.NewKVStoreService(keys[pubkeytypes.StoreKey]),
 		stakingKeeper,

x/pubkey/keeper/endblock_test.go

@@ -45,9 +45,11 @@ import (
 	"github.com/sedaprotocol/seda-chain/integration"
 	"github.com/sedaprotocol/seda-chain/x/pubkey"
 	"github.com/sedaprotocol/seda-chain/x/pubkey/keeper"
+	pubkeykeeper "github.com/sedaprotocol/seda-chain/x/pubkey/keeper"
 	"github.com/sedaprotocol/seda-chain/x/pubkey/types"
 	"github.com/sedaprotocol/seda-chain/x/staking"
 	stakingkeeper "github.com/sedaprotocol/seda-chain/x/staking/keeper"
+	stakingtypes "github.com/sedaprotocol/seda-chain/x/staking/types"
 	"github.com/sedaprotocol/seda-chain/x/vesting"
 )
 
@@ -125,8 +127,9 @@ func initFixture(tb testing.TB) *fixture {
 		log.NewNopLogger(),
 	)
 
-	sdkstakingKeeper := sdkstakingkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(params.Bech32PrefixValAddr), addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr))
-	stakingKeeper := stakingkeeper.NewKeeper(sdkstakingKeeper)
+	var pubKeyKeeper *pubkeykeeper.Keeper
+	sdkStakingKeeper := sdkstakingkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(params.Bech32PrefixValAddr), addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr))
+	stakingKeeper := stakingkeeper.NewKeeper(sdkStakingKeeper, pubKeyKeeper, addresscodec.NewBech32Codec(params.Bech32PrefixValAddr))
 
 	stakingParams := sdkstakingtypes.DefaultParams()
 	stakingParams.BondDenom = bondDenom
@@ -141,7 +144,7 @@ func initFixture(tb testing.TB) *fixture {
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
-	pubkeyKeeper := keeper.NewKeeper(
+	pubKeyKeeper = keeper.NewKeeper(
 		cdc,
 		runtime.NewKVStoreService(keys[types.StoreKey]),
 		stakingKeeper,
@@ -152,8 +155,8 @@ func initFixture(tb testing.TB) *fixture {
 
 	authModule := auth.NewAppModule(cdc, accountKeeper, app.RandomGenesisAccounts, nil)
 	bankModule := bank.NewAppModule(cdc, bankKeeper, accountKeeper, nil)
-	stakingModule := staking.NewAppModule(cdc, stakingKeeper, accountKeeper, bankKeeper, pubkeyKeeper)
-	pubkeyModule := pubkey.NewAppModule(cdc, *pubkeyKeeper)
+	stakingModule := staking.NewAppModule(cdc, stakingKeeper, accountKeeper, bankKeeper, pubKeyKeeper)
+	pubkeyModule := pubkey.NewAppModule(cdc, *pubKeyKeeper)
 
 	integrationApp := integration.NewIntegrationApp(newCtx, logger, keys, cdc, map[string]appmodule.AppModule{
 		authtypes.ModuleName:       authModule,
@@ -162,16 +165,18 @@ func initFixture(tb testing.TB) *fixture {
 		types.ModuleName:           pubkeyModule,
 	})
 
-	types.RegisterMsgServer(integrationApp.MsgServiceRouter(), keeper.NewMsgServerImpl(*pubkeyKeeper))
-	sdkstakingtypes.RegisterMsgServer(integrationApp.MsgServiceRouter(), sdkstakingkeeper.NewMsgServerImpl(sdkstakingKeeper))
+	types.RegisterMsgServer(integrationApp.MsgServiceRouter(), keeper.NewMsgServerImpl(*pubKeyKeeper))
+	sdkStakingMsgServer := sdkstakingkeeper.NewMsgServerImpl(sdkStakingKeeper)
+	stakingMsgServer := stakingkeeper.NewMsgServerImpl(sdkStakingMsgServer, stakingKeeper)
+	stakingtypes.RegisterMsgServer(integrationApp.MsgServiceRouter(), stakingMsgServer)
 
 	return &fixture{
 		IntegationApp: integrationApp,
 		cdc:           cdc,
 		accountKeeper: accountKeeper,
 		bankKeeper:    bankKeeper,
 		stakingKeeper: *stakingKeeper,
-		keeper:        *pubkeyKeeper,
+		keeper:        *pubKeyKeeper,
 	}
 }
 

x/staking/keeper/keeper.go

@@ -5,24 +5,31 @@ import (
 	"fmt"
 	"math"
 
+	addresscodec "cosmossdk.io/core/address"
 	sdkmath "cosmossdk.io/math"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
-	"github.com/cosmos/cosmos-sdk/x/staking/types"
+	sdktypes "github.com/cosmos/cosmos-sdk/x/staking/types"
+
+	"github.com/sedaprotocol/seda-chain/x/staking/types"
 )
 
 type Keeper struct {
 	*sdkkeeper.Keeper
+	pubKeyKeeper          types.PubKeyKeeper
+	validatorAddressCodec addresscodec.Codec
 }
 
-func NewKeeper(sdkStakingKeeper *sdkkeeper.Keeper) *Keeper {
+func NewKeeper(sdkStakingKeeper *sdkkeeper.Keeper, pubKeyKeeper types.PubKeyKeeper, valAddrCdc addresscodec.Codec) *Keeper {
 	return &Keeper{
-		Keeper: sdkStakingKeeper,
+		Keeper:                sdkStakingKeeper,
+		pubKeyKeeper:          pubKeyKeeper,
+		validatorAddressCodec: valAddrCdc,
 	}
 }
 
-func (k *Keeper) SetHooks(sh types.StakingHooks) {
+func (k *Keeper) SetHooks(sh sdktypes.StakingHooks) {
 	k.Keeper.SetHooks(sh)
 }
 
@@ -55,7 +62,7 @@ func (k Keeper) TransferDelegation(ctx context.Context, fromAddr, toAddr sdk.Acc
 	// Assume the worst case that we need to transfer all redelegation entries
 	mightExceedLimit := false
 	var cbErr error
-	err = k.IterateDelegatorRedelegations(ctx, fromAddr, func(toRedelegation types.Redelegation) (stop bool) {
+	err = k.IterateDelegatorRedelegations(ctx, fromAddr, func(toRedelegation sdktypes.Redelegation) (stop bool) {
 		// There's no redelegation index by delegator and dstVal or vice-versa.
 		// The minimum cardinality is to look up by delegator, so scan and skip.
 		if toRedelegation.ValidatorDstAddress != valAddr.String() {
@@ -90,7 +97,7 @@ func (k Keeper) TransferDelegation(ctx context.Context, fromAddr, toAddr sdk.Acc
 		return transferred, err
 	}
 	if mightExceedLimit {
-		// avoid types.ErrMaxRedelegationEntries
+		// avoid sdktypes.ErrMaxRedelegationEntries
 		return transferred, nil
 	}
 
@@ -104,8 +111,8 @@ func (k Keeper) TransferDelegation(ctx context.Context, fromAddr, toAddr sdk.Acc
 	// Update or create the delTo object, calling appropriate hooks
 	delTo, err := k.GetDelegation(ctx, toAddr, valAddr)
 	if err != nil {
-		if err == types.ErrNoDelegation {
-			delTo = types.NewDelegation(toAddr.String(), validator.GetOperator(), sdkmath.LegacyZeroDec())
+		if err == sdktypes.ErrNoDelegation {
+			delTo = sdktypes.NewDelegation(toAddr.String(), validator.GetOperator(), sdkmath.LegacyZeroDec())
 			err = k.Hooks().BeforeDelegationCreated(ctx, toAddr, valAddr)
 		} else {
 			return transferred, err

x/staking/keeper/msg_server.go

@@ -3,8 +3,6 @@ package keeper
 import (
 	"context"
 
-	addresscodec "cosmossdk.io/core/address"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -17,15 +15,13 @@ var _ types.MsgServer = msgServer{}
 
 type msgServer struct {
 	stakingtypes.MsgServer
-	pubkeyKeeper          types.PubkeyKeeper
-	validatorAddressCodec addresscodec.Codec
+	*Keeper
 }
 
-func NewMsgServerImpl(sdkMsgServer stakingtypes.MsgServer, pubKeyKeeper types.PubkeyKeeper, valAddrCdc addresscodec.Codec) types.MsgServer {
+func NewMsgServerImpl(sdkMsgServer stakingtypes.MsgServer, keeper *Keeper) types.MsgServer {
 	ms := &msgServer{
-		MsgServer:             sdkMsgServer,
-		pubkeyKeeper:          pubKeyKeeper,
-		validatorAddressCodec: valAddrCdc,
+		MsgServer: sdkMsgServer,
+		Keeper:    keeper,
 	}
 	return ms
 }
@@ -43,7 +39,7 @@ func (m msgServer) CreateSEDAValidator(ctx context.Context, msg *types.MsgCreate
 	}
 
 	// Validate and store the public keys.
-	activated, err := m.pubkeyKeeper.IsProvingSchemeActivated(ctx, utils.SEDAKeyIndexSecp256k1)
+	activated, err := m.pubKeyKeeper.IsProvingSchemeActivated(ctx, utils.SEDAKeyIndexSecp256k1)
 	if err != nil {
 		return nil, err
 	}
@@ -52,7 +48,7 @@ func (m msgServer) CreateSEDAValidator(ctx context.Context, msg *types.MsgCreate
 		if err != nil {
 			return nil, sdkerrors.ErrInvalidRequest.Wrapf("invalid SEDA keys: %s", err)
 		}
-		err = m.pubkeyKeeper.StoreIndexedPubKeys(sdkCtx, valAddr, msg.IndexedPubKeys)
+		err = m.pubKeyKeeper.StoreIndexedPubKeys(sdkCtx, valAddr, msg.IndexedPubKeys)
 		if err != nil {
 			return nil, err
 		}

x/staking/module.go

@@ -98,7 +98,7 @@ type AppModule struct {
 	keeper        *keeper.Keeper
 	accountKeeper sdktypes.AccountKeeper
 	bankKeeper    sdktypes.BankKeeper
-	pubKeyKeeper  types.PubkeyKeeper
+	pubKeyKeeper  types.PubKeyKeeper
 }
 
 // NewAppModule creates a new AppModule object
@@ -107,7 +107,7 @@ func NewAppModule(
 	keeper *keeper.Keeper,
 	ak sdktypes.AccountKeeper,
 	bk sdktypes.BankKeeper,
-	pk types.PubkeyKeeper,
+	pk types.PubKeyKeeper,
 ) AppModule {
 	return AppModule{
 		AppModuleBasic: AppModuleBasic{cdc: cdc},
@@ -121,7 +121,7 @@ func NewAppModule(
 // RegisterServices registers module services.
 func (am AppModule) RegisterServices(cfg module.Configurator) {
 	sdkMsgServer := sdkkeeper.NewMsgServerImpl(am.keeper.Keeper)
-	msgServer := keeper.NewMsgServerImpl(sdkMsgServer, am.pubKeyKeeper, am.keeper.ValidatorAddressCodec())
+	msgServer := keeper.NewMsgServerImpl(sdkMsgServer, am.keeper)
 	types.RegisterMsgServer(cfg.MsgServer(), msgServer)
 
 	querier := sdkkeeper.Querier{Keeper: am.keeper.Keeper}
@@ -136,7 +136,7 @@ func (am AppModule) IsAppModule() {}
 
 // RegisterInvariants registers the staking module invariants.
 func (am AppModule) RegisterInvariants(ir sdk.InvariantRegistry) {
-	sdkkeeper.RegisterInvariants(ir, am.keeper.Keeper)
+	keeper.RegisterInvariants(ir, am.keeper)
 }
 
 // InitGenesis performs genesis initialization for the staking module.

x/staking/types/expected_keepers.go

@@ -9,7 +9,8 @@ import (
 	"github.com/sedaprotocol/seda-chain/x/pubkey/types"
 )
 
-type PubkeyKeeper interface {
+type PubKeyKeeper interface {
 	StoreIndexedPubKeys(ctx sdk.Context, valAddr sdk.ValAddress, pubKeys []types.IndexedPubKey) error
 	IsProvingSchemeActivated(ctx context.Context, index utils.SEDAKeyIndex) (bool, error)
+	GetValidatorKeyAtIndex(ctx context.Context, validatorAddr sdk.ValAddress, index utils.SEDAKeyIndex) ([]byte, error)
 }

x/tally/keeper/integration_test.go

@@ -137,8 +137,9 @@ func initFixture(tb testing.TB) *fixture {
 		log.NewNopLogger(),
 	)
 
-	sdkstakingKeeper := sdkstakingkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(params.Bech32PrefixValAddr), addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr))
-	stakingKeeper := stakingkeeper.NewKeeper(sdkstakingKeeper)
+	var pubKeyKeeper *pubkeykeeper.Keeper
+	sdkStakingKeeper := sdkstakingkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(params.Bech32PrefixValAddr), addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr))
+	stakingKeeper := stakingkeeper.NewKeeper(sdkStakingKeeper, pubKeyKeeper, addresscodec.NewBech32Codec(params.Bech32PrefixValAddr))
 
 	stakingParams := sdkstakingtypes.DefaultParams()
 	stakingParams.BondDenom = bondDenom
@@ -191,7 +192,7 @@ func initFixture(tb testing.TB) *fixture {
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
-	pubKeyKeeper := pubkeykeeper.NewKeeper(
+	pubKeyKeeper = pubkeykeeper.NewKeeper(
 		cdc,
 		runtime.NewKVStoreService(keys[pubkeytypes.StoreKey]),
 		stakingKeeper,

x/vesting/keeper/integration_test.go

@@ -41,6 +41,7 @@ import (
 	pubkeytypes "github.com/sedaprotocol/seda-chain/x/pubkey/types"
 	"github.com/sedaprotocol/seda-chain/x/staking"
 	stakingkeeper "github.com/sedaprotocol/seda-chain/x/staking/keeper"
+	stakingtypes "github.com/sedaprotocol/seda-chain/x/staking/types"
 	"github.com/sedaprotocol/seda-chain/x/vesting"
 	"github.com/sedaprotocol/seda-chain/x/vesting/keeper"
 	"github.com/sedaprotocol/seda-chain/x/vesting/types"
@@ -119,8 +120,9 @@ func initFixture(tb testing.TB) *fixture {
 		log.NewNopLogger(),
 	)
 
-	sdkstakingKeeper := sdkstakingkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(params.Bech32PrefixValAddr), addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr))
-	stakingKeeper := stakingkeeper.NewKeeper(sdkstakingKeeper)
+	var pubKeyKeeper *pubkeykeeper.Keeper
+	sdkStakingKeeper := sdkstakingkeeper.NewKeeper(cdc, runtime.NewKVStoreService(keys[sdkstakingtypes.StoreKey]), accountKeeper, bankKeeper, authority.String(), addresscodec.NewBech32Codec(params.Bech32PrefixValAddr), addresscodec.NewBech32Codec(params.Bech32PrefixConsAddr))
+	stakingKeeper := stakingkeeper.NewKeeper(sdkStakingKeeper, pubKeyKeeper, addresscodec.NewBech32Codec(params.Bech32PrefixValAddr))
 
 	stakingParams := sdkstakingtypes.DefaultParams()
 	stakingParams.BondDenom = bondDenom
@@ -135,7 +137,7 @@ func initFixture(tb testing.TB) *fixture {
 		authtypes.NewModuleAddress(govtypes.ModuleName).String(),
 	)
 
-	pubKeyKeeper := pubkeykeeper.NewKeeper(
+	pubKeyKeeper = pubkeykeeper.NewKeeper(
 		cdc,
 		runtime.NewKVStoreService(keys[pubkeytypes.StoreKey]),
 		stakingKeeper,
@@ -157,7 +159,9 @@ func initFixture(tb testing.TB) *fixture {
 	})
 
 	types.RegisterMsgServer(integrationApp.MsgServiceRouter(), keeper.NewMsgServerImpl(accountKeeper, bankKeeper, stakingKeeper))
-	sdkstakingtypes.RegisterMsgServer(integrationApp.MsgServiceRouter(), sdkstakingkeeper.NewMsgServerImpl(sdkstakingKeeper))
+	sdkStakingMsgServer := sdkstakingkeeper.NewMsgServerImpl(sdkStakingKeeper)
+	stakingMsgServer := stakingkeeper.NewMsgServerImpl(sdkStakingMsgServer, stakingKeeper)
+	stakingtypes.RegisterMsgServer(integrationApp.MsgServiceRouter(), stakingMsgServer)
 
 	return &fixture{
 		IntegationApp: integrationApp,