Improve HexToPubkey & HexToSignature (#122)

* Improve HexToPubkey & HexToSignature * Fix lint issues & use tabled tests * Add more tests * Disable dupl linter * Fix typo * Remove "valid" field from test table
flashbots · Feb 29, 2024 · 6bf58ee · 6bf58ee
1 parent 09199fc
commit 6bf58ee
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 12 deletions.
diff --git a/.golangci.yaml b/.golangci.yaml
@@ -1,6 +1,7 @@
 linters:
   enable-all: true
   disable:
+    - dupl
     - exhaustruct
     - funlen
     - gochecknoglobals

diff --git a/utils/utils.go b/utils/utils.go
@@ -15,6 +15,7 @@ import (
 	"github.com/attestantio/go-eth2-client/spec/phase0"
 	utilbellatrix "github.com/attestantio/go-eth2-client/util/bellatrix"
 	utilcapella "github.com/attestantio/go-eth2-client/util/capella"
+	bls12381 "github.com/consensys/gnark-crypto/ecc/bls12-381"
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -28,6 +29,8 @@ var (
 	ErrNilPayload         = errors.New("nil payload")
 	ErrUnsupportedVersion = errors.New("unsupported version")
 	ErrUnknownVersion     = errors.New("unknown version")
+	ErrInvalidPubkey      = errors.New("invalid pubkey")
+	ErrInvalidSignature   = errors.New("invalid signature")
 )
 
 func BlsPublicKeyToPublicKey(blsPubKey *bls.PublicKey) (ret phase0.BLSPubKey, err error) {
@@ -61,19 +64,33 @@ func HexToAddress(s string) (ret bellatrix.ExecutionAddress, err error) {
 // HexToPubkey takes a hex string and returns a PublicKey
 func HexToPubkey(s string) (ret phase0.BLSPubKey, err error) {
 	bytes, err := hexutil.Decode(s)
+	if err != nil {
+		return phase0.BLSPubKey{}, err
+	}
 	if len(bytes) != len(ret) {
 		return phase0.BLSPubKey{}, ErrLength
 	}
+	_, err = new(bls12381.G1Affine).SetBytes(bytes)
+	if err != nil {
+		return phase0.BLSPubKey{}, ErrInvalidPubkey
+	}
 	copy(ret[:], bytes)
 	return
 }
 
 // HexToSignature takes a hex string and returns a Signature
 func HexToSignature(s string) (ret phase0.BLSSignature, err error) {
 	bytes, err := hexutil.Decode(s)
+	if err != nil {
+		return phase0.BLSSignature{}, err
+	}
 	if len(bytes) != len(ret) {
 		return phase0.BLSSignature{}, ErrLength
 	}
+	_, err = new(bls12381.G2Affine).SetBytes(bytes)
+	if err != nil {
+		return phase0.BLSSignature{}, ErrInvalidSignature
+	}
 	copy(ret[:], bytes)
 	return
 }

diff --git a/utils/utils_test.go b/utils/utils_test.go
@@ -32,21 +32,93 @@ func TestHexToAddress(t *testing.T) {
 }
 
 func TestHexToPubkey(t *testing.T) {
-	_, err := HexToPubkey("0x01")
-	require.Error(t, err)
-
-	a, err := HexToPubkey("0xed7f862045422bd51ba732730ce993c94d2545e5db1112102026343904fcdf6f5cf37926a3688444703772ed80fa223f")
-	require.NoError(t, err)
-	require.Equal(t, "0xed7f862045422bd51ba732730ce993c94d2545e5db1112102026343904fcdf6f5cf37926a3688444703772ed80fa223f", a.String())
+	testCases := []struct {
+		name        string
+		pubkey      string
+		expectedErr string
+	}{
+		{
+			name:   "Valid pubkey",
+			pubkey: "0xac6e77dfe25ecd6110b8e780608cce0dab71fdd5ebea22a16c0205200f2f8e2e3ad3b71d3499c54ad14d6c21b41a37ae",
+		},
+		{
+			name:        "Invalid pubkey (wrong length)",
+			pubkey:      "0x123456",
+			expectedErr: "invalid length",
+		},
+		{
+			name:        "Invalid pubkey (not on the curve)",
+			pubkey:      "0xed7f862045422bd51ba732730ce993c94d2545e5db1112102026343904fcdf6f5cf37926a3688444703772ed80fa223f",
+			expectedErr: "invalid pubkey",
+		},
+		{
+			name:        "Invalid pubkey (no 0x prefix)",
+			pubkey:      "this is not hex",
+			expectedErr: "hex string without 0x prefix",
+		},
+		{
+			name:        "Invalid pubkey (not hex)",
+			pubkey:      "0xthisisnothex",
+			expectedErr: "invalid hex string",
+		},
+	}
+
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := HexToPubkey(tt.pubkey)
+			if tt.expectedErr != "" {
+				require.EqualError(t, err, tt.expectedErr)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.pubkey, result.String())
+			}
+		})
+	}
 }
 
 func TestHexToSignature(t *testing.T) {
-	_, err := HexToSignature("0x01")
-	require.Error(t, err)
-
-	a, err := HexToSignature("0xb8f03e639b91fa8e9892f66c798f07f6e7b3453234f643b2c06a35c5149cf6d85e4e1572c33549fe749292445fbff9e0739c78159324c35dc1a90e5745ca70c8caf1b63fb6678d81bd2d5cb6baeb1462df7a93877d0e22a31dd6438334536d9a")
-	require.NoError(t, err)
-	require.Equal(t, "0xb8f03e639b91fa8e9892f66c798f07f6e7b3453234f643b2c06a35c5149cf6d85e4e1572c33549fe749292445fbff9e0739c78159324c35dc1a90e5745ca70c8caf1b63fb6678d81bd2d5cb6baeb1462df7a93877d0e22a31dd6438334536d9a", a.String())
+	testCases := []struct {
+		name        string
+		signature   string
+		expectedErr string
+	}{
+		{
+			name:      "Valid signature",
+			signature: "0x8069aa021666163aae46d353c348aa913fb5050062e05ab764c8eef99407a8befcd46190b30cc40e40ee8c197356959816799d62e85f640ef76b4be1b08a741949230fbde49589125537daad06c23a66838725d89e3504bc21559a91534f6712",
+		},
+		{
+			name:        "Invalid signature (wrong length)",
+			signature:   "0x123456",
+			expectedErr: "invalid length",
+		},
+		{
+			name:        "Invalid signature (not on the curve)",
+			signature:   "0xb8f03e639b91fa8e9892f66c798f07f6e7b3453234f643b2c06a35c5149cf6d85e4e1572c33549fe749292445fbff9e0739c78159324c35dc1a90e5745ca70c8caf1b63fb6678d81bd2d5cb6baeb1462df7a93877d0e22a31dd6438334536d9a",
+			expectedErr: "invalid signature",
+		},
+		{
+			name:        "Invalid signature (no 0x prefix)",
+			signature:   "this is not hex",
+			expectedErr: "hex string without 0x prefix",
+		},
+		{
+			name:        "Invalid signature (not hex)",
+			signature:   "0xthisisnothex",
+			expectedErr: "invalid hex string",
+		},
+	}
+
+	for _, tt := range testCases {
+		t.Run(tt.name, func(t *testing.T) {
+			result, err := HexToSignature(tt.signature)
+			if tt.expectedErr != "" {
+				require.EqualError(t, err, tt.expectedErr)
+			} else {
+				require.NoError(t, err)
+				require.Equal(t, tt.signature, result.String())
+			}
+		})
+	}
 }
 
 func TestComputeHash(t *testing.T) {