feat: epoching to replace staking msg server (#457)

PRs - #435 - #436 --------- Co-authored-by: Runchao Han <[email protected]> Co-authored-by: Mauro Lacy <[email protected]> Co-authored-by: Mauro Lacy <[email protected]> Co-authored-by: Gurjot Singh <[email protected]> Co-authored-by: hiepmai-babylonchain <[email protected]> Co-authored-by: lesterli <[email protected]> Co-authored-by: Cirrus Gai <[email protected]> Co-authored-by: Runchao Han <[email protected]> Co-authored-by: Mauro Lacy <[email protected]> Co-authored-by: Filippos Malandrakis <[email protected]> Co-authored-by: Lazar <[email protected]> Co-authored-by: KonradStaniec <[email protected]>
babylonlabs-io · Feb 3, 2025 · 68a6d44 · 68a6d44
1 parent 83af4ad
commit 68a6d44
Show file tree

Hide file tree

Showing 48 changed files with 3,081 additions and 921 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -47,6 +47,8 @@ check of rewards
 
 - [#402](https://github.com/babylonlabs-io/babylon/pull/402) **Babylon multi-staking support**.
 This PR contains a series of PRs on multi-staking support and BTC stakingintegration.
+- [#457](https://github.com/babylonlabs-io/babylon/pull/457) Remove staking msg server and update gentx to generate
+`MsgWrappedCreateValidator`
 
 ### Bug fixes
 

diff --git a/app/ante/ante.go b/app/ante/ante.go
@@ -64,7 +64,6 @@ func NewAnteHandler(
 
 	anteHandler := sdk.ChainAnteDecorators(
 		NewWrappedAnteHandler(authAnteHandler),
-		epochingkeeper.NewDropValidatorMsgDecorator(epochingKeeper),
 		NewBtcValidationDecorator(btcConfig, btccKeeper),
 	)
 

diff --git a/app/app.go b/app/app.go
@@ -103,6 +103,7 @@ import (
 	"github.com/cosmos/cosmos-sdk/x/slashing"
 	slashingtypes "github.com/cosmos/cosmos-sdk/x/slashing/types"
 	"github.com/cosmos/cosmos-sdk/x/staking"
+	stakingkeeper "github.com/cosmos/cosmos-sdk/x/staking/keeper"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/cosmos/gogoproto/proto"
 	"github.com/cosmos/ibc-go/modules/capability"
@@ -329,7 +330,7 @@ func NewBabylonApp(
 	app.BasicModuleManager = module.NewBasicManagerFromManager(
 		app.ModuleManager,
 		map[string]module.AppModuleBasic{
-			genutiltypes.ModuleName: genutil.NewAppModuleBasic(genutiltypes.DefaultMessageValidator),
+			genutiltypes.ModuleName: genutil.NewAppModuleBasic(checkpointingtypes.GenTxMessageValidatorWrappedCreateValidator),
 			govtypes.ModuleName: gov.NewAppModuleBasic(
 				[]govclient.ProposalHandler{
 					paramsclient.ProposalHandler,
@@ -454,10 +455,7 @@ func NewBabylonApp(
 
 	app.ModuleManager.RegisterInvariants(app.CrisisKeeper)
 	app.configurator = module.NewConfigurator(app.appCodec, app.MsgServiceRouter(), app.GRPCQueryRouter())
-	if err := app.ModuleManager.RegisterServices(app.configurator); err != nil {
-		panic(err)
-	}
-
+	app.RegisterServicesWithoutStaking()
 	autocliv1.RegisterQueryServer(app.GRPCQueryRouter(), runtimeservices.NewAutoCLIQueryService(app.ModuleManager.Modules))
 
 	reflectionSvc, err := runtimeservices.NewReflectionService()
@@ -588,6 +586,47 @@ func NewBabylonApp(
 	return app
 }
 
+// RegisterServicesWithoutStaking calls the module manager
+// registration services without the staking module.
+func (app *BabylonApp) RegisterServicesWithoutStaking() {
+	// removes the staking module from the register services
+	stkModTemp := app.ModuleManager.Modules[stakingtypes.ModuleName]
+	delete(app.ModuleManager.Modules, stakingtypes.ModuleName)
+
+	if err := app.ModuleManager.RegisterServices(app.configurator); err != nil {
+		panic(err)
+	}
+
+	app.RegisterStakingQueryAndMigrations()
+
+	// adds the staking module back it back
+	app.ModuleManager.Modules[stakingtypes.ModuleName] = stkModTemp
+}
+
+// RegisterStakingQueryAndMigrations registrates in the configurator
+// the x/staking query server and its migrations
+func (app *BabylonApp) RegisterStakingQueryAndMigrations() {
+	cfg, stkK := app.configurator, app.StakingKeeper
+	stkq := stakingkeeper.NewQuerier(stkK)
+
+	stakingtypes.RegisterQueryServer(cfg.QueryServer(), stkq)
+
+	ls := app.GetSubspace(stakingtypes.ModuleName)
+	m := stakingkeeper.NewMigrator(stkK, ls)
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 1, m.Migrate1to2); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/%s from version 1 to 2: %v", stakingtypes.ModuleName, err))
+	}
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 2, m.Migrate2to3); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/%s from version 2 to 3: %v", stakingtypes.ModuleName, err))
+	}
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 3, m.Migrate3to4); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/%s from version 3 to 4: %v", stakingtypes.ModuleName, err))
+	}
+	if err := cfg.RegisterMigration(stakingtypes.ModuleName, 4, m.Migrate4to5); err != nil {
+		panic(fmt.Sprintf("failed to migrate x/%s from version 4 to 5: %v", stakingtypes.ModuleName, err))
+	}
+}
+
 // GetBaseApp returns the BaseApp of BabylonApp
 // required by ibctesting
 func (app *BabylonApp) GetBaseApp() *baseapp.BaseApp {

diff --git a/app/app_test.go b/app/app_test.go
@@ -10,6 +10,7 @@ import (
 	dbm "github.com/cosmos/cosmos-db"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/module"
+	stktypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 	"github.com/stretchr/testify/require"
 )
 
@@ -94,3 +95,32 @@ func TestUpgradeStateOnGenesis(t *testing.T) {
 		}
 	}
 }
+
+func TestStakingRouterDisabled(t *testing.T) {
+	db := dbm.NewMemDB()
+	signer, _ := signer.SetupTestPrivSigner()
+	logger := log.NewTestLogger(t)
+
+	app := NewBabylonAppWithCustomOptions(t, false, signer, SetupOptions{
+		Logger:             logger,
+		DB:                 db,
+		InvCheckPeriod:     0,
+		SkipUpgradeHeights: map[int64]bool{},
+		AppOpts:            TmpAppOptions(),
+	})
+
+	msgs := []sdk.Msg{
+		&stktypes.MsgCreateValidator{},
+		&stktypes.MsgBeginRedelegate{},
+		&stktypes.MsgCancelUnbondingDelegation{},
+		&stktypes.MsgDelegate{},
+		&stktypes.MsgEditValidator{},
+		&stktypes.MsgUndelegate{},
+		&stktypes.MsgUpdateParams{},
+	}
+
+	for _, msg := range msgs {
+		msgHandler := app.MsgServiceRouter().HandlerByTypeURL(sdk.MsgTypeURL(msg))
+		require.Nil(t, msgHandler)
+	}
+}
diff --git a/app/keepers/keepers.go b/app/keepers/keepers.go
@@ -259,6 +259,7 @@ func (ak *AppKeepers) InitKeepers(
 		runtime.NewKVStoreService(keys[epochingtypes.StoreKey]),
 		bankKeeper,
 		stakingKeeper,
+		stakingkeeper.NewMsgServerImpl(stakingKeeper),
 		appparams.AccGov.String(),
 	)
 
@@ -411,7 +412,7 @@ func (ak *AppKeepers) InitKeepers(
 		appparams.AccGov.String(),
 	)
 
-	wasmOpts = append(owasm.RegisterCustomPlugins(&ak.EpochingKeeper, &ak.CheckpointingKeeper, &ak.BTCLightClientKeeper, &ak.ZoneConciergeKeeper), wasmOpts...)
+	wasmOpts = append(owasm.RegisterCustomPlugins(&epochingKeeper, &ak.CheckpointingKeeper, &ak.BTCLightClientKeeper, &ak.ZoneConciergeKeeper), wasmOpts...)
 	wasmOpts = append(owasm.RegisterGrpcQueries(*bApp.GRPCQueryRouter(), appCodec), wasmOpts...)
 
 	ak.WasmKeeper = wasmkeeper.NewKeeper(
@@ -572,8 +573,6 @@ func (ak *AppKeepers) InitKeepers(
 		appparams.AccGov.String(),
 	)
 
-	// add msgServiceRouter so that the epoching module can forward unwrapped messages to the staking module
-	epochingKeeper.SetMsgServiceRouter(bApp.MsgServiceRouter())
 	// make ZoneConcierge and Monitor to subscribe to the epoching's hooks
 	epochingKeeper.SetHooks(
 		epochingtypes.NewMultiEpochingHooks(zcKeeper.Hooks(), monitorKeeper.Hooks()),

diff --git a/app/test_helpers.go b/app/test_helpers.go
@@ -169,7 +169,7 @@ func genesisStateWithValSet(t *testing.T,
 			Description:       stakingtypes.Description{},
 			UnbondingHeight:   int64(0),
 			UnbondingTime:     time.Unix(0, 0).UTC(),
-			Commission:        stakingtypes.NewCommission(math.LegacyZeroDec(), math.LegacyZeroDec(), math.LegacyZeroDec()),
+			Commission:        stakingtypes.NewCommission(math.LegacyZeroDec(), math.LegacyNewDec(100), math.LegacyNewDec(2)),
 			MinSelfDelegation: math.ZeroInt(),
 		}
 

diff --git a/client/client/tx.go b/client/client/tx.go
@@ -5,12 +5,11 @@ import (
 	"fmt"
 	"sync"
 
-	"github.com/babylonlabs-io/babylon/client/babylonclient"
-
 	signingv1beta1 "cosmossdk.io/api/cosmos/tx/signing/v1beta1"
 	"cosmossdk.io/errors"
 	txsigning "cosmossdk.io/x/tx/signing"
 	"github.com/avast/retry-go/v4"
+	"github.com/babylonlabs-io/babylon/client/babylonclient"
 	btcctypes "github.com/babylonlabs-io/babylon/x/btccheckpoint/types"
 	btclctypes "github.com/babylonlabs-io/babylon/x/btclightclient/types"
 	abci "github.com/cometbft/cometbft/abci/types"

diff --git a/cmd/babylond/cmd/create_bls_key.go → cmd/babylond/cmd/bls/create_bls_key.go b/cmd/babylond/cmd/create_bls_key.go → cmd/babylond/cmd/bls/create_bls_key.go
@@ -1,4 +1,4 @@
-package cmd
+package bls
 
 import (
 	"errors"
@@ -47,7 +47,8 @@ $ babylond create-bls-key %s1f5tnl46mk4dfp4nx3n2vnrvyw2h2ydz6ykhk3r --home ./
 				return err
 			}
 
-			return CreateBlsKey(homeDir, addr)
+			_, err = CreateBlsKey(homeDir, addr)
+			return err
 		},
 	}
 
@@ -56,20 +57,20 @@ $ babylond create-bls-key %s1f5tnl46mk4dfp4nx3n2vnrvyw2h2ydz6ykhk3r --home ./
 	return cmd
 }
 
-func CreateBlsKey(home string, addr sdk.AccAddress) error {
+func CreateBlsKey(home string, addr sdk.AccAddress) (*privval.WrappedFilePV, error) {
 	nodeCfg := cmtconfig.DefaultConfig()
 	keyPath := filepath.Join(home, nodeCfg.PrivValidatorKeyFile())
 	statePath := filepath.Join(home, nodeCfg.PrivValidatorStateFile())
 
 	pv, err := LoadWrappedFilePV(keyPath, statePath)
 	if err != nil {
-		return err
+		return nil, err
 	}
 
 	wrappedPV := privval.NewWrappedFilePV(pv.GetValPrivKey(), bls12381.GenPrivKey(), keyPath, statePath)
 	wrappedPV.SetAccAddress(addr)
 
-	return nil
+	return wrappedPV, nil
 }
 
 // LoadWrappedFilePV loads the wrapped file private key from the file path.

diff --git a/cmd/babylond/cmd/genhelpers/bls_add.go b/cmd/babylond/cmd/genhelpers/bls_add.go