From 4517d00cc1f3f3a1462db62c09e770e1931cd27c Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Tue, 13 Aug 2024 13:34:51 +0800
Subject: [PATCH] sm2: fix RecoverPublicKeysFromSM2Signature

---
 sm2/sm2.go      |  2 +-
 sm2/sm2_test.go | 35 +++++++++++++++++++++++++++++++++--
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/sm2/sm2.go b/sm2/sm2.go
index 9af620d7..642ed9a9 100644
--- a/sm2/sm2.go
+++ b/sm2/sm2.go
@@ -762,7 +762,7 @@ func RecoverPublicKeysFromSM2Signature(hash, sig []byte) ([]*ecdsa.PublicKey, er
 		pointRx = append(pointRx, s)
 	}
 	pubs := make([]*ecdsa.PublicKey, 0, 4)
-	bytes := make([]byte, len(rBytes)+1)
+	bytes := make([]byte, len(32)+1)
 	compressFlags := []byte{compressed02, compressed03}
 	// Rx has one or two possible values, so point R has two or four possible values
 	for _, x := range pointRx {
diff --git a/sm2/sm2_test.go b/sm2/sm2_test.go
index 97c79ea7..1311324c 100644
--- a/sm2/sm2_test.go
+++ b/sm2/sm2_test.go
@@ -470,8 +470,7 @@ func TestSignVerify(t *testing.T) {
 	}
 }
 
-func TestRecoverPublicKeysFromSM2Signature(t *testing.T) {
-	priv, _ := GenerateKey(rand.Reader)
+func testRecoverPublicKeysFromSM2Signature(t *testing.T, priv *PrivateKey) {
 	tests := []struct {
 		name      string
 		plainText string
@@ -511,6 +510,38 @@ func TestRecoverPublicKeysFromSM2Signature(t *testing.T) {
 	}
 }
 
+func TestRecoverPublicKeysFromSM2Signature(t *testing.T) {
+	priv, _ := GenerateKey(rand.Reader)
+	testRecoverPublicKeysFromSM2Signature(t, priv)
+	keyInt := bigFromHex("d6833540d019e0438a5dd73b414f26ab43d8064b99671206944e284dbd969093")
+	priv, _ = NewPrivateKeyFromInt(keyInt)
+	testRecoverPublicKeysFromSM2Signature(t, priv)
+
+	// failed case
+	hashValue, _ := CalculateSM2Hash(&priv.PublicKey, []byte("encryption standard encryption "), nil)
+	signature, _ := hex.DecodeString("3045022000cd0b56bf6be810032d28ff27d6f3468f1f1a09bcf8581f30a5de6692c85ea602210096ba29c086134af1be139dd572f2f2908f30e01fd0c28e06a687cbb0ff6e33ce")
+	// verify signature with public key
+	if !VerifyASN1(&priv.PublicKey, hashValue, signature) {
+		t.Errorf("failed to verify hash for sig=%x, priv=%x", signature, priv.D.Bytes())
+	}
+	pubs, err := RecoverPublicKeysFromSM2Signature(hashValue, signature)
+	if err != nil {
+		t.Fatalf("recover failed %v", err)
+	}
+	found := false
+	for _, pub := range pubs {
+		if !VerifyASN1(pub, hashValue, signature) {
+			t.Errorf("failed to verify hash for sig=%x, priv=%x", signature, priv.D.Bytes())
+		}
+		if pub.Equal(&priv.PublicKey) {
+			found = true
+		}
+	}
+	if !found {
+		t.Errorf("recover failed, not found public key for sig=%x, priv=%x", signature, priv.D.Bytes())
+	}
+}
+
 func TestSignVerifyLegacy(t *testing.T) {
 	priv, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
 	tests := []struct {