update NeedKms logic

Author: qingyang-hu

mongo/client_encryption.go

@@ -188,12 +188,14 @@ func (ce *ClientEncryption) CreateDataKey(
 	}
 
 	// create data key document
+	fmt.Println("CreateDataKey")
 	dataKeyDoc, err := ce.crypt.CreateDataKey(ctx, kmsProvider, co)
 	if err != nil {
 		return bson.Binary{}, err
 	}
 
 	// insert key into key vault
+	fmt.Println("InsertOne")
 	_, err = ce.keyVaultColl.InsertOne(ctx, dataKeyDoc)
 	if err != nil {
 		return bson.Binary{}, err

x/mongo/driver/crypt.go

@@ -341,6 +341,7 @@ func (c *crypt) retrieveKeys(ctx context.Context, cryptCtx *mongocrypt.Context)
 }
 
 func (c *crypt) decryptKeys(cryptCtx *mongocrypt.Context) error {
+	c.mongoCrypt.EnableRetry()
 	for {
 		kmsCtx := cryptCtx.NextKmsContext()
 		if kmsCtx == nil {
@@ -400,6 +401,11 @@ func (c *crypt) decryptKey(kmsCtx *mongocrypt.KmsContext) error {
 		res := make([]byte, bytesNeeded)
 		bytesRead, err := conn.Read(res)
 		if err != nil && !errors.Is(err, io.EOF) {
+			fail := kmsCtx.Fail()
+			fmt.Println("conn read", err, fail)
+			if fail {
+				continue
+			}
 			return err
 		}
 

x/mongo/driver/mongocrypt/mongocrypt.go

@@ -522,3 +522,8 @@ func (m *MongoCrypt) GetKmsProviders(ctx context.Context) (bsoncore.Document, er
 	}
 	return builder.Build(), nil
 }
+
+// EnableRetry enables retry.
+func (m *MongoCrypt) EnableRetry() {
+	_ = C.mongocrypt_setopt_retry_kms(m.wrapped, true)
+}

x/mongo/driver/mongocrypt/mongocrypt_kms_context.go

@@ -11,6 +11,7 @@ package mongocrypt
 
 // #include <mongocrypt.h>
 import "C"
+import "time"
 
 // KmsContext represents a mongocrypt_kms_ctx_t handle.
 type KmsContext struct {
@@ -41,6 +42,8 @@ func (kc *KmsContext) KMSProvider() string {
 
 // Message returns the message to send to the KMS.
 func (kc *KmsContext) Message() ([]byte, error) {
+	time.Sleep(C.mongocrypt_kms_ctx_usleep(kc.wrapped) * time.Microsecond)
+
 	msgBinary := newBinary()
 	defer msgBinary.close()
 
@@ -74,3 +77,8 @@ func (kc *KmsContext) createErrorFromStatus() error {
 	C.mongocrypt_kms_ctx_status(kc.wrapped, status)
 	return errorFromStatus(status)
 }
+
+// Fail returns a boolean indicating whether the failed request may be retried.
+func (kc *KmsContext) Fail() bool {
+	return C.mongocrypt_kms_ctx_fail(kc.wrapped)
+}

x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go

@@ -37,3 +37,8 @@ func (kc *KmsContext) BytesNeeded() int32 {
 func (kc *KmsContext) FeedResponse([]byte) error {
 	panic(cseNotSupportedMsg)
 }
+
+// Fail returns a boolean indicating whether the failed request may be retried.
+func (kc *KmsContext) Fail() bool {
+	panic(cseNotSupportedMsg)
+}

x/mongo/driver/mongocrypt/mongocrypt_not_enabled.go

@@ -95,3 +95,8 @@ func (m *MongoCrypt) Close() {
 func (m *MongoCrypt) GetKmsProviders(context.Context) (bsoncore.Document, error) {
 	panic(cseNotSupportedMsg)
 }
+
+// EnableRetry enables retry.
+func (m *MongoCrypt) EnableRetry() {
+	panic(cseNotSupportedMsg)
+}