修改了设置对端参数名称，增加了防止重复设置的检查。

sm2/sm2_keyexchange.go

@@ -37,8 +37,8 @@ func (ke *KeyExchange) GetSharedKey() []byte {
 // NewKeyExchange create one new KeyExchange object
 //
 // 在部分场景中，在初始  KeyExchange 时暂时没有对端的公开信息（如公钥、UID），这些信息可能需要在后续的交换中得到。
-// 这种情况下，可设置 peerPub、peerUID 参数为 nil，并在合适的时候通过 KeyExchange.SetPeerPub 方法配置相关参数。
+// 这种情况下，可设置 peerPub、peerUID 参数为 nil，并在合适的时候通过 KeyExchange.SetPeerParameters 方法配置相关参数。
-// 注意 KeyExchange.SetPeerPub 方法必须要在 KeyExchange.RepondKeyExchange 或 KeyExchange.RepondKeyExchange 方法调用。
+// 注意 KeyExchange.SetPeerParameters 方法必须要在 KeyExchange.RepondKeyExchange 或 KeyExchange.RepondKeyExchange 方法之前调用。
 func NewKeyExchange(priv *PrivateKey, peerPub *ecdsa.PublicKey, uid, peerUID []byte, keyLen int, genSignature bool) (ke *KeyExchange, err error) {
 	ke = &KeyExchange{}
 	ke.genSignature = genSignature
@@ -60,7 +60,7 @@ func NewKeyExchange(priv *PrivateKey, peerPub *ecdsa.PublicKey, uid, peerUID []b
 		return nil, err
 	}
 
-	err = ke.SetPeerPub(peerPub, peerUID)
+	err = ke.SetPeerParameters(peerPub, peerUID)
 	if err != nil {
 		return nil, err
 	}
@@ -74,15 +74,21 @@ func NewKeyExchange(priv *PrivateKey, peerPub *ecdsa.PublicKey, uid, peerUID []b
 	return
 }
 
-// SetPeerPub 设置对端公开信息，该方法用于某些初期状态无法取得对端公开参数的场景。
+// SetPeerParameters 设置对端公开信息，该方法用于某些初期状态无法取得对端公开参数的场景。
 // 例如：在TLCP协议中，基于SM2算法ECDHE过程。
-func (ke *KeyExchange) SetPeerPub(peerPub *ecdsa.PublicKey, peerUID []byte) error {
+//
+// 注意该方法仅在 NewKeyExchange 没有提供 peerPub、peerUID参数时允许被调用，
+// 且该方法只能调用一次不可重复调用，若多次调用或peerPub、peerUID已经存在则会发生错误。
+func (ke *KeyExchange) SetPeerParameters(peerPub *ecdsa.PublicKey, peerUID []byte) error {
 	if peerPub == nil {
 		return nil
 	}
 	if len(peerUID) == 0 {
 		peerUID = defaultUID
 	}
+	if ke.peerPub != nil {
+		return errors.New("sm2: 'peerPub' already exists, please do not set it")
+	}
 
 	var err error
 	ke.peerPub = peerPub
@@ -184,7 +190,7 @@ func respondKeyExchange(ke *KeyExchange, r *big.Int, rA *ecdsa.PublicKey) (*ecds
 // RepondKeyExchange when responder receive rA, for responder's step B1-B8
 func (ke *KeyExchange) RepondKeyExchange(rand io.Reader, rA *ecdsa.PublicKey) (*ecdsa.PublicKey, []byte, error) {
 	if ke.peerPub == nil {
-		return nil, nil, errors.New("sm2: peer public not set, you probable need call KeyExchange.SetPeerPub")
+		return nil, nil, errors.New("sm2: peer public not set, you probable need call KeyExchange.SetPeerParameters")
 	}
 	if !ke.privateKey.IsOnCurve(rA.X, rA.Y) {
 		return nil, nil, errors.New("sm2: received invalid random from initiator")
@@ -200,7 +206,7 @@ func (ke *KeyExchange) RepondKeyExchange(rand io.Reader, rA *ecdsa.PublicKey) (*
 // ConfirmResponder for initiator's step A4-A10
 func (ke *KeyExchange) ConfirmResponder(rB *ecdsa.PublicKey, sB []byte) ([]byte, error) {
 	if ke.peerPub == nil {
-		return nil, errors.New("sm2: peer public not set, you probable need call KeyExchange.SetPeerPub")
+		return nil, errors.New("sm2: peer public not set, you probable need call KeyExchange.SetPeerParameters")
 	}
 	if !ke.privateKey.IsOnCurve(rB.X, rB.Y) {
 		return nil, errors.New("sm2: received invalid random from responder")

sm2/sm2_keyexchange_test.go

@@ -3,6 +3,7 @@ package sm2
 import (
 	"crypto/rand"
 	"encoding/hex"
+	"errors"
 	"testing"
 )
 
@@ -42,7 +43,7 @@ func TestKeyExchangeSample(t *testing.T) {
 	}
 }
 
-func TestKeyExchangeNoPeerPubInit(t *testing.T) {
+func TestSetPeerParameters(t *testing.T) {
 	priv1, _ := GenerateKey(rand.Reader)
 	priv2, _ := GenerateKey(rand.Reader)
 	uidA := []byte("Alice")
@@ -63,11 +64,11 @@ func TestKeyExchangeNoPeerPubInit(t *testing.T) {
 	}
 
 	// 设置对端参数
-	err = initiator.SetPeerPub(&priv2.PublicKey, uidB)
+	err = initiator.SetPeerParameters(&priv2.PublicKey, uidB)
 	if err != nil {
 		t.Fatal(err)
 	}
-	err = responder.SetPeerPub(&priv1.PublicKey, uidA)
+	err = responder.SetPeerParameters(&priv1.PublicKey, uidA)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -91,3 +92,98 @@ func TestKeyExchangeNoPeerPubInit(t *testing.T) {
 		t.Errorf("got different key")
 	}
 }
+
+func TestKeyExchange_SetPeerParameters(t *testing.T) {
+	priv1, _ := GenerateKey(rand.Reader)
+	priv2, _ := GenerateKey(rand.Reader)
+	uidA := []byte("Alice")
+	uidB := []byte("Bob")
+
+	initiator, err := NewKeyExchange(priv1, nil, uidA, nil, 32, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+	responder, err := NewKeyExchange(priv2, nil, uidB, nil, 32, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rA, err := initiator.InitKeyExchange(rand.Reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// 设置对端参数
+	err = initiator.SetPeerParameters(&priv2.PublicKey, uidB)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = responder.SetPeerParameters(&priv1.PublicKey, uidA)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rB, s2, err := responder.RepondKeyExchange(rand.Reader, rA)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	s1, err := initiator.ConfirmResponder(rB, s2)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = responder.ConfirmInitiator(s1)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if hex.EncodeToString(initiator.key) != hex.EncodeToString(responder.key) {
+		t.Errorf("got different key")
+	}
+}
+
+func TestKeyExchange_SetPeerParameters_ErrCase(t *testing.T) {
+	priv1, _ := GenerateKey(rand.Reader)
+	priv2, _ := GenerateKey(rand.Reader)
+	uidA := []byte("Alice")
+	uidB := []byte("Bob")
+
+	initiator, err := NewKeyExchange(priv1, nil, uidA, nil, 32, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+	responder, err := NewKeyExchange(priv2, &priv1.PublicKey, uidB, uidA, 32, true)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	rA, err := initiator.InitKeyExchange(rand.Reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+	rB, s2, err := responder.RepondKeyExchange(rand.Reader, rA)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = initiator.ConfirmResponder(rB, s2)
+	if err == nil {
+		t.Fatal(errors.New("expect call ConfirmResponder got a error, but not"))
+	}
+
+	err = initiator.SetPeerParameters(&priv2.PublicKey, uidB)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = initiator.SetPeerParameters(&priv2.PublicKey, uidB)
+	if err == nil {
+		t.Fatal(errors.New("expect call SetPeerParameters repeat got a error, but not"))
+	}
+
+	err = responder.SetPeerParameters(&priv1.PublicKey, uidA)
+	if err == nil {
+		t.Fatal(errors.New("expect responder call SetPeerParameters got a error, but not"))
+	}
+}