From cd8d60fc161147372fddb6fd005837bb0acb0345 Mon Sep 17 00:00:00 2001
From: Emman <emman.sun@cargosmart.com>
Date: Wed, 20 Jan 2021 15:05:24 +0800
Subject: [PATCH] MAGIC - implement MarshalPKIXPublicKey

---
 sm2/x509.go      | 43 ++++++++++++++++++++++++++++++++++++++++---
 sm2/x509_test.go | 20 ++++++++++++++++++--
 2 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/sm2/x509.go b/sm2/x509.go
index 85402c5..ea2347e 100644
--- a/sm2/x509.go
+++ b/sm2/x509.go
@@ -3,12 +3,20 @@ package sm2
 import (
 	"crypto/ecdsa"
 	"crypto/elliptic"
+	"crypto/x509"
 	"crypto/x509/pkix"
 	"encoding/asn1"
 	"errors"
 	"math/big"
 )
 
+// pkixPublicKey reflects a PKIX public key structure. See SubjectPublicKeyInfo
+// in RFC 3280.
+type pkixPublicKey struct {
+	Algo      pkix.AlgorithmIdentifier
+	BitString asn1.BitString
+}
+
 type publicKeyInfo struct {
 	Raw       asn1.RawContent
 	Algorithm pkix.AlgorithmIdentifier
@@ -45,13 +53,12 @@ func ParsePKIXPublicKey(derBytes []byte) (interface{}, error) {
 	}
 
 	if !pki.Algorithm.Algorithm.Equal(oidPublicKeyECDSA) {
-		return nil, errors.New("x509: invalid public key algorithm")
+		return x509.ParsePKIXPublicKey(derBytes)
 	}
 	keyData := &pki
 	asn1Data := keyData.PublicKey.RightAlign()
 	paramsData := keyData.Algorithm.Parameters.FullBytes
 	namedCurveOID := new(asn1.ObjectIdentifier)
-	namedCurve := P256()
 	rest, err := asn1.Unmarshal(paramsData, namedCurveOID)
 	if err != nil {
 		return nil, errors.New("x509: failed to parse ECDSA parameters as named curve")
@@ -60,8 +67,9 @@ func ParsePKIXPublicKey(derBytes []byte) (interface{}, error) {
 		return nil, errors.New("x509: trailing data after ECDSA parameters")
 	}
 	if !namedCurveOID.Equal(oidNamedCurveP256SM2) {
-		return nil, errors.New("x509: it's not SM2 elliptic curve")
+		return x509.ParsePKIXPublicKey(derBytes)
 	}
+	namedCurve := P256()
 	x, y := elliptic.Unmarshal(namedCurve, asn1Data)
 	if x == nil {
 		return nil, errors.New("x509: failed to unmarshal elliptic curve point")
@@ -73,3 +81,32 @@ func ParsePKIXPublicKey(derBytes []byte) (interface{}, error) {
 	}
 	return pub, nil
 }
+
+// MarshalPKIXPublicKey converts a public key to PKIX, ASN.1 DER form.
+//
+// The following key types are currently supported: *rsa.PublicKey, *ecdsa.PublicKey
+// and ed25519.PublicKey. Unsupported key types result in an error.
+//
+// This kind of key is commonly encoded in PEM blocks of type "PUBLIC KEY".
+func MarshalPKIXPublicKey(pub interface{}) ([]byte, error) {
+	ecdPub, ok := pub.(*ecdsa.PublicKey)
+	if !ok || ecdPub.Curve != P256() {
+		return x509.MarshalPKIXPublicKey(pub)
+	}
+
+	publicKeyAlgorithm := pkix.AlgorithmIdentifier{Algorithm: oidPublicKeyECDSA}
+	publicKeyBytes := elliptic.Marshal(ecdPub.Curve, ecdPub.X, ecdPub.Y)
+	paramBytes, _ := asn1.Marshal(oidNamedCurveP256SM2)
+	publicKeyAlgorithm.Parameters.FullBytes = paramBytes
+
+	pkix := pkixPublicKey{
+		Algo: publicKeyAlgorithm,
+		BitString: asn1.BitString{
+			Bytes:     publicKeyBytes,
+			BitLength: 8 * len(publicKeyBytes),
+		},
+	}
+
+	ret, _ := asn1.Marshal(pkix)
+	return ret, nil
+}
diff --git a/sm2/x509_test.go b/sm2/x509_test.go
index 3bcb961..f433b8d 100644
--- a/sm2/x509_test.go
+++ b/sm2/x509_test.go
@@ -5,11 +5,11 @@ import (
 	"crypto/rand"
 	"encoding/pem"
 	"errors"
+	"strings"
 	"testing"
 )
 
-const publicKeyPemFromAliKms = `
------BEGIN PUBLIC KEY-----
+const publicKeyPemFromAliKms = `-----BEGIN PUBLIC KEY-----
 MFkwEwYHKoZIzj0CAQYIKoEcz1UBgi0DQgAELfjZP28bYfGSvbODYlXiB5bcoXE+
 2LRjjpIH3DcCCct9FuVhi9cm60nDFrbW49k2D3GJco2iWPlr0+5LV+t4AQ==
 -----END PUBLIC KEY-----
@@ -34,3 +34,19 @@ func TestParsePKIXPublicKey(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestMarshalPKIXPublicKey(t *testing.T) {
+	pub, err := getPublicKey([]byte(publicKeyPemFromAliKms))
+	if err != nil {
+		t.Fatal(err)
+	}
+	result, err := MarshalPKIXPublicKey(pub)
+	if err != nil {
+		t.Fatal(err)
+	}
+	block := &pem.Block{Bytes: result, Type: "PUBLIC KEY"}
+	pemContent := string(pem.EncodeToMemory(block))
+	if !strings.EqualFold(publicKeyPemFromAliKms, pemContent) {
+		t.Errorf("expected=%s, result=%s", publicKeyPemFromAliKms, pemContent)
+	}
+}