From 6ecb7c2e606175a16c94b6ce00a0d3487ef14adc Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Wed, 20 Nov 2024 11:45:13 -0500
Subject: [PATCH] update NeedKms logic

---
 etc/install-libmongocrypt.sh                              | 2 +-
 mongo/client_encryption.go                                | 2 ++
 x/mongo/driver/crypt.go                                   | 6 ++++++
 x/mongo/driver/mongocrypt/mongocrypt.go                   | 5 +++++
 x/mongo/driver/mongocrypt/mongocrypt_kms_context.go       | 8 ++++++++
 .../mongocrypt/mongocrypt_kms_context_not_enabled.go      | 5 +++++
 x/mongo/driver/mongocrypt/mongocrypt_not_enabled.go       | 5 +++++
 7 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/etc/install-libmongocrypt.sh b/etc/install-libmongocrypt.sh
index 646721a8f7..a94d648eae 100755
--- a/etc/install-libmongocrypt.sh
+++ b/etc/install-libmongocrypt.sh
@@ -3,7 +3,7 @@
 # This script installs libmongocrypt into an "install" directory.
 set -eux
 
-LIBMONGOCRYPT_TAG="1.11.0"
+LIBMONGOCRYPT_TAG="1.12.0"
 
 # Install libmongocrypt based on OS.
 if [ "Windows_NT" = "${OS:-}" ]; then
diff --git a/mongo/client_encryption.go b/mongo/client_encryption.go
index 07c18529a8..66bb73259c 100644
--- a/mongo/client_encryption.go
+++ b/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
diff --git a/x/mongo/driver/crypt.go b/x/mongo/driver/crypt.go
index 4368fd125d..890e22886e 100644
--- a/x/mongo/driver/crypt.go
+++ b/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
 		}
 
diff --git a/x/mongo/driver/mongocrypt/mongocrypt.go b/x/mongo/driver/mongocrypt/mongocrypt.go
index 5f34f5cd71..85b4dd1b65 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt.go
+++ b/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)
+}
diff --git a/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go b/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go
index 296a22315c..f5242650af 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go
+++ b/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)
+}
diff --git a/x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go b/x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go
index 6bce2f0299..608d3784f1 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go
+++ b/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)
+}
diff --git a/x/mongo/driver/mongocrypt/mongocrypt_not_enabled.go b/x/mongo/driver/mongocrypt/mongocrypt_not_enabled.go
index 6e21e64917..f64e5cb7f6 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt_not_enabled.go
+++ b/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)
+}