feat: verify migration after saving files and etc

cmd/babylond/cmd/genhelpers/bls_create.go

@@ -58,7 +58,7 @@ $ babylond gen-helpers create-bls %s1f5tnl46mk4dfp4nx3n2vnrvyw2h2ydz6ykhk3r --ho
 				}
 				return nil
 			}(cmtPvKeyFile, cmtPvStateFile, blsKeyFile, blsPasswordFile); err != nil {
-				return err
+				return fmt.Errorf("failed to check files: %w", err)
 			}
 
 			cmtPV := cmtprivval.LoadFilePV(cmtPvKeyFile, cmtPvStateFile)

cmd/babylond/cmd/init.go

@@ -1,23 +1,29 @@
 package cmd
 
 import (
+	"fmt"
+
 	"github.com/spf13/cobra"
 
 	"github.com/cosmos/cosmos-sdk/client/flags"
 	"github.com/cosmos/cosmos-sdk/types/module"
 	genutil "github.com/cosmos/cosmos-sdk/x/genutil/client/cli"
 )
 
+// InitCmd returns the command to initialize the config.
+// It runs InitCmd of cosmos-sdk first, then runs createBlsKeyAndSave.
 func InitCmd(mbm module.BasicManager, defaultNodeHome string) *cobra.Command {
 	cosmosInitCmd := genutil.InitCmd(mbm, defaultNodeHome)
 	cmd := &cobra.Command{
 		Use:   cosmosInitCmd.Use,
 		Short: cosmosInitCmd.Short,
-		Long:  cosmosInitCmd.Long,
-		Args:  cosmosInitCmd.Args,
+		Long: `Initializes the configuration files for the validator and node.
+		 This command also asks for a password to 
+		 generate the BLS key and encrypt it into an erc2335 structure.`,
+		Args: cosmosInitCmd.Args,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			if err := cosmosInitCmd.RunE(cmd, args); err != nil {
-				return err
+				return fmt.Errorf("failed to run init command: %w", err)
 			}
 
 			homeDir, _ := cmd.Flags().GetString(flags.FlagHome)
@@ -27,6 +33,6 @@ func InitCmd(mbm module.BasicManager, defaultNodeHome string) *cobra.Command {
 		},
 	}
 	cmd.Flags().AddFlagSet(cosmosInitCmd.Flags())
-	cmd.Flags().String(flagBlsPassword, "", "The password for the BLS key. If a flag is set, the non-empty password should be provided. If a flag is not set, the password will be read from the prompt.")
+	cmd.Flags().String(flagBlsPassword, "", "The password for the BLS key. If the flag is not set, the password will be read from the prompt.")
 	return cmd
 }

cmd/babylond/cmd/migrate.go → cmd/babylond/cmd/migrate_bls_key.go

@@ -24,14 +24,15 @@ type PrevWrappedFilePV struct {
 	BlsPrivKey bls12381.PrivateKey `json:"bls_priv_key"`
 }
 
+// MigrateBlsKeyCmd returns a command to migrate the bls keys
 func MigrateBlsKeyCmd() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "migrate-bls-key",
 		Short: "Migrate the contents of the priv_validator_key.json file into separate files of bls and comet",
-		Long: strings.TrimSpace(`migrate splits the contents of the priv_validator_key.json file, 
+		Long: strings.TrimSpace(`Migration splits the contents of the priv_validator_key.json file, 
 		which contained both the bls and comet keys used in previous versions, into separate files.
 
-BLS keys are stored along with other validator keys in priv_validator_key.json in previous version,
+BLS keys are stored along with the Ed25519 validator key in priv_validator_key.json in the previous version,
 which should exist before running the command (via babylond init or babylond testnet).
 
 NOTE: Before proceeding with the migration, ensure you back up the priv_validator_key.json file to a secure location.
@@ -55,20 +56,21 @@ $ babylond migrate-bls-key --home ./
 }
 
 // migrate splits the contents of the priv_validator_key.json file,
-// which contained both the bls and comet keys used in previous versions, into separate files
+// which contained both the bls and comet keys used in previous versions, into separate files.
+// After saving keys to separate files, it verifies if the migrated keys match
 func migrate(homeDir, password string) error {
 	cmtcfg := cmtcfg.DefaultConfig()
 	cmtcfg.SetRoot(homeDir)
 
 	filepath := cmtcfg.PrivValidatorKeyFile()
 
 	if !cmtos.FileExists(filepath) {
-		return fmt.Errorf("priv_validator_key.json of previous version not found")
+		return fmt.Errorf("priv_validator_key.json of previous version not found in %s", filepath)
 	}
 
 	pv, err := loadPrevWrappedFilePV(filepath)
 	if err != nil {
-		return err
+		return fmt.Errorf("failed to load previous version of priv_validator_key.json in %s", filepath)
 	}
 
 	prevCmtPrivKey := pv.PrivKey
@@ -82,22 +84,25 @@ func migrate(homeDir, password string) error {
 		password = privval.NewBlsPassword()
 	}
 
-	cmtPv := cmtprivval.NewFilePV(prevCmtPrivKey, cmtcfg.PrivValidatorKeyFile(), cmtcfg.PrivValidatorStateFile())
-	blsPv := privval.NewBlsPV(prevBlsPrivKey, privval.DefaultBlsKeyFile(homeDir), privval.DefaultBlsPasswordFile(homeDir))
+	cmtKeyFilePath := cmtcfg.PrivValidatorKeyFile()
+	cmtStateFilePath := cmtcfg.PrivValidatorStateFile()
+	blsKeyFilePath := privval.DefaultBlsKeyFile(homeDir)
+	blsPasswordFilePath := privval.DefaultBlsPasswordFile(homeDir)
 
-	// before saving keys to files, verify that the migrated keys match
-	if err := verifyAfterMigration(
-		prevCmtPrivKey,
-		cmtPv.Key.PrivKey,
-		prevBlsPrivKey,
-		blsPv.Key.PrivKey,
-	); err != nil {
-		return fmt.Errorf("failed to verify after migration: %w", err)
-	}
+	cmtPv := cmtprivval.NewFilePV(prevCmtPrivKey, cmtKeyFilePath, cmtStateFilePath)
+	blsPv := privval.NewBlsPV(prevBlsPrivKey, blsKeyFilePath, blsPasswordFilePath)
 
 	// save key to files after verification
 	cmtPv.Save()
 	blsPv.Key.Save(password)
+
+	if err := verifySeparateFiles(
+		cmtKeyFilePath, cmtStateFilePath, blsKeyFilePath, blsPasswordFilePath,
+		prevCmtPrivKey, prevBlsPrivKey,
+	); err != nil {
+		return fmt.Errorf("failed to verify separate files: %w", err)
+	}
+
 	return nil
 }
 
@@ -115,9 +120,17 @@ func loadPrevWrappedFilePV(filePath string) (*PrevWrappedFilePV, error) {
 	return &pvKey, nil
 }
 
-// verifyAfterMigration checks if the migrated keys match
-func verifyAfterMigration(prevCmtPrivKey, newCmtPrivKey cmtcrypto.PrivKey, prevBlsPrivKey, newBlsPrivKey bls12381.PrivateKey) error {
-	if bytes.Equal(prevCmtPrivKey.Bytes(), newCmtPrivKey.Bytes()) && bytes.Equal(prevBlsPrivKey, newBlsPrivKey) {
+// verifySeparateFiles checks if the migrated keys match
+// after saving keys to separate files
+func verifySeparateFiles(
+	cmtKeyFilePath, cmtStateFilePath, blsKeyFilePath, blsPasswordFilePath string,
+	prevCmtPrivKey cmtcrypto.PrivKey,
+	prevBlsPrivKey bls12381.PrivateKey,
+) error {
+	cmtPv := cmtprivval.LoadFilePV(cmtKeyFilePath, cmtStateFilePath)
+	blsPv := privval.LoadBlsPV(blsKeyFilePath, blsPasswordFilePath)
+
+	if bytes.Equal(prevCmtPrivKey.Bytes(), cmtPv.Key.PrivKey.Bytes()) && bytes.Equal(prevBlsPrivKey, blsPv.Key.PrivKey) {
 		return nil
 	}
 	return fmt.Errorf("migrated keys do not match")

cmd/babylond/cmd/migrate_test.go → cmd/babylond/cmd/migrate_bls_key_test.go

@@ -6,10 +6,8 @@ import (
 	"testing"
 
 	"github.com/babylonlabs-io/babylon/crypto/bls12381"
-	"github.com/babylonlabs-io/babylon/privval"
 	"github.com/cometbft/cometbft/crypto/ed25519"
 	cmtjson "github.com/cometbft/cometbft/libs/json"
-	cmtprivval "github.com/cometbft/cometbft/privval"
 	"github.com/stretchr/testify/require"
 )
 
@@ -35,7 +33,6 @@ func TestMigrate(t *testing.T) {
 
 		err = migrate(tempDir, "")
 		require.Error(t, err)
-		require.Contains(t, err.Error(), "Error reading PrivValidator key")
 	})
 
 	t.Run("missing keys", func(t *testing.T) {
@@ -92,16 +89,11 @@ func TestMigrate(t *testing.T) {
 		require.FileExists(t, newBlsKeyFile)
 		require.FileExists(t, newBlsPasswordFile)
 
-		t.Run("verify after migration", func(t *testing.T) {
-			newCmtPv := cmtprivval.LoadFilePV(newPvKeyFile, newPvStateFile)
-			newBlsPv := privval.LoadBlsPV(newBlsKeyFile, newBlsPasswordFile)
-			err := verifyAfterMigration(
-				pvKey.PrivKey,
-				newCmtPv.Key.PrivKey,
-				pvKey.BlsPrivKey,
-				newBlsPv.Key.PrivKey,
+		t.Run("verify separated files", func(t *testing.T) {
+			verifySeparateFiles(
+				newPvKeyFile, newPvStateFile, newBlsKeyFile, newBlsPasswordFile,
+				pvKey.PrivKey, pvKey.BlsPrivKey,
 			)
-			require.NoError(t, err)
 		})
 	})
 }
@@ -140,33 +132,3 @@ func TestLoadPrevWrappedFilePV(t *testing.T) {
 		require.NotNil(t, loadedPvKey.BlsPrivKey)
 	})
 }
-
-func TestVerifyAfterMigration(t *testing.T) {
-	t.Run("matching keys", func(t *testing.T) {
-		cmtKey := ed25519.GenPrivKey()
-		blsKey := bls12381.GenPrivKey()
-
-		err := verifyAfterMigration(cmtKey, cmtKey, blsKey, blsKey)
-		require.NoError(t, err)
-	})
-
-	t.Run("non-matching comet keys", func(t *testing.T) {
-		cmtKey1 := ed25519.GenPrivKey()
-		cmtKey2 := ed25519.GenPrivKey()
-		blsKey := bls12381.GenPrivKey()
-
-		err := verifyAfterMigration(cmtKey1, cmtKey2, blsKey, blsKey)
-		require.Error(t, err)
-		require.Contains(t, err.Error(), "migrated keys do not match")
-	})
-
-	t.Run("non-matching bls keys", func(t *testing.T) {
-		cmtKey := ed25519.GenPrivKey()
-		blsKey1 := bls12381.GenPrivKey()
-		blsKey2 := bls12381.GenPrivKey()
-
-		err := verifyAfterMigration(cmtKey, cmtKey, blsKey1, blsKey2)
-		require.Error(t, err)
-		require.Contains(t, err.Error(), "migrated keys do not match")
-	})
-}

cmd/babylond/cmd/root.go

@@ -2,6 +2,7 @@ package cmd
 
 import (
 	"errors"
+	"fmt"
 	"io"
 	"os"
 
@@ -72,12 +73,12 @@ func NewRootCmd() *cobra.Command {
 			initClientCtx = initClientCtx.WithCmdContext(cmd.Context())
 			initClientCtx, err := client.ReadPersistentCommandFlags(initClientCtx, cmd.Flags())
 			if err != nil {
-				return err
+				return fmt.Errorf("failed to read command flags: %w", err)
 			}
 
 			initClientCtx, err = config.ReadFromClientConfig(initClientCtx)
 			if err != nil {
-				return err
+				return fmt.Errorf("failed to read client config: %w", err)
 			}
 
 			if !initClientCtx.Offline {
@@ -94,14 +95,14 @@ func NewRootCmd() *cobra.Command {
 					txConfigOpts,
 				)
 				if err != nil {
-					return err
+					return fmt.Errorf("failed to create tx config: %w", err)
 				}
 
 				initClientCtx = initClientCtx.WithTxConfig(txConfig)
 			}
 
 			if err := client.SetCmdClientContextHandler(initClientCtx, cmd); err != nil {
-				return err
+				return fmt.Errorf("failed to set cmd client context handler: %w", err)
 			}
 
 			customAppTemplate, customAppConfig := initAppConfig()
@@ -110,7 +111,7 @@ func NewRootCmd() *cobra.Command {
 			err = server.InterceptConfigsPreRunHandler(cmd, customAppTemplate, customAppConfig, customCometConfig)
 
 			if err != nil {
-				return err
+				return fmt.Errorf("failed to intercept configs: %w", err)
 			}
 
 			return nil
@@ -130,7 +131,7 @@ func NewRootCmd() *cobra.Command {
 	}
 
 	if err := autoCliOpts.EnhanceRootCommand(rootCmd); err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to enhance root command: %w", err))
 	}
 
 	return rootCmd
@@ -270,7 +271,7 @@ func newApp(logger log.Logger, db dbm.DB, traceStore io.Writer, appOpts serverty
 
 	privSigner, err := signer.InitPrivSigner(homeDir)
 	if err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to initialize priv signer: %w", err))
 	}
 
 	var wasmOpts []wasmkeeper.Option
@@ -308,13 +309,13 @@ func appExport(
 
 	privSigner, err := signer.InitPrivSigner(homePath)
 	if err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to initialize priv signer: %w", err))
 	}
 	if height != -1 {
 		babylonApp = app.NewBabylonApp(logger, db, traceStore, false, map[int64]bool{}, uint(1), privSigner, appOpts, app.EmptyWasmOpts)
 
 		if err = babylonApp.LoadHeight(height); err != nil {
-			return servertypes.ExportedApp{}, err
+			return servertypes.ExportedApp{}, fmt.Errorf("failed to load height: %w", err)
 		}
 	} else {
 		babylonApp = app.NewBabylonApp(logger, db, traceStore, true, map[int64]bool{}, uint(1), privSigner, appOpts, app.EmptyWasmOpts)

privval/bls.go

@@ -113,29 +113,29 @@ func (k *BlsPVKey) Save(password string) {
 	// encrypt the bls12381 key to erc2335 type
 	erc2335BlsPvKey, err := erc2335.Encrypt(k.PrivKey, k.PubKey.Bytes(), password)
 	if err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to encrypt BLS key: %w", err))
 	}
 
 	// Parse the encrypted key back to Erc2335KeyStore structure
 	var keystore erc2335.Erc2335KeyStore
 	if err := json.Unmarshal(erc2335BlsPvKey, &keystore); err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to unmarshal BLS key: %w", err))
 	}
 
 	// convert keystore to json
 	jsonBytes, err := json.MarshalIndent(keystore, "", "  ")
 	if err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to marshal BLS key: %w", err))
 	}
 
 	// write generated erc2335 keystore to file
 	if err := tempfile.WriteFileAtomic(k.filePath, jsonBytes, 0600); err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to write BLS key: %w", err))
 	}
 
 	// save used password to file
 	if err := tempfile.WriteFileAtomic(k.passwordPath, []byte(password), 0600); err != nil {
-		panic(err)
+		panic(fmt.Errorf("failed to write BLS password: %w", err))
 	}
 }
 

privval/file.go

@@ -4,11 +4,14 @@ import (
 	"fmt"
 
 	"github.com/babylonlabs-io/babylon/crypto/bls12381"
+	"github.com/babylonlabs-io/babylon/x/checkpointing/keeper"
 	checkpointingtypes "github.com/babylonlabs-io/babylon/x/checkpointing/types"
 	cmtcrypto "github.com/cometbft/cometbft/crypto"
 	cmtprivval "github.com/cometbft/cometbft/privval"
 )
 
+var _ keeper.BlsSigner = &WrappedFilePV{}
+
 // WrappedFilePV is a wrapper around cmtprivval.FilePV
 type WrappedFilePV struct {
 	Comet cmtprivval.FilePVKey
@@ -34,7 +37,7 @@ func (pv *WrappedFilePV) SignMsgWithBls(msg []byte) (bls12381.Signature, error)
 // GetBlsPubkey returns the public key of the BLS, implementing the BlsSigner interface
 func (pv *WrappedFilePV) GetBlsPubkey() (bls12381.PublicKey, error) {
 	if pv.Bls.PrivKey == nil {
-		return nil, checkpointingtypes.ErrBlsPrivKeyDoesNotExist
+		return nil, fmt.Errorf("Error while getting BLS public key: %w", checkpointingtypes.ErrBlsPrivKeyDoesNotExist)
 	}
 	return pv.Bls.PrivKey.PubKey(), nil
 }