From e34c65eb5f7ca61eab47da0e3433e245e30f22c6 Mon Sep 17 00:00:00 2001
From: Roman Behma <13855864+begmaroman@users.noreply.github.com>
Date: Wed, 11 Oct 2023 15:15:14 +0800
Subject: [PATCH] Store BLS private key in hex instead of raw big int  (#1973)

---
 command/polybftsecrets/params.go      |  2 +-
 command/polybftsecrets/params_test.go |  4 +--
 consensus/polybft/signer/private.go   | 46 +++++++++++++++++++--------
 crypto/crypto.go                      | 12 -------
 4 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/command/polybftsecrets/params.go b/command/polybftsecrets/params.go
index ec521ecbfc..aaa6a35808 100644
--- a/command/polybftsecrets/params.go
+++ b/command/polybftsecrets/params.go
@@ -226,7 +226,7 @@ func (ip *initParams) getResult(
 				return nil, err
 			}
 
-			res.BLSPrivateKey = hex.EncodeToString(blspk)
+			res.BLSPrivateKey = string(blspk)
 		}
 	}
 
diff --git a/command/polybftsecrets/params_test.go b/command/polybftsecrets/params_test.go
index 2d140973b6..c9cbd0aa15 100644
--- a/command/polybftsecrets/params_test.go
+++ b/command/polybftsecrets/params_test.go
@@ -100,9 +100,7 @@ func Test_getResult(t *testing.T) {
 	assert.Equal(t, sir.Address.String(), pubKey)
 
 	// Test BLS public key serialization
-	blsPrivKeyRaw, err := hex.DecodeString(sir.BLSPrivateKey)
-	require.NoError(t, err)
-	blsPrivKey, err := bls.UnmarshalPrivateKey(blsPrivKeyRaw)
+	blsPrivKey, err := bls.UnmarshalPrivateKey([]byte(sir.BLSPrivateKey))
 	require.NoError(t, err)
 
 	blsPubKey := hex.EncodeToString(blsPrivKey.PublicKey().Marshal())
diff --git a/consensus/polybft/signer/private.go b/consensus/polybft/signer/private.go
index e122ed5999..c696338d14 100644
--- a/consensus/polybft/signer/private.go
+++ b/consensus/polybft/signer/private.go
@@ -2,8 +2,11 @@ package bls
 
 import (
 	"math/big"
+	"strings"
 
 	bn256 "github.com/umbracle/go-eth-bn256"
+
+	"github.com/0xPolygon/polygon-edge/helper/hex"
 )
 
 // PrivateKey holds private key for bls implementation
@@ -11,6 +14,34 @@ type PrivateKey struct {
 	s *big.Int
 }
 
+// NewZeroPrivateKey is the constructor of an empty PrivateKey
+func NewZeroPrivateKey() *PrivateKey {
+	return &PrivateKey{
+		s: new(big.Int),
+	}
+}
+
+// UnmarshalPrivateKey unmarshals the private key from the given byte slice.
+// This function supports both raw big int string and hex-encoded big int string.
+func UnmarshalPrivateKey(data []byte) (*PrivateKey, error) {
+	pk := NewZeroPrivateKey()
+
+	// First trying to use a default unmarshaling function of big int.
+	// It works for the raw big int string format and for hex with 0x prefix.
+	if err := pk.s.UnmarshalText(data); err == nil {
+		return pk, nil
+	}
+
+	// Otherwise, trying to assume the given data is a hex encoded big int represented as a bytes array.
+	// This is needed in order to be compatible with the currently stored polybft BLS keys.
+	var err error
+	if pk.s, err = hex.DecodeHexToBig(string(data)); err != nil {
+		return nil, err
+	}
+
+	return pk, nil
+}
+
 // PublicKey returns the public key from the PrivateKey
 func (p *PrivateKey) PublicKey() *PublicKey {
 	g2 := new(bn256.G2).ScalarMult(g2Point, p.s)
@@ -18,9 +49,9 @@ func (p *PrivateKey) PublicKey() *PublicKey {
 	return &PublicKey{g2: g2}
 }
 
-// Marshal marshals private key to byte slice
+// Marshal marshals private key hex (without 0x prefix) represented as a byte slice
 func (p *PrivateKey) Marshal() ([]byte, error) {
-	return p.s.MarshalText()
+	return []byte(strings.TrimPrefix(hex.EncodeBig(p.s), "0x")), nil
 }
 
 // Sign generates a simple BLS signature of the given message
@@ -34,14 +65,3 @@ func (p *PrivateKey) Sign(message, domain []byte) (*Signature, error) {
 
 	return &Signature{g1: g1}, nil
 }
-
-// UnmarshalPrivateKey unmarshals the private key from the given byte slice
-func UnmarshalPrivateKey(data []byte) (*PrivateKey, error) {
-	s := new(big.Int)
-
-	if err := s.UnmarshalText(data); err != nil {
-		return nil, err
-	}
-
-	return &PrivateKey{s: s}, nil
-}
diff --git a/crypto/crypto.go b/crypto/crypto.go
index ae2a9cf420..daf8f858c0 100644
--- a/crypto/crypto.go
+++ b/crypto/crypto.go
@@ -369,18 +369,6 @@ func BLSSecretKeyToPubkeyBytes(key *bls_sig.SecretKey) ([]byte, error) {
 	return marshalled, nil
 }
 
-// BytesToBLSPublicKey decodes given hex string and returns BLS Public Key
-func BytesToBLSPublicKey(input string) (*bls_sig.PublicKey, error) {
-	// The key file on disk should be encoded in Base64,
-	// so it must be decoded before it can be parsed by ParsePrivateKey
-	decoded, err := hex.DecodeString(input)
-	if err != nil {
-		return nil, err
-	}
-
-	return UnmarshalBLSPublicKey(decoded)
-}
-
 // UnmarshalBLSPublicKey unmarshal bytes data into BLS Public Key
 func UnmarshalBLSPublicKey(input []byte) (*bls_sig.PublicKey, error) {
 	pk := &bls_sig.PublicKey{}