From 86cfa509412e22aa6eea33b5de29f9a933d37e92 Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Fri, 29 Aug 2025 10:34:47 +0800
Subject: [PATCH] sm9: add back SetMasterPublic methods

---
 sm9/sm9_key.go  | 14 ++++++++++++++
 sm9/sm9_test.go | 14 ++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/sm9/sm9_key.go b/sm9/sm9_key.go
index bd63679..f449365 100644
--- a/sm9/sm9_key.go
+++ b/sm9/sm9_key.go
@@ -235,6 +235,13 @@ func (priv *SignPrivateKey) MasterPublic() *SignMasterPublicKey {
 	return &SignMasterPublicKey{internal: masterKey, publicKey: masterKey.Bytes()}
 }
 
+// SetMasterPublic sets the master public key for the SignPrivateKey.
+// The caller should ensure that the provided master public key is valid and corresponds
+// to the SignPrivateKey. This method is NOT safe for concurrent use.
+func (priv *SignPrivateKey) SetMasterPublic(master *SignMasterPublicKey) {
+	priv.internal.SetMasterPublicKey(master.internal)
+}
+
 // MarshalASN1 marshal signature private key to asn.1 format data according
 // SM9 cryptographic algorithm application specification
 func (priv *SignPrivateKey) MarshalASN1() ([]byte, error) {
@@ -463,6 +470,13 @@ func (priv *EncryptPrivateKey) MasterPublic() *EncryptMasterPublicKey {
 	return &EncryptMasterPublicKey{publicKey: master.Bytes(), internal: master}
 }
 
+// SetMasterPublic sets the master public key for the EncryptPrivateKey.
+// The caller should ensure that the provided master public key is valid and corresponds
+// to the EncryptPrivateKey. This method is NOT safe for concurrent use.
+func (priv *EncryptPrivateKey) SetMasterPublic(master *EncryptMasterPublicKey) {
+	priv.internal.SetMasterPublicKey(master.internal)
+}
+
 // MarshalASN1 marshal encryption private key to asn.1 format data according
 // SM9 cryptographic algorithm application specification
 func (priv *EncryptPrivateKey) MarshalASN1() ([]byte, error) {
diff --git a/sm9/sm9_test.go b/sm9/sm9_test.go
index efa1c22..2913d9e 100644
--- a/sm9/sm9_test.go
+++ b/sm9/sm9_test.go
@@ -21,6 +21,13 @@ func TestSignASN1(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	// Test Marshal and Unmarshal
+	userKeyBytes := userKey.Bytes()
+	userKey, err = sm9.UnmarshalSignPrivateKeyRaw(userKeyBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	userKey.SetMasterPublic(masterKey.PublicKey())
 	sig, err := userKey.Sign(rand.Reader, hashed, nil)
 	if err != nil {
 		t.Fatal(err)
@@ -103,6 +110,13 @@ func TestEncryptDecrypt(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	// Test Marshal and Unmarshal
+	userKeyBytes := userKey.Bytes()
+	userKey, err = sm9.UnmarshalEncryptPrivateKeyRaw(userKeyBytes)
+	if err != nil {
+		t.Fatal(err)
+	}
+	userKey.SetMasterPublic(masterKey.PublicKey())
 	encTypes := []sm9.EncrypterOpts{
 		sm9.DefaultEncrypterOpts, sm9.SM4ECBEncrypterOpts, sm9.SM4CBCEncrypterOpts, sm9.SM4CFBEncrypterOpts, sm9.SM4OFBEncrypterOpts,
 	}