From fe532e12b421e6853d2224c344e3d23f99f957cf Mon Sep 17 00:00:00 2001
From: YuanHongYe <yuanhongye@chinatelecom.cn>
Date: Tue, 15 Apr 2025 15:41:20 +0800
Subject: [PATCH] sm9: Prevent PublicKey() returning nil after unmarshaling
 master private key

During unmarshaling of SignMasterPrivateKey and EncryptMasterPrivateKey,
now generate the corresponding public key.
This ensures that PublicKey() does not return nil.

Test cases included to validate the changes.

Signed-off-by: YuanHongYe <yuanhongye@chinatelecom.cn>
---
 sm9/sm9_key.go      | 16 ++++++++++++++--
 sm9/sm9_key_test.go | 10 ++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/sm9/sm9_key.go b/sm9/sm9_key.go
index a5b6de5..bd63679 100644
--- a/sm9/sm9_key.go
+++ b/sm9/sm9_key.go
@@ -114,7 +114,13 @@ func UnmarshalSignMasterPrivateKeyASN1(der []byte) (*SignMasterPrivateKey, error
 	if err != nil {
 		return nil, err
 	}
-	return &SignMasterPrivateKey{privateKey: priv.Bytes(), internal: priv}, nil
+
+	master := &SignMasterPrivateKey{privateKey: priv.Bytes(), internal: priv}
+	master.publicKey = &SignMasterPublicKey{
+		publicKey: priv.PublicKey().Bytes(),
+		internal:  priv.PublicKey(),
+	}
+	return master, nil
 }
 
 // GenerateUserKey generate a signature private key for the given user.
@@ -370,7 +376,13 @@ func UnmarshalEncryptMasterPrivateKeyASN1(der []byte) (*EncryptMasterPrivateKey,
 	if err != nil {
 		return nil, err
 	}
-	return &EncryptMasterPrivateKey{privateKey: privateKey.Bytes(), internal: privateKey}, nil
+
+	master := &EncryptMasterPrivateKey{privateKey: privateKey.Bytes(), internal: privateKey}
+	master.publicKey = &EncryptMasterPublicKey{
+		publicKey: privateKey.PublicKey().Bytes(),
+		internal:  privateKey.PublicKey(),
+	}
+	return master, nil
 }
 
 // Equal compares the receiver EncryptMasterPublicKey with another EncryptMasterPublicKey
diff --git a/sm9/sm9_key_test.go b/sm9/sm9_key_test.go
index 70cc12b..66f16e6 100644
--- a/sm9/sm9_key_test.go
+++ b/sm9/sm9_key_test.go
@@ -27,6 +27,11 @@ func TestSignMasterPrivateKeyMarshalASN1(t *testing.T) {
 	if !masterKey.Equal(masterKey2) {
 		t.Errorf("expected %v, got %v", hex.EncodeToString(masterKey.Bytes()), hex.EncodeToString(masterKey2.Bytes()))
 	}
+
+	masterPubKey := masterKey2.PublicKey()
+	if masterPubKey == nil {
+		t.Fatal("cannot export public key")
+	}
 }
 
 func TestSignMasterPublicKeyMarshalASN1(t *testing.T) {
@@ -129,6 +134,11 @@ func TestEncryptMasterPrivateKeyMarshalASN1(t *testing.T) {
 	if !masterKey.Equal(masterKey2) {
 		t.Errorf("expected %v, got %v", hex.EncodeToString(masterKey.Bytes()), hex.EncodeToString(masterKey2.Bytes()))
 	}
+
+	masterPubKey := masterKey2.PublicKey()
+	if masterPubKey == nil {
+		t.Fatal("cannot export public key")
+	}
 }
 
 func TestEncryptMasterPublicKeyMarshalASN1(t *testing.T) {