fix: add comments, wrapping errors and remove unused functions

app/signer/private.go

@@ -1,6 +1,8 @@
 package signer
 
 import (
+	"fmt"
+
 	cmtconfig "github.com/cometbft/cometbft/config"
 
 	"github.com/babylonlabs-io/babylon/privval"
@@ -21,7 +23,7 @@ func InitPrivSigner(nodeDir string) (*PrivSigner, error) {
 	blsPasswordFile := privval.DefaultBlsPasswordFile(nodeDir)
 
 	if err := privval.EnsureDirs(pvKeyFile, pvStateFile, blsKeyFile, blsPasswordFile); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to ensure dirs: %w", err)
 	}
 
 	cometPV := cmtprivval.LoadFilePV(pvKeyFile, pvStateFile)

crypto/erc2335/erc2335.go

@@ -1,39 +1,35 @@
 package erc2335
 
 import (
-	"crypto/rand"
-	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"os"
 
-	"github.com/pkg/errors"
 	keystorev4 "github.com/wealdtech/go-eth2-wallet-encryptor-keystorev4"
 )
 
+// Erc2335KeyStore represents an ERC-2335 compatible keystore used in keystorev4.
 type Erc2335KeyStore struct {
-	Crypto      map[string]interface{} `json:"crypto"`
-	Version     uint                   `json:"version"`
-	UUID        string                 `json:"uuid"`
-	Path        string                 `json:"path"`
-	Pubkey      string                 `json:"pubkey"`
-	Description string                 `json:"description"`
+	Crypto      map[string]interface{} `json:"crypto"`      // Map containing the encryption details for the keystore such as checksum, cipher, and kdf.
+	Version     uint                   `json:"version"`     // Version of the keystore format (e.g., 4 for keystorev4).
+	UUID        string                 `json:"uuid"`        // Unique identifier for the keystore.
+	Path        string                 `json:"path"`        // File path where the keystore is stored.
+	Pubkey      string                 `json:"pubkey"`      // Public key associated with the keystore, stored as a hexadecimal string.
+	Description string                 `json:"description"` // Optional description of the keystore, currently used to store the delegator address.
 }
 
-// wonjoon: encrypt key pair to erc2335 keystore
-// available to handle all keys in []byte format
+// Encrypt encrypts the private key using the keystorev4 encryptor.
 func Encrypt(privKey, pubKey []byte, password string) ([]byte, error) {
 	if privKey == nil {
-		return nil, errors.New("private key cannot be nil")
+		return nil, fmt.Errorf("private key cannot be nil")
 	}
 
 	encryptor := keystorev4.New()
 	cryptoFields, err := encryptor.Encrypt(privKey, password)
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to encrypt private key")
+		return nil, fmt.Errorf("failed to encrypt private key: %w", err)
 	}
 
-	// Create the keystore json structure
 	keystoreJSON := Erc2335KeyStore{
 		Crypto:  cryptoFields,
 		Version: 4,
@@ -43,41 +39,24 @@ func Encrypt(privKey, pubKey []byte, password string) ([]byte, error) {
 	return json.Marshal(keystoreJSON)
 }
 
+// Decrypt decrypts the private key from the keystore using the given password.
+func Decrypt(keystore Erc2335KeyStore, password string) ([]byte, error) {
+	encryptor := keystorev4.New()
+	return encryptor.Decrypt(keystore.Crypto, password)
+}
+
+// LoadKeyStore loads a keystore from a file.
 func LoadKeyStore(filePath string) (Erc2335KeyStore, error) {
 	var keystore Erc2335KeyStore
 
 	keyJSONBytes, err := os.ReadFile(filePath)
 	if err != nil {
-		return Erc2335KeyStore{}, err
+		return Erc2335KeyStore{}, fmt.Errorf("failed to read keystore file: %w", err)
 	}
 
 	if err := json.Unmarshal(keyJSONBytes, &keystore); err != nil {
-		return Erc2335KeyStore{}, err
+		return Erc2335KeyStore{}, fmt.Errorf("failed to unmarshal keystore: %w", err)
 	}
 
 	return keystore, nil
 }
-
-// decrypt private key from erc2335 keystore
-func Decrypt(keystore Erc2335KeyStore, password string) ([]byte, error) {
-	encryptor := keystorev4.New()
-	return encryptor.Decrypt(keystore.Crypto, password)
-}
-
-func SavePasswordToFile(password, filePath string) error {
-	return os.WriteFile(filePath, []byte(password), 0600)
-}
-
-func LoadPaswordFromFile(filePath string) (string, error) {
-	password, err := os.ReadFile(filePath)
-	return string(password), err
-}
-
-func CreateRandomPassword() string {
-	password := make([]byte, 32)
-	_, err := rand.Read(password)
-	if err != nil {
-		panic(err)
-	}
-	return hex.EncodeToString(password)
-}

crypto/erc2335/erc2335_test.go

@@ -6,17 +6,18 @@ import (
 	"testing"
 
 	"github.com/babylonlabs-io/babylon/crypto/bls12381"
+	"github.com/cometbft/cometbft/libs/tempfile"
 	"github.com/test-go/testify/require"
 )
 
+const password string = "password"
+
 func TestEncryptBLS(t *testing.T) {
 	t.Run("create bls key", func(t *testing.T) {
 		blsPrivKey := bls12381.GenPrivKey()
 		blsPubKey := blsPrivKey.PubKey().Bytes()
 
 		t.Run("encrypt bls key", func(t *testing.T) {
-			password := CreateRandomPassword()
-			t.Logf("password: %s", password)
 			encryptedBlsKey, err := Encrypt(blsPrivKey, blsPubKey, password)
 			require.NoError(t, err)
 			t.Logf("encrypted bls key: %s", encryptedBlsKey)
@@ -41,18 +42,16 @@ func TestEncryptBLS(t *testing.T) {
 		})
 
 		t.Run("save password and encrypt bls key", func(t *testing.T) {
-			password := CreateRandomPassword()
-			t.Logf("password: %s", password)
-
 			encryptedBlsKey, err := Encrypt(blsPrivKey, blsPubKey, password)
 			require.NoError(t, err)
 			t.Logf("encrypted bls key: %s", encryptedBlsKey)
-			err = SavePasswordToFile(password, "password.txt")
+			err = tempfile.WriteFileAtomic("password.txt", []byte(password), 0600)
 			require.NoError(t, err)
 
 			t.Run("load password and decrypt bls key", func(t *testing.T) {
-				password, err := LoadPaswordFromFile("password.txt")
+				passwordBytes, err := os.ReadFile("password.txt")
 				require.NoError(t, err)
+				password := string(passwordBytes)
 
 				var keystore Erc2335KeyStore
 				err = json.Unmarshal(encryptedBlsKey, &keystore)
@@ -64,15 +63,15 @@ func TestEncryptBLS(t *testing.T) {
 			})
 
 			t.Run("save new password into same file", func(t *testing.T) {
-				newPassword := CreateRandomPassword()
-				t.Logf("new password: %s", newPassword)
-				err = SavePasswordToFile(newPassword, "password.txt")
+				newPassword := "new password"
+				err = tempfile.WriteFileAtomic("password.txt", []byte(newPassword), 0600)
 				require.NoError(t, err)
 			})
 
 			t.Run("failed when load different password and decrypt bls key", func(t *testing.T) {
-				password, err := LoadPaswordFromFile("password.txt")
+				passwordBytes, err := os.ReadFile("password.txt")
 				require.NoError(t, err)
+				password := string(passwordBytes)
 
 				var keystore Erc2335KeyStore
 				err = json.Unmarshal(encryptedBlsKey, &keystore)
@@ -83,7 +82,7 @@ func TestEncryptBLS(t *testing.T) {
 			})
 
 			t.Run("failed when password file don't exist", func(t *testing.T) {
-				_, err := LoadPaswordFromFile("nopassword.txt")
+				_, err := os.ReadFile("nopassword.txt")
 				require.Error(t, err)
 			})
 		})

privval/bls.go

@@ -16,29 +16,34 @@ import (
 )
 
 const (
-	DefaultBlsKeyName      = "bls_key.json"
-	DefaultBlsPasswordName = "bls_password.txt"
+	DefaultBlsKeyName      = "bls_key.json"     // Default file name for BLS key
+	DefaultBlsPasswordName = "bls_password.txt" // Default file name for BLS password
 )
 
 var (
-	defaultBlsKeyFilePath  = filepath.Join(cmtcfg.DefaultConfigDir, DefaultBlsKeyName)
-	defaultBlsPasswordPath = filepath.Join(cmtcfg.DefaultConfigDir, DefaultBlsPasswordName)
+	defaultBlsKeyFilePath  = filepath.Join(cmtcfg.DefaultConfigDir, DefaultBlsKeyName)      // Default file path for BLS key
+	defaultBlsPasswordPath = filepath.Join(cmtcfg.DefaultConfigDir, DefaultBlsPasswordName) // Default file path for BLS password
 )
 
+// BlsPV is a wrapper around BlsPVKey
 type BlsPV struct {
+	// Key is a structure containing bls12381 keys,
+	// paths of both key and password files,
+	// and delegator address
 	Key BlsPVKey
 }
 
+// BlsPVKey is a wrapper containing bls12381 keys,
+// paths of both key and password files, and delegator address.
 type BlsPVKey struct {
-	PubKey  bls12381.PublicKey  `json:"bls_pub_key"`
-	PrivKey bls12381.PrivateKey `json:"bls_priv_key"`
-
-	DelegatorAddress string
-
-	filePath     string
-	passwordPath string
+	PubKey           bls12381.PublicKey  `json:"bls_pub_key"`       // Public Key of BLS
+	PrivKey          bls12381.PrivateKey `json:"bls_priv_key"`      // Private Key of BLS
+	DelegatorAddress string              `json:"delegator_address"` // Delegate Address
+	filePath         string              // File Path of BLS Key
+	passwordPath     string              // File Path of BLS Password
 }
 
+// NewBlsPV returns a new BlsPV.
 func NewBlsPV(privKey bls12381.PrivateKey, keyFilePath, passwordFilePath, delegatorAddress string) *BlsPV {
 	return &BlsPV{
 		Key: BlsPVKey{
@@ -51,17 +56,21 @@ func NewBlsPV(privKey bls12381.PrivateKey, keyFilePath, passwordFilePath, delega
 	}
 }
 
+// GenBlsPV returns a new BlsPV after saving it to the file.
 func GenBlsPV(keyFilePath, passwordFilePath, password, delegatorAddress string) *BlsPV {
 	pv := NewBlsPV(bls12381.GenPrivKey(), keyFilePath, passwordFilePath, delegatorAddress)
 	pv.Key.Save(password, delegatorAddress)
 	return pv
 }
 
+// LoadBlsPV returns a BlsPV after loading the erc2335 type of structure
+// from the file and decrypt it using a password.
 func LoadBlsPV(keyFilePath, passwordFilePath string) *BlsPV {
-	password, err := erc2335.LoadPaswordFromFile(passwordFilePath)
+	passwordBytes, err := os.ReadFile(passwordFilePath)
 	if err != nil {
 		cmtos.Exit(fmt.Sprintf("failed to read BLS password file: %v", err.Error()))
 	}
+	password := string(passwordBytes)
 
 	keystore, err := erc2335.LoadKeyStore(keyFilePath)
 	if err != nil {
@@ -86,6 +95,7 @@ func LoadBlsPV(keyFilePath, passwordFilePath string) *BlsPV {
 	}
 }
 
+// NewBlsPassword returns a password from the user prompt.
 func NewBlsPassword() string {
 	inBuf := bufio.NewReader(os.Stdin)
 	password, err := input.GetString("Enter your bls password", inBuf)
@@ -95,8 +105,8 @@ func NewBlsPassword() string {
 	return password
 }
 
-// Save bls key using password
-// Check both paths of bls key and password inside function
+// Save saves the bls12381 key to the file.
+// The file stores an erc2335 structure containing the encrypted bls private key.
 func (k *BlsPVKey) Save(password, addr string) {
 	// encrypt the bls12381 key to erc2335 type
 	erc2335BlsPvKey, err := erc2335.Encrypt(k.PrivKey, k.PubKey.Bytes(), password)
@@ -125,16 +135,17 @@ func (k *BlsPVKey) Save(password, addr string) {
 	}
 
 	// save used password to file
-	err = erc2335.SavePasswordToFile(password, k.passwordPath)
-	if err != nil {
+	if err := tempfile.WriteFileAtomic(k.passwordPath, []byte(password), 0600); err != nil {
 		panic(err)
 	}
 }
 
+// DefaultBlsKeyFile returns the default BLS key file path.
 func DefaultBlsKeyFile(home string) string {
 	return filepath.Join(home, defaultBlsKeyFilePath)
 }
 
+// DefaultBlsPasswordFile returns the default BLS password file path.
 func DefaultBlsPasswordFile(home string) string {
 	return filepath.Join(home, defaultBlsPasswordPath)
 }

privval/bls_test.go

@@ -7,7 +7,6 @@ import (
 	"github.com/cometbft/cometbft/crypto/ed25519"
 
 	"github.com/babylonlabs-io/babylon/crypto/bls12381"
-	"github.com/babylonlabs-io/babylon/crypto/erc2335"
 	"github.com/cosmos/cosmos-sdk/types"
 	"github.com/test-go/testify/assert"
 )
@@ -26,7 +25,7 @@ func TestNewBlsPV(t *testing.T) {
 		pv := NewBlsPV(bls12381.GenPrivKey(), keyFilePath, passwordFilePath, "")
 		assert.NotNil(t, pv)
 
-		password := erc2335.CreateRandomPassword()
+		password := "password"
 		pv.Key.Save(password, "")
 
 		t.Run("load bls key from file", func(t *testing.T) {
@@ -42,7 +41,7 @@ func TestNewBlsPV(t *testing.T) {
 		pv := NewBlsPV(bls12381.GenPrivKey(), keyFilePath, passwordFilePath, "")
 		assert.NotNil(t, pv)
 
-		password := erc2335.CreateRandomPassword()
+		password := "password"
 
 		delegatorAddress := types.AccAddress(ed25519.GenPrivKey().PubKey().Address()).String()
 		pv.Key.Save(password, delegatorAddress)

test/e2e/initialization/node.go

@@ -168,7 +168,7 @@ func (n *internalNode) createConsensusKey() error {
 	blsPasswordFile := privval.DefaultBlsPasswordFile(n.configDir())
 
 	if err := privval.EnsureDirs(pvKeyFile, pvStateFile, blsKeyFile, blsPasswordFile); err != nil {
-		return err
+		return fmt.Errorf("failed to ensure dirs: %w", err)
 	}
 	accAddress, _ := n.keyInfo.GetAddress()
 

testutil/datagen/init_val.go

@@ -36,7 +36,7 @@ func InitializeNodeValidatorFilesFromMnemonic(config *cfg.Config, mnemonic strin
 	blsKeyFile := privval.DefaultBlsKeyFile(config.RootDir)
 	blsPasswordFile := privval.DefaultBlsPasswordFile(config.RootDir)
 	if err := privval.EnsureDirs(cmtKeyFile, cmtStateFile, blsKeyFile, blsPasswordFile); err != nil {
-		return "", nil, err
+		return "", nil, fmt.Errorf("failed to ensure dirs: %w", err)
 	}
 
 	var filePV *cmtprivval.FilePV

testutil/signer/private.go

@@ -1,6 +1,7 @@
 package signer
 
 import (
+	"fmt"
 	"os"
 
 	cryptocodec "github.com/cosmos/cosmos-sdk/crypto/codec"
@@ -60,7 +61,7 @@ func GeneratePrivSigner(nodeDir string) error {
 	blsPasswordFile := privval.DefaultBlsPasswordFile(nodeDir)
 
 	if err := privval.EnsureDirs(cmtKeyFile, cmtStateFile, blsKeyFile, blsPasswordFile); err != nil {
-		return err
+		return fmt.Errorf("failed to ensure dirs: %w", err)
 	}
 
 	cometPV := cmtprivval.GenFilePV(cmtKeyFile, cmtStateFile)

x/checkpointing/client/cli/utils.go

@@ -212,7 +212,7 @@ func getValKeyFromFile(homeDir string) (*privval.ValidatorKeys, error) {
 	blsPasswordPath := privval.DefaultBlsPasswordFile(homeDir)
 
 	if err := privval.EnsureDirs(cmtKeyPath, cmtStatePath, blsKeyPath, blsPasswordPath); err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to ensure dirs: %w", err)
 	}
 
 	filePV := cmtprivval.LoadFilePV(cmtKeyPath, cmtStatePath)