From 4a059c11af4dc0a8e4f047563443ce5a4c77d9b3 Mon Sep 17 00:00:00 2001
From: Haaai <55118568+Liuhaai@users.noreply.github.com>
Date: Wed, 10 Apr 2024 18:51:07 -0700
Subject: [PATCH] refactor secp256k1 pubkey (#40)

Co-authored-by: dustinxie <dahuaxie@gmail.com>
---
 crypto/secp256k1.go      | 36 ++++++++++++++++++++++++++++--------
 crypto/secp256k1_test.go | 16 ++++++++++++++++
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/crypto/secp256k1.go b/crypto/secp256k1.go
index 6bbb5eb..1f1aaf8 100644
--- a/crypto/secp256k1.go
+++ b/crypto/secp256k1.go
@@ -24,7 +24,7 @@ type (
 	}
 	// secp256k1PubKey implements the SECP256K1 public key
 	secp256k1PubKey struct {
-		*ecdsa.PublicKey
+		raw []byte
 	}
 )
 
@@ -72,7 +72,7 @@ func (k *secp256k1PrvKey) EcdsaPrivateKey() interface{} {
 // PublicKey returns the public key corresponding to private key
 func (k *secp256k1PrvKey) PublicKey() PublicKey {
 	return &secp256k1PubKey{
-		PublicKey: &k.PrivateKey.PublicKey,
+		raw: crypto.FromECDSAPub(&k.PrivateKey.PublicKey),
 	}
 }
 
@@ -93,20 +93,35 @@ func (k *secp256k1PrvKey) Zero() {
 // PublicKey function
 //======================================
 
+var (
+	secp256k1PubKeyByteLen = (crypto.S256().Params().BitSize + 7) >> 3
+)
+
 // newSecp256k1PubKeyFromBytes converts bytes format to PublicKey
 func newSecp256k1PubKeyFromBytes(b []byte) (PublicKey, error) {
-	pk, err := crypto.UnmarshalPubkey(b)
-	if err != nil {
-		return nil, err
+	if !validateP256k1PubkeyBytes(b) {
+		return nil, ErrPublicKey
 	}
+
 	return &secp256k1PubKey{
-		PublicKey: pk,
+		raw: b,
 	}, nil
 }
 
+func validateP256k1PubkeyBytes(data []byte) bool {
+	if len(data) != 1+2*secp256k1PubKeyByteLen {
+		return false
+	}
+	if data[0] != 4 { // uncompressed form
+		return false
+	}
+
+	return true
+}
+
 // Bytes returns the public key in bytes representation
 func (k *secp256k1PubKey) Bytes() []byte {
-	return crypto.FromECDSAPub(k.PublicKey)
+	return k.raw
 }
 
 // HexString returns the public key in hex string
@@ -116,7 +131,12 @@ func (k *secp256k1PubKey) HexString() string {
 
 // EcdsaPublicKey returns the embedded ecdsa publick key
 func (k *secp256k1PubKey) EcdsaPublicKey() interface{} {
-	return k.PublicKey
+	pk, err := crypto.UnmarshalPubkey(k.raw)
+	if err != nil {
+		// it should be validated when initialized
+		panic(err)
+	}
+	return pk
 }
 
 // Hash is the last 20-byte of keccak hash of public key (X, Y) co-ordinate, same as Ethereum address generation
diff --git a/crypto/secp256k1_test.go b/crypto/secp256k1_test.go
index cc0694c..642aeba 100644
--- a/crypto/secp256k1_test.go
+++ b/crypto/secp256k1_test.go
@@ -69,3 +69,19 @@ func TestSecp256k1(t *testing.T) {
 	require.Equal("53fbc28faf9a52dfe5f591948a23189e900381b5", hex.EncodeToString(pk.Hash()))
 
 }
+
+func BenchmarkSecp256k1(b *testing.B) {
+	require := require.New(b)
+
+	sk, err := newSecp256k1PrvKey()
+	require.NoError(err)
+	defer sk.Zero()
+	pk := sk.PublicKey()
+	require.Equal(secp256pubKeyLength, len(pk.Bytes()))
+	b.ResetTimer()
+	pkData := pk.Bytes()
+	for n := 0; n < b.N; n++ {
+		_, err := newSecp256k1PubKeyFromBytes(pkData)
+		require.NoError(err)
+	}
+}