crypto add kms request retry

Author: jojoliang

crypto/master_kms_cipher.go

@@ -9,6 +9,7 @@ import (
 	"github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/common/profile"
 	kms "github.com/tencentcloud/tencentcloud-sdk-go/tencentcloud/kms/v20190118"
 	"github.com/tencentyun/cos-go-sdk-v5"
+	"time"
 )
 
 var (
@@ -19,6 +20,11 @@ type MasterKMSCipher struct {
 	Client  *kms.Client
 	KmsId   string
 	MatDesc string
+	Conf    MasterKMSCipherConf
+}
+type MasterKMSCipherConf struct {
+	Retry         int
+	RetryInternal int64 // ms
 }
 
 type KMSClientOptions = func(*profile.HttpProfile)
@@ -29,6 +35,20 @@ func KMSEndpoint(endpoint string) KMSClientOptions {
 	}
 }
 
+type MasterKMSCipherOptions = func(*MasterKMSCipher)
+
+func KMSRetry(retry int) MasterKMSCipherOptions {
+	return func(c *MasterKMSCipher) {
+		c.Conf.Retry = retry
+	}
+}
+
+func KMSRetryInternal(internal int64) MasterKMSCipherOptions {
+	return func(c *MasterKMSCipher) {
+		c.Conf.RetryInternal = internal
+	}
+}
+
 func NewKMSClient(cred *cos.Credential, region string, opt ...KMSClientOptions) (*kms.Client, error) {
 	if cred == nil {
 		fmt.Errorf("credential is nil")
@@ -48,11 +68,10 @@ func NewKMSClient(cred *cos.Credential, region string, opt ...KMSClientOptions)
 	return client, err
 }
 
-func CreateMasterKMS(client *kms.Client, kmsId string, desc map[string]string) (MasterCipher, error) {
+func CreateMasterKMS(client *kms.Client, kmsId string, desc map[string]string, opts ...MasterKMSCipherOptions) (MasterCipher, error) {
 	if kmsId == "" || client == nil {
 		return nil, fmt.Errorf("KMS ID is empty or kms client is nil")
 	}
-	var kmsCipher MasterKMSCipher
 	var jdesc string
 	if len(desc) > 0 {
 		bs, err := json.Marshal(desc)
@@ -61,33 +80,62 @@ func CreateMasterKMS(client *kms.Client, kmsId string, desc map[string]string) (
 		}
 		jdesc = string(bs)
 	}
-	kmsCipher.Client = client
-	kmsCipher.KmsId = kmsId
-	kmsCipher.MatDesc = jdesc
-	return &kmsCipher, nil
+	kmsCipher := &MasterKMSCipher{
+		Client:  client,
+		KmsId:   kmsId,
+		MatDesc: jdesc,
+		Conf: MasterKMSCipherConf{
+			Retry:         3,
+			RetryInternal: 100,
+		},
+	}
+	for _, fn := range opts {
+		fn(kmsCipher)
+	}
+	return kmsCipher, nil
 }
 
 func (kc *MasterKMSCipher) Encrypt(plaintext []byte) ([]byte, error) {
 	request := kms.NewEncryptRequest()
 	request.KeyId = common.StringPtr(kc.KmsId)
 	request.EncryptionContext = common.StringPtr(kc.MatDesc)
 	request.Plaintext = common.StringPtr(base64.StdEncoding.EncodeToString(plaintext))
-	resp, err := kc.Client.Encrypt(request)
-	if err != nil {
-		return nil, err
+
+	var retry int
+	for {
+		resp, err := kc.Client.Encrypt(request)
+		if err != nil {
+			retry++
+			if retry < kc.Conf.Retry {
+				time.Sleep(time.Duration(kc.Conf.RetryInternal) * time.Millisecond)
+				continue
+			}
+			return nil, err
+		}
+		return []byte(*resp.Response.CiphertextBlob), nil
 	}
-	return []byte(*resp.Response.CiphertextBlob), nil
+	return nil, nil
 }
 
 func (kc *MasterKMSCipher) Decrypt(ciphertext []byte) ([]byte, error) {
 	request := kms.NewDecryptRequest()
 	request.CiphertextBlob = common.StringPtr(string(ciphertext))
 	request.EncryptionContext = common.StringPtr(kc.MatDesc)
-	resp, err := kc.Client.Decrypt(request)
-	if err != nil {
-		return nil, err
+
+	var retry int
+	for {
+		resp, err := kc.Client.Decrypt(request)
+		if err != nil {
+			retry++
+			if retry < kc.Conf.Retry {
+				time.Sleep(time.Duration(kc.Conf.RetryInternal) * time.Millisecond)
+				continue
+			}
+			return nil, err
+		}
+		return base64.StdEncoding.DecodeString(*resp.Response.Plaintext)
 	}
-	return base64.StdEncoding.DecodeString(*resp.Response.Plaintext)
+	return nil, nil
 }
 
 func (kc *MasterKMSCipher) GetWrapAlgorithm() string {