feat: modify logic in migrate to verify before saving keys to file.

cmd/babylond/cmd/create_bls_key.go

@@ -7,7 +7,6 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/babylonlabs-io/babylon/app"
-	"github.com/babylonlabs-io/babylon/crypto/bls12381"
 	"github.com/babylonlabs-io/babylon/privval"
 )
 
@@ -29,7 +28,7 @@ $ babylond create-bls-key --home ./
 		RunE: func(cmd *cobra.Command, args []string) error {
 			homeDir, _ := cmd.Flags().GetString(flags.FlagHome)
 			password, _ := cmd.Flags().GetString(flagBlsPassword)
-			createBlsKey(bls12381.GenPrivKey(), homeDir, password)
+			createBlsKeyAndSave(homeDir, password)
 			return nil
 		},
 	}
@@ -39,11 +38,10 @@ $ babylond create-bls-key --home ./
 	return cmd
 }
 
-func createBlsKey(privKey bls12381.PrivateKey, homeDir, password string) *privval.BlsPV {
+// createBlsKeyAndSave creates a pair of BLS keys and saves them to files
+func createBlsKeyAndSave(homeDir, password string) {
 	if password == "" {
 		password = privval.NewBlsPassword()
 	}
-	pv := privval.NewBlsPV(privKey, privval.DefaultBlsKeyFile(homeDir), privval.DefaultBlsPasswordFile(homeDir))
-	pv.Key.Save(password)
-	return pv
+	privval.GenBlsPV(privval.DefaultBlsKeyFile(homeDir), privval.DefaultBlsPasswordFile(homeDir), password)
 }

cmd/babylond/cmd/init.go

@@ -1,7 +1,6 @@
 package cmd
 
 import (
-	"github.com/babylonlabs-io/babylon/crypto/bls12381"
 	"github.com/spf13/cobra"
 
 	"github.com/cosmos/cosmos-sdk/client/flags"
@@ -23,7 +22,7 @@ func InitCmd(mbm module.BasicManager, defaultNodeHome string) *cobra.Command {
 
 			homeDir, _ := cmd.Flags().GetString(flags.FlagHome)
 			password, _ := cmd.Flags().GetString(flagBlsPassword)
-			createBlsKey(bls12381.GenPrivKey(), homeDir, password)
+			createBlsKeyAndSave(homeDir, password)
 			return nil
 		},
 	}

cmd/babylond/cmd/migrate.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/babylonlabs-io/babylon/app"
 	"github.com/babylonlabs-io/babylon/crypto/bls12381"
+	"github.com/babylonlabs-io/babylon/privval"
 	cmtcfg "github.com/cometbft/cometbft/config"
 	cmtcrypto "github.com/cometbft/cometbft/crypto"
 	cmtjson "github.com/cometbft/cometbft/libs/json"
@@ -77,10 +78,14 @@ func migrate(homeDir, password string) error {
 		return fmt.Errorf("priv_validator_key.json of previous version does not contain both the comet and bls keys")
 	}
 
+	if password == "" {
+		password = privval.NewBlsPassword()
+	}
+
 	cmtPv := cmtprivval.NewFilePV(prevCmtPrivKey, cmtcfg.PrivValidatorKeyFile(), cmtcfg.PrivValidatorStateFile())
-	cmtPv.Key.Save()
-	blsPv := createBlsKey(prevBlsPrivKey, homeDir, password)
+	blsPv := privval.NewBlsPV(prevBlsPrivKey, privval.DefaultBlsKeyFile(homeDir), privval.DefaultBlsPasswordFile(homeDir))
 
+	// before saving keys to files, verify that the migrated keys match
 	if err := verifyAfterMigration(
 		prevCmtPrivKey,
 		cmtPv.Key.PrivKey,
@@ -89,6 +94,10 @@ func migrate(homeDir, password string) error {
 	); err != nil {
 		return fmt.Errorf("failed to verify after migration: %w", err)
 	}
+
+	// save key to files after verification
+	cmtPv.Save()
+	blsPv.Key.Save(password)
 	return nil
 }
 

cmd/babylond/cmd/migrate_test.go

@@ -6,8 +6,10 @@ 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"
 )
 
@@ -61,6 +63,10 @@ func TestMigrate(t *testing.T) {
 		err := os.MkdirAll(configDir, 0755)
 		require.NoError(t, err)
 
+		dataDir := filepath.Join(tempDir, "data")
+		err = os.MkdirAll(dataDir, 0755)
+		require.NoError(t, err)
+
 		pvKeyFile := filepath.Join(configDir, "priv_validator_key.json")
 		pvKey := PrevWrappedFilePV{
 			PrivKey:    ed25519.GenPrivKey(),
@@ -78,9 +84,25 @@ func TestMigrate(t *testing.T) {
 
 		// Check if new files are created
 		newPvKeyFile := filepath.Join(configDir, "priv_validator_key.json")
+		newPvStateFile := filepath.Join(dataDir, "priv_validator_state.json")
 		newBlsKeyFile := filepath.Join(configDir, "bls_key.json")
+		newBlsPasswordFile := filepath.Join(configDir, "bls_password.txt")
 		require.FileExists(t, newPvKeyFile)
+		require.FileExists(t, newPvStateFile)
 		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,
+			)
+			require.NoError(t, err)
+		})
 	})
 }