From ef7ebafb8e88c97687d01f2ed7bcb4373ca06717 Mon Sep 17 00:00:00 2001
From: shawn <shawn.huang@tron.network>
Date: Wed, 22 Mar 2023 20:56:52 +0800
Subject: [PATCH] fix: key secp256 marshal issue

---
 crypto/crypto.go       |  2 +-
 crypto/keys.go         |  2 +-
 crypto/tron_address.go | 21 ++++++++++++++++-----
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/crypto/crypto.go b/crypto/crypto.go
index 4b23d3d..0ea60d8 100644
--- a/crypto/crypto.go
+++ b/crypto/crypto.go
@@ -176,7 +176,7 @@ func Decrypt(key, text []byte) ([]byte, error) {
 }
 
 func GetPubKeyFromPeerId(pid string) (ic.PubKey, error) {
-	peerId, err := peer.IDFromBytes([]byte(pid))
+	peerId, err := peer.Decode(pid)
 	if err != nil {
 		return nil, err
 	}
diff --git a/crypto/keys.go b/crypto/keys.go
index c5a3c57..d112e1b 100644
--- a/crypto/keys.go
+++ b/crypto/keys.go
@@ -36,7 +36,7 @@ func FromIcPrivateKey(privKey ic.PrivKey) (*Keys, error) {
 		return nil, err
 	}
 
-	pubKeyRaw, err := ic.MarshalPublicKey(pubKey)
+	pubKeyRaw, err := Secp256k1PublicKeyRaw(pubKey)
 	if err != nil {
 		return nil, err
 	}
diff --git a/crypto/tron_address.go b/crypto/tron_address.go
index 53f4c95..bac5af8 100644
--- a/crypto/tron_address.go
+++ b/crypto/tron_address.go
@@ -4,7 +4,9 @@ import (
 	"crypto/ecdsa"
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
 
+	"github.com/decred/dcrd/dcrec/secp256k1/v4"
 	"github.com/tron-us/go-common/v2/crypto"
 	"github.com/tron-us/protobuf/proto"
 
@@ -41,7 +43,7 @@ func GetTronPubKeyFromPubkey(pubkeyS string) (*string, error) {
 }
 
 func GetTronPubKeyFromPeerIdPretty(peerId string) (*string, error) {
-	pid, err := peer.IDFromBytes([]byte(peerId))
+	pid, err := peer.Decode(peerId)
 	if err != nil {
 		return nil, err
 	}
@@ -80,12 +82,11 @@ func TronSignRaw(privKey ic.PrivKey, data []byte) ([]byte, error) {
 }
 
 func GetTronPubKeyFromIcPubKey(pubkey ic.PubKey) (*string, error) {
-	pubkeyRaw, err := ic.MarshalPublicKey(pubkey)
+	rawPubKey, err := Secp256k1PublicKeyRaw(pubkey)
 	if err != nil {
 		return nil, err
 	}
-
-	ethPubkey, err := eth.UnmarshalPubkey(pubkeyRaw)
+	ethPubkey, err := eth.UnmarshalPubkey(rawPubKey)
 	if err != nil {
 		return nil, err
 	}
@@ -117,7 +118,7 @@ func EcdsaPublicKeyToAddress(p ecdsa.PublicKey) (Address, error) {
 }
 
 func GetRawFullFromPeerIdPretty(peerid string) ([]byte, error) {
-	peerId, err := peer.IDFromBytes([]byte(peerid))
+	peerId, err := peer.Decode(peerid)
 	if err != nil {
 		return nil, err
 	}
@@ -127,3 +128,13 @@ func GetRawFullFromPeerIdPretty(peerid string) ([]byte, error) {
 	}
 	return pubkey.Raw()
 }
+
+// Raw returns the bytes of the key
+func Secp256k1PublicKeyRaw(pk ic.PubKey) (res []byte, err error) {
+	// defer func() { catch.HandlePanic(recover(), &err, "secp256k1 public key marshaling") }()
+	k, ok := pk.(*ic.Secp256k1PublicKey)
+	if !ok {
+		return nil, fmt.Errorf("only secp256k1 keys support full public key bytes")
+	}
+	return (*secp256k1.PublicKey)(k).SerializeUncompressed(), nil
+}