feat: BLS keystore improvement (#467)

Closes #336 --------- Co-authored-by: wonjoon <[email protected]> Co-authored-by: dongsam <[email protected]>
babylonlabs-io · Feb 5, 2025 · d065cdd · d065cdd
1 parent 2108dce
commit d065cdd
Show file tree

Hide file tree

Showing 60 changed files with 1,418 additions and 871 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -44,11 +44,12 @@ randomness commit at `TestBTCRewardsDistribution`.
 - [#391](https://github.com/babylonlabs-io/babylon/pull/391) Fix e2e `TestBTCRewardsDistribution` flunky
 check of rewards
 - [#458](https://github.com/babylonlabs-io/babylon/pull/458) Set `CosmosProvider` functions as public
+- [#467](https://github.com/babylonlabs-io/babylon/pull/467) BLS keystore improvement
 
 ### State Machine Breaking
 
 - [#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.
+This PR contains a series of PRs on multi-staking support and BTC staking integration.
 - [#457](https://github.com/babylonlabs-io/babylon/pull/457) Remove staking msg server and update gentx to generate
 `MsgWrappedCreateValidator`
 

diff --git a/Makefile b/Makefile
@@ -86,6 +86,11 @@ else
   BUILD_TAGS += mainnet
 endif
 
+# Handles the inclusion of e2e upgrade in binary
+ifeq (e2e_upgrade,$(findstring e2e_upgrade,$(BABYLON_BUILD_OPTIONS)))
+  BUILD_TAGS += e2e_upgrade
+endif
+
 # DB backend selection
 ifeq (cleveldb,$(findstring cleveldb,$(BABYLON_BUILD_OPTIONS)))
   ldflags += -X github.com/cosmos/cosmos-sdk/types.DBBackend=cleveldb

diff --git a/app/app.go b/app/app.go
@@ -23,36 +23,6 @@ import (
 	"github.com/CosmWasm/wasmd/x/wasm"
 	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
 	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
-	"github.com/babylonlabs-io/babylon/app/ante"
-	appkeepers "github.com/babylonlabs-io/babylon/app/keepers"
-	appparams "github.com/babylonlabs-io/babylon/app/params"
-	"github.com/babylonlabs-io/babylon/app/signer"
-	"github.com/babylonlabs-io/babylon/app/upgrades"
-	"github.com/babylonlabs-io/babylon/client/docs"
-	bbn "github.com/babylonlabs-io/babylon/types"
-	"github.com/babylonlabs-io/babylon/x/btccheckpoint"
-	btccheckpointtypes "github.com/babylonlabs-io/babylon/x/btccheckpoint/types"
-	"github.com/babylonlabs-io/babylon/x/btclightclient"
-	btclightclienttypes "github.com/babylonlabs-io/babylon/x/btclightclient/types"
-	"github.com/babylonlabs-io/babylon/x/btcstaking"
-	btcstakingtypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
-	"github.com/babylonlabs-io/babylon/x/btcstkconsumer"
-	bsctypes "github.com/babylonlabs-io/babylon/x/btcstkconsumer/types"
-	"github.com/babylonlabs-io/babylon/x/checkpointing"
-	checkpointingtypes "github.com/babylonlabs-io/babylon/x/checkpointing/types"
-	"github.com/babylonlabs-io/babylon/x/epoching"
-	epochingtypes "github.com/babylonlabs-io/babylon/x/epoching/types"
-	"github.com/babylonlabs-io/babylon/x/finality"
-	finalitytypes "github.com/babylonlabs-io/babylon/x/finality/types"
-	"github.com/babylonlabs-io/babylon/x/incentive"
-	incentivekeeper "github.com/babylonlabs-io/babylon/x/incentive/keeper"
-	incentivetypes "github.com/babylonlabs-io/babylon/x/incentive/types"
-	"github.com/babylonlabs-io/babylon/x/mint"
-	minttypes "github.com/babylonlabs-io/babylon/x/mint/types"
-	"github.com/babylonlabs-io/babylon/x/monitor"
-	monitortypes "github.com/babylonlabs-io/babylon/x/monitor/types"
-	"github.com/babylonlabs-io/babylon/x/zoneconcierge"
-	zctypes "github.com/babylonlabs-io/babylon/x/zoneconcierge/types"
 	abci "github.com/cometbft/cometbft/abci/types"
 	cmtos "github.com/cometbft/cometbft/libs/os"
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
@@ -119,6 +89,36 @@ import (
 	ibcexported "github.com/cosmos/ibc-go/v8/modules/core/exported"
 	ibctm "github.com/cosmos/ibc-go/v8/modules/light-clients/07-tendermint"
 	"github.com/spf13/cast"
+
+	"github.com/babylonlabs-io/babylon/app/ante"
+	appkeepers "github.com/babylonlabs-io/babylon/app/keepers"
+	appparams "github.com/babylonlabs-io/babylon/app/params"
+	"github.com/babylonlabs-io/babylon/app/upgrades"
+	"github.com/babylonlabs-io/babylon/client/docs"
+	bbn "github.com/babylonlabs-io/babylon/types"
+	"github.com/babylonlabs-io/babylon/x/btccheckpoint"
+	btccheckpointtypes "github.com/babylonlabs-io/babylon/x/btccheckpoint/types"
+	"github.com/babylonlabs-io/babylon/x/btclightclient"
+	btclightclienttypes "github.com/babylonlabs-io/babylon/x/btclightclient/types"
+	"github.com/babylonlabs-io/babylon/x/btcstaking"
+	btcstakingtypes "github.com/babylonlabs-io/babylon/x/btcstaking/types"
+	"github.com/babylonlabs-io/babylon/x/btcstkconsumer"
+	bsctypes "github.com/babylonlabs-io/babylon/x/btcstkconsumer/types"
+	"github.com/babylonlabs-io/babylon/x/checkpointing"
+	checkpointingtypes "github.com/babylonlabs-io/babylon/x/checkpointing/types"
+	"github.com/babylonlabs-io/babylon/x/epoching"
+	epochingtypes "github.com/babylonlabs-io/babylon/x/epoching/types"
+	"github.com/babylonlabs-io/babylon/x/finality"
+	finalitytypes "github.com/babylonlabs-io/babylon/x/finality/types"
+	"github.com/babylonlabs-io/babylon/x/incentive"
+	incentivekeeper "github.com/babylonlabs-io/babylon/x/incentive/keeper"
+	incentivetypes "github.com/babylonlabs-io/babylon/x/incentive/types"
+	"github.com/babylonlabs-io/babylon/x/mint"
+	minttypes "github.com/babylonlabs-io/babylon/x/mint/types"
+	"github.com/babylonlabs-io/babylon/x/monitor"
+	monitortypes "github.com/babylonlabs-io/babylon/x/monitor/types"
+	"github.com/babylonlabs-io/babylon/x/zoneconcierge"
+	zctypes "github.com/babylonlabs-io/babylon/x/zoneconcierge/types"
 )
 
 const (
@@ -210,7 +210,7 @@ func NewBabylonApp(
 	loadLatest bool,
 	skipUpgradeHeights map[int64]bool,
 	invCheckPeriod uint,
-	privSigner *signer.PrivSigner,
+	blsSigner *checkpointingtypes.BlsSigner,
 	appOpts servertypes.AppOptions,
 	wasmOpts []wasmkeeper.Option,
 	baseAppOptions ...func(*baseapp.BaseApp),
@@ -261,7 +261,7 @@ func NewBabylonApp(
 		homePath,
 		invCheckPeriod,
 		skipUpgradeHeights,
-		privSigner,
+		*blsSigner,
 		appOpts,
 		wasmConfig,
 		wasmOpts,

diff --git a/app/app_test.go b/app/app_test.go
@@ -5,21 +5,27 @@ import (
 	"testing"
 
 	"cosmossdk.io/log"
-	"github.com/babylonlabs-io/babylon/testutil/signer"
 	abci "github.com/cometbft/cometbft/abci/types"
 	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"
+
+	testsigner "github.com/babylonlabs-io/babylon/testutil/signer"
+	checkpointingtypes "github.com/babylonlabs-io/babylon/x/checkpointing/types"
 )
 
 func TestBabylonBlockedAddrs(t *testing.T) {
 	db := dbm.NewMemDB()
-	signer, _ := signer.SetupTestPrivSigner()
+
+	tbs, err := testsigner.SetupTestBlsSigner()
+	require.NoError(t, err)
+	blsSigner := checkpointingtypes.BlsSigner(tbs)
+
 	logger := log.NewTestLogger(t)
 
-	app := NewBabylonAppWithCustomOptions(t, false, signer, SetupOptions{
+	app := NewBabylonAppWithCustomOptions(t, false, blsSigner, SetupOptions{
 		Logger:             logger,
 		DB:                 db,
 		InvCheckPeriod:     0,
@@ -42,7 +48,7 @@ func TestBabylonBlockedAddrs(t *testing.T) {
 		)
 	}
 
-	_, err := app.FinalizeBlock(&abci.RequestFinalizeBlock{
+	_, err = app.FinalizeBlock(&abci.RequestFinalizeBlock{
 		Height: 1,
 	})
 	require.NoError(t, err)
@@ -58,7 +64,7 @@ func TestBabylonBlockedAddrs(t *testing.T) {
 		true,
 		map[int64]bool{},
 		0,
-		signer,
+		&blsSigner,
 		TmpAppOptions(),
 		EmptyWasmOpts,
 	)
@@ -73,11 +79,14 @@ func TestGetMaccPerms(t *testing.T) {
 
 func TestUpgradeStateOnGenesis(t *testing.T) {
 	db := dbm.NewMemDB()
-	privSigner, err := signer.SetupTestPrivSigner()
+
+	tbs, err := testsigner.SetupTestBlsSigner()
 	require.NoError(t, err)
+	blsSigner := checkpointingtypes.BlsSigner(tbs)
+
 	logger := log.NewTestLogger(t)
 
-	app := NewBabylonAppWithCustomOptions(t, false, privSigner, SetupOptions{
+	app := NewBabylonAppWithCustomOptions(t, false, blsSigner, SetupOptions{
 		Logger:             logger,
 		DB:                 db,
 		InvCheckPeriod:     0,
@@ -98,10 +107,10 @@ func TestUpgradeStateOnGenesis(t *testing.T) {
 
 func TestStakingRouterDisabled(t *testing.T) {
 	db := dbm.NewMemDB()
-	signer, _ := signer.SetupTestPrivSigner()
+	tbs, _ := testsigner.SetupTestBlsSigner()
 	logger := log.NewTestLogger(t)
 
-	app := NewBabylonAppWithCustomOptions(t, false, signer, SetupOptions{
+	app := NewBabylonAppWithCustomOptions(t, false, tbs, SetupOptions{
 		Logger:             logger,
 		DB:                 db,
 		InvCheckPeriod:     0,

diff --git a/app/build_tags.go b/app/build_tags.go
@@ -0,0 +1,5 @@
+package app
+
+// This variable is only set true
+// when build tag is set to "e2e_upgrade"
+var IsE2EUpgradeBuildFlag bool
diff --git a/app/encoding.go b/app/encoding.go
@@ -12,6 +12,7 @@ import (
 	appparams "github.com/babylonlabs-io/babylon/app/params"
 	"github.com/babylonlabs-io/babylon/testutil/signer"
 	bbn "github.com/babylonlabs-io/babylon/types"
+	checkpointingtypes "github.com/babylonlabs-io/babylon/x/checkpointing/types"
 )
 
 // TmpAppOptions returns an app option with tmp dir and btc network
@@ -27,16 +28,19 @@ func TmpAppOptions() simsutils.AppOptionsMap {
 	return appOpts
 }
 
+// NewTmpBabylonApp returns a new BabylonApp
 func NewTmpBabylonApp() *BabylonApp {
-	signer, _ := signer.SetupTestPrivSigner()
+	tbs, _ := signer.SetupTestBlsSigner()
+	blsSigner := checkpointingtypes.BlsSigner(tbs)
+
 	return NewBabylonApp(
 		log.NewNopLogger(),
 		dbm.NewMemDB(),
 		nil,
 		true,
 		map[int64]bool{},
 		0,
-		signer,
+		&blsSigner,
 		TmpAppOptions(),
 		[]wasmkeeper.Option{})
 }

diff --git a/app/include_upgrade_e2e.go b/app/include_upgrade_e2e.go
@@ -0,0 +1,9 @@
+//go:build e2e_upgrade
+
+package app
+
+// init is used to include v1 upgrade testnet data
+// it is also used for e2e testing
+func init() {
+	IsE2EUpgradeBuildFlag = true
+}
diff --git a/app/keepers/keepers.go b/app/keepers/keepers.go
@@ -17,9 +17,6 @@ import (
 	"github.com/CosmWasm/wasmd/x/wasm"
 	wasmkeeper "github.com/CosmWasm/wasmd/x/wasm/keeper"
 	wasmtypes "github.com/CosmWasm/wasmd/x/wasm/types"
-	mintkeeper "github.com/babylonlabs-io/babylon/x/mint/keeper"
-	minttypes "github.com/babylonlabs-io/babylon/x/mint/types"
-	"github.com/babylonlabs-io/babylon/x/zoneconcierge"
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/runtime"
@@ -62,7 +59,6 @@ import (
 	ibckeeper "github.com/cosmos/ibc-go/v8/modules/core/keeper"
 
 	appparams "github.com/babylonlabs-io/babylon/app/params"
-	"github.com/babylonlabs-io/babylon/app/signer"
 	bbn "github.com/babylonlabs-io/babylon/types"
 	owasm "github.com/babylonlabs-io/babylon/wasmbinding"
 	btccheckpointkeeper "github.com/babylonlabs-io/babylon/x/btccheckpoint/keeper"
@@ -81,8 +77,11 @@ import (
 	finalitytypes "github.com/babylonlabs-io/babylon/x/finality/types"
 	incentivekeeper "github.com/babylonlabs-io/babylon/x/incentive/keeper"
 	incentivetypes "github.com/babylonlabs-io/babylon/x/incentive/types"
+	mintkeeper "github.com/babylonlabs-io/babylon/x/mint/keeper"
+	minttypes "github.com/babylonlabs-io/babylon/x/mint/types"
 	monitorkeeper "github.com/babylonlabs-io/babylon/x/monitor/keeper"
 	monitortypes "github.com/babylonlabs-io/babylon/x/monitor/types"
+	"github.com/babylonlabs-io/babylon/x/zoneconcierge"
 	zckeeper "github.com/babylonlabs-io/babylon/x/zoneconcierge/keeper"
 	zctypes "github.com/babylonlabs-io/babylon/x/zoneconcierge/types"
 )
@@ -173,7 +172,7 @@ func (ak *AppKeepers) InitKeepers(
 	homePath string,
 	invCheckPeriod uint,
 	skipUpgradeHeights map[int64]bool,
-	privSigner *signer.PrivSigner,
+	blsSigner checkpointingtypes.BlsSigner,
 	appOpts servertypes.AppOptions,
 	wasmConfig wasmtypes.WasmConfig,
 	wasmOpts []wasmkeeper.Option,
@@ -266,7 +265,7 @@ func (ak *AppKeepers) InitKeepers(
 	checkpointingKeeper := checkpointingkeeper.NewKeeper(
 		appCodec,
 		runtime.NewKVStoreService(keys[checkpointingtypes.StoreKey]),
-		privSigner.WrappedPV,
+		blsSigner,
 		epochingKeeper,
 	)