From 586ab84f33e556069a2fe4f2e8e7b0a886c55f43 Mon Sep 17 00:00:00 2001
From: Emman <emman.sun@cargosmart.com>
Date: Thu, 25 Feb 2021 13:37:14 +0800
Subject: [PATCH] MAGIC - move public key check out of loop

---
 sm2/sm2.go | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/sm2/sm2.go b/sm2/sm2.go
index cad4f3c..0bb2688 100644
--- a/sm2/sm2.go
+++ b/sm2/sm2.go
@@ -192,6 +192,10 @@ func Encrypt(random io.Reader, pub *ecdsa.PublicKey, msg []byte, opts *Encrypter
 	if opts == nil {
 		opts = &defaultEncrypterOpts
 	}
+	//A3, requirement is to check if h*P is infinite point, h is 1
+	if (pub.X.Sign() == 0 && pub.Y.Sign() == 0) || !curve.IsOnCurve(pub.X, pub.Y) {
+		return nil, errors.New("SM2: invalid public key")
+	}
 	for {
 		//A1, generate random k
 		k, err := randFieldElement(curve, random)
@@ -203,11 +207,6 @@ func Encrypt(random io.Reader, pub *ecdsa.PublicKey, msg []byte, opts *Encrypter
 		x1, y1 := curve.ScalarBaseMult(k.Bytes())
 		c1 := opts.PointMarshalMode.mashal(curve, x1, y1)
 
-		//A3, requirement is to check if h*P is infinite point, h is 1
-		if !curve.IsOnCurve(pub.X, pub.Y) {
-			return nil, errors.New("SM2: invalid public key")
-		}
-
 		//A4, calculate k * P (point of Public Key)
 		x2, y2 := curve.ScalarMult(pub.X, pub.Y, k.Bytes())