From 02a60a50c9a5aa399492f014a64fac7f7367edf8 Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Mon, 11 Nov 2024 14:44:39 -0500
Subject: [PATCH 1/9] GODRIVER-3168 Retry KMS requests on transient errors.

---
 .../client_side_encryption_prose_test.go      | 66 ++++++++++++++++++-
 1 file changed, 65 insertions(+), 1 deletion(-)

diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index 0620b9ff10..4b6966c47b 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -13,7 +13,9 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
+	"crypto/x509"
 	"encoding/base64"
+	"encoding/json"
 	"fmt"
 	"io/ioutil"
 	"net"
@@ -30,6 +32,7 @@ import (
 	"go.mongodb.org/mongo-driver/v2/internal/handshake"
 	"go.mongodb.org/mongo-driver/v2/internal/integration/mtest"
 	"go.mongodb.org/mongo-driver/v2/internal/integtest"
+	"go.mongodb.org/mongo-driver/v2/internal/require"
 	"go.mongodb.org/mongo-driver/v2/mongo"
 	"go.mongodb.org/mongo-driver/v2/mongo/options"
 	"go.mongodb.org/mongo-driver/v2/mongo/writeconcern"
@@ -2918,7 +2921,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		}
 	})
 
-	mt.RunOpts("22. range explicit encryption applies defaults", qeRunOpts22, func(mt *mtest.T) {
+	mt.RunOpts("23. range explicit encryption applies defaults", qeRunOpts22, func(mt *mtest.T) {
 		err := mt.Client.Database("keyvault").Collection("datakeys").Drop(context.Background())
 		assert.Nil(mt, err, "error on Drop: %v", err)
 
@@ -2979,6 +2982,67 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			assert.Greater(t, len(payload.Data), len(payloadDefaults.Data), "the returned payload size is expected to be greater than %d", len(payloadDefaults.Data))
 		})
 	})
+
+	mt.RunOpts("24. KMS Retry Tests", qeRunOpts22, func(mt *mtest.T) {
+		setFailPoint := func(failure string, count int) error {
+			url := fmt.Sprintf("https://localhost:9003/set_failpoint/%s", failure)
+			var payloadBuf bytes.Buffer
+			body := map[string]int{"count": count}
+			json.NewEncoder(&payloadBuf).Encode(body)
+			req, err := http.NewRequest(http.MethodPost, url, &payloadBuf)
+			if err != nil {
+				return err
+			}
+
+			cert, err := ioutil.ReadFile(os.Getenv("CSFLE_TLS_CA_FILE"))
+			if err != nil {
+				return err
+			}
+
+			certPool := x509.NewCertPool()
+			certPool.AppendCertsFromPEM(cert)
+
+			client := &http.Client{
+				Transport: &http.Transport{
+					TLSClientConfig: &tls.Config{
+						RootCAs: certPool,
+					},
+				},
+			}
+			_, err = client.Do(req)
+			return err
+		}
+
+		keyVaultClient, err := mongo.Connect(options.Client().ApplyURI(mtest.ClusterURI()))
+		require.NoError(mt, err, "error on Connect: %v", err)
+
+		ceo := options.ClientEncryption().
+			SetKeyVaultNamespace("keyvault.datakeys").
+			SetKmsProviders(fullKmsProvidersMap)
+		clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
+		require.NoError(mt, err, "error on NewClientEncryption: %v", err)
+
+		err = setFailPoint("http", 1)
+		require.NoError(mt, err, "mock server error: %v", err)
+
+		dkOpts := options.DataKey().SetMasterKey(
+			bson.D{
+				{"region", "foo"},
+				{"key", "bar"},
+				{"endpoint", "127.0.0.1:9003"},
+			},
+		)
+		var keyID bson.Binary
+		keyID, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
+		require.NoError(mt, err, "error in CreateDataKey: %v", err)
+
+		testVal := bson.RawValue{Type: bson.TypeInt32, Value: bsoncore.AppendInt32(nil, 123)}
+		eo := options.Encrypt().
+			SetKeyID(keyID).
+			SetAlgorithm("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic")
+		_, err = clientEncryption.Encrypt(context.Background(), testVal, eo)
+		assert.NoError(mt, err, "error in Encrypt: %v", err)
+	})
 }
 
 func getWatcher(mt *mtest.T, streamType mongo.StreamType, cpt *cseProseTest) watcher {

From 12a15309ee70e347c692cefa9ac9c865ab8b5f0b Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Wed, 13 Nov 2024 13:09:34 -0500
Subject: [PATCH 2/9] add task

---
 .evergreen/config.yml                         | 57 +++++++++++++++++++
 Taskfile.yml                                  |  3 +
 .../client_side_encryption_prose_test.go      | 21 ++++---
 3 files changed, 70 insertions(+), 11 deletions(-)

diff --git a/.evergreen/config.yml b/.evergreen/config.yml
index 90da768350..d2bf4b101b 100644
--- a/.evergreen/config.yml
+++ b/.evergreen/config.yml
@@ -553,6 +553,42 @@ functions:
           KMS_MOCK_SERVERS_RUNNING: "true"
         args: [*task-runner, evg-test-kmip]
 
+  start-kms-failpoint-server:
+    - command: ec2.assume_role
+      params:
+        role_arn: ${aws_test_secrets_role}
+    - command: subprocess.exec
+      params:
+        working_dir: src/go.mongodb.org/mongo-driver
+        binary: bash
+        background: true
+        include_expansions_in_env: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "DRIVERS_TOOLS"]
+        # This cannot use task because it will hang on Windows.
+        args: [etc/setup-encryption.sh]
+    - command: subprocess.exec
+      params:
+        binary: python3
+        background: true
+        args: ["-u", "${DRIVERS_TOOLS}/.evergreen/csfle/kms_failpoint_server.py", "--port", "9003"]
+
+  run-retry-kms-requests:
+    - command: subprocess.exec
+      type: test
+      params:
+        binary: "bash"
+        env: 
+          GO_BUILD_TAGS: cse
+        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, 
+          MONGO_GO_DRIVER_COMPRESSOR]
+        args: [*task-runner, setup-test]
+    - command: subprocess.exec
+      type: test
+      params:
+        binary: "bash"
+        env:
+          KMS_FAILPOINT_SERVERS_RUNNING: "true"
+        args: [*task-runner, evg-test-retry-kms-requests]
+
   run-fuzz-tests:
     - command: subprocess.exec
       type: test
@@ -1486,6 +1522,21 @@ tasks:
           AUTH: "noauth"
           SSL: "nossl"
 
+  - name: "test-retry-kms-requests"
+    tags: ["retry-kms-requests"]
+    commands:
+      - func: bootstrap-mongo-orchestration
+        vars:
+          TOPOLOGY: "server"
+          AUTH: "noauth"
+          SSL: "nossl"
+      - func: start-kms-failpoint-server
+      - func: run-retry-kms-requests
+        vars:
+          TOPOLOGY: "server"
+          AUTH: "noauth"
+          SSL: "nossl"
+
   - name: "test-serverless"
     tags: ["serverless"]
     commands:
@@ -2195,6 +2246,12 @@ buildvariants:
     tasks:
       - name: ".kms-kmip"
 
+  - matrix_name: "retry-kms-requests-test"
+    matrix_spec: { version: ["7.0"], os-ssl-40: ["rhel87-64"] }
+    display_name: "Retry KMS Requests ${os-ssl-40}"
+    tasks:
+      - name: ".retry-kms-requests"
+
   - matrix_name: "fuzz-test"
     matrix_spec: { version: ["5.0"], os-ssl-40: ["rhel87-64"] }
     display_name: "Fuzz ${version} ${os-ssl-40}"
diff --git a/Taskfile.yml b/Taskfile.yml
index f22427a640..d93a7e4e9b 100644
--- a/Taskfile.yml
+++ b/Taskfile.yml
@@ -141,6 +141,9 @@ tasks:
   evg-test-kms:
     - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_tls_tests >> test.suite
 
+  evg-test-retry-kms-requests:
+    - go test -exec "env PKG_CONFIG_PATH=${PKG_CONFIG_PATH} LD_LIBRARY_PATH=${LD_LIBRARY_PATH}" ${BUILD_TAGS} -v -timeout {{.TEST_TIMEOUT}}s ./internal/integration -run TestClientSideEncryptionProse/kms_retry_tests >> test.suite
+
   evg-test-load-balancers:
     # Load balancer should be tested with all unified tests as well as tests in the following
     # components: retryable reads, retryable writes, change streams, initial DNS seedlist discovery.
diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index 4b6966c47b..a476644f85 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -13,7 +13,6 @@ import (
 	"bytes"
 	"context"
 	"crypto/tls"
-	"crypto/x509"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
@@ -2983,7 +2982,12 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		})
 	})
 
-	mt.RunOpts("24. KMS Retry Tests", qeRunOpts22, func(mt *mtest.T) {
+	mt.RunOpts("24. kms retry tests", noClientOpts, func(mt *mtest.T) {
+		kmsTlsTestcase := os.Getenv("KMS_FAILPOINT_SERVERS_RUNNING")
+		if kmsTlsTestcase == "" {
+			mt.Skipf("Skipping test as KMS_FAILPOINT_SERVERS_RUNNING is not set")
+		}
+
 		setFailPoint := func(failure string, count int) error {
 			url := fmt.Sprintf("https://localhost:9003/set_failpoint/%s", failure)
 			var payloadBuf bytes.Buffer
@@ -2994,18 +2998,10 @@ func TestClientSideEncryptionProse(t *testing.T) {
 				return err
 			}
 
-			cert, err := ioutil.ReadFile(os.Getenv("CSFLE_TLS_CA_FILE"))
-			if err != nil {
-				return err
-			}
-
-			certPool := x509.NewCertPool()
-			certPool.AppendCertsFromPEM(cert)
-
 			client := &http.Client{
 				Transport: &http.Transport{
 					TLSClientConfig: &tls.Config{
-						RootCAs: certPool,
+						InsecureSkipVerify: true,
 					},
 				},
 			}
@@ -3036,6 +3032,9 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		keyID, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
 		require.NoError(mt, err, "error in CreateDataKey: %v", err)
 
+		err = setFailPoint("http", 1)
+		require.NoError(mt, err, "mock server error: %v", err)
+
 		testVal := bson.RawValue{Type: bson.TypeInt32, Value: bsoncore.AppendInt32(nil, 123)}
 		eo := options.Encrypt().
 			SetKeyID(keyID).

From 00e69f5ed0ee7764ecbad04cf76a3e717b3f9058 Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Fri, 22 Nov 2024 14:13:31 -0500
Subject: [PATCH 3/9] update NeedKms logic

---
 etc/install-libmongocrypt.sh                  |  2 +-
 .../client_side_encryption_prose_test.go      | 24 +++++++++++--------
 x/mongo/driver/crypt.go                       |  7 +++---
 x/mongo/driver/mongocrypt/mongocrypt.go       |  3 ++-
 .../mongocrypt/mongocrypt_kms_context.go      |  8 +++++++
 .../mongocrypt_kms_context_not_enabled.go     |  5 ++++
 6 files changed, 34 insertions(+), 15 deletions(-)

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/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index a476644f85..3222b79467 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -2988,6 +2988,10 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			mt.Skipf("Skipping test as KMS_FAILPOINT_SERVERS_RUNNING is not set")
 		}
 
+		tlsCfg := &tls.Config{
+			InsecureSkipVerify: true,
+		}
+
 		setFailPoint := func(failure string, count int) error {
 			url := fmt.Sprintf("https://localhost:9003/set_failpoint/%s", failure)
 			var payloadBuf bytes.Buffer
@@ -2999,14 +3003,13 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			}
 
 			client := &http.Client{
-				Transport: &http.Transport{
-					TLSClientConfig: &tls.Config{
-						InsecureSkipVerify: true,
-					},
-				},
+				Transport: &http.Transport{TLSClientConfig: tlsCfg},
 			}
-			_, err = client.Do(req)
-			return err
+			res, err := client.Do(req)
+			if err != nil {
+				return err
+			}
+			return res.Body.Close()
 		}
 
 		keyVaultClient, err := mongo.Connect(options.Client().ApplyURI(mtest.ClusterURI()))
@@ -3014,11 +3017,12 @@ func TestClientSideEncryptionProse(t *testing.T) {
 
 		ceo := options.ClientEncryption().
 			SetKeyVaultNamespace("keyvault.datakeys").
-			SetKmsProviders(fullKmsProvidersMap)
+			SetKmsProviders(fullKmsProvidersMap).
+			SetTLSConfig(map[string]*tls.Config{"aws": tlsCfg})
 		clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
 		require.NoError(mt, err, "error on NewClientEncryption: %v", err)
 
-		err = setFailPoint("http", 1)
+		err = setFailPoint("network", 1)
 		require.NoError(mt, err, "mock server error: %v", err)
 
 		dkOpts := options.DataKey().SetMasterKey(
@@ -3032,7 +3036,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		keyID, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
 		require.NoError(mt, err, "error in CreateDataKey: %v", err)
 
-		err = setFailPoint("http", 1)
+		err = setFailPoint("network", 1)
 		require.NoError(mt, err, "mock server error: %v", err)
 
 		testVal := bson.RawValue{Type: bson.TypeInt32, Value: bsoncore.AppendInt32(nil, 123)}
diff --git a/x/mongo/driver/crypt.go b/x/mongo/driver/crypt.go
index 4368fd125d..5a94d298cc 100644
--- a/x/mongo/driver/crypt.go
+++ b/x/mongo/driver/crypt.go
@@ -9,9 +9,7 @@ package driver
 import (
 	"context"
 	"crypto/tls"
-	"errors"
 	"fmt"
-	"io"
 	"strings"
 	"time"
 
@@ -399,7 +397,10 @@ 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) {
+		if err != nil {
+			if kmsCtx.Fail() {
+				err = nil
+			}
 			return err
 		}
 
diff --git a/x/mongo/driver/mongocrypt/mongocrypt.go b/x/mongo/driver/mongocrypt/mongocrypt.go
index 5f34f5cd71..7f7c3e8fc9 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt.go
+++ b/x/mongo/driver/mongocrypt/mongocrypt.go
@@ -53,6 +53,7 @@ func NewMongoCrypt(opts *options.MongoCryptOptions) (*MongoCrypt, error) {
 	if wrapped == nil {
 		return nil, errors.New("could not create new mongocrypt object")
 	}
+	C.mongocrypt_setopt_retry_kms(wrapped, true)
 	httpClient := opts.HTTPClient
 	if httpClient == nil {
 		httpClient = httputil.DefaultHTTPClient
@@ -85,7 +86,7 @@ func NewMongoCrypt(opts *options.MongoCryptOptions) (*MongoCrypt, error) {
 	}
 
 	if opts.BypassQueryAnalysis {
-		C.mongocrypt_setopt_bypass_query_analysis(wrapped)
+		C.mongocrypt_setopt_bypass_query_analysis(crypt.wrapped)
 	}
 
 	// If loading the crypt_shared library isn't disabled, set the default library search path "$SYSTEM"
diff --git a/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go b/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go
index 296a22315c..d3b028f777 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(time.Duration(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 bool(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)
+}

From 701bd46de447dcb896c51b62ab77a43e00cea705 Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Mon, 2 Dec 2024 17:21:51 -0500
Subject: [PATCH 4/9] fix tests

---
 .evergreen/config.yml                         |  13 +-
 .../client_side_encryption_prose_test.go      | 142 ++++++++++++++----
 2 files changed, 114 insertions(+), 41 deletions(-)

diff --git a/.evergreen/config.yml b/.evergreen/config.yml
index d2bf4b101b..9fe2ff5490 100644
--- a/.evergreen/config.yml
+++ b/.evergreen/config.yml
@@ -554,9 +554,6 @@ functions:
         args: [*task-runner, evg-test-kmip]
 
   start-kms-failpoint-server:
-    - command: ec2.assume_role
-      params:
-        role_arn: ${aws_test_secrets_role}
     - command: subprocess.exec
       params:
         working_dir: src/go.mongodb.org/mongo-driver
@@ -578,15 +575,15 @@ functions:
         binary: "bash"
         env: 
           GO_BUILD_TAGS: cse
-        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, 
-          MONGO_GO_DRIVER_COMPRESSOR]
+        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, MONGO_GO_DRIVER_COMPRESSOR]
         args: [*task-runner, setup-test]
     - command: subprocess.exec
       type: test
       params:
         binary: "bash"
         env:
-          KMS_FAILPOINT_SERVERS_RUNNING: "true"
+          KMS_FAILPOINT_CA_FILE: "${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem"
+          KMS_FAILPOINT_SERVER_RUNNING: "true"
         args: [*task-runner, evg-test-retry-kms-requests]
 
   run-fuzz-tests:
@@ -1532,10 +1529,6 @@ tasks:
           SSL: "nossl"
       - func: start-kms-failpoint-server
       - func: run-retry-kms-requests
-        vars:
-          TOPOLOGY: "server"
-          AUTH: "noauth"
-          SSL: "nossl"
 
   - name: "test-serverless"
     tags: ["serverless"]
diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index 3222b79467..29218284cb 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -2983,13 +2983,23 @@ func TestClientSideEncryptionProse(t *testing.T) {
 	})
 
 	mt.RunOpts("24. kms retry tests", noClientOpts, func(mt *mtest.T) {
-		kmsTlsTestcase := os.Getenv("KMS_FAILPOINT_SERVERS_RUNNING")
+		kmsTlsTestcase := os.Getenv("KMS_FAILPOINT_SERVER_RUNNING")
 		if kmsTlsTestcase == "" {
-			mt.Skipf("Skipping test as KMS_FAILPOINT_SERVERS_RUNNING is not set")
+			mt.Skipf("Skipping test as KMS_FAILPOINT_SERVER_RUNNING is not set")
 		}
 
-		tlsCfg := &tls.Config{
-			InsecureSkipVerify: true,
+		mt.Parallel()
+
+		var tlsCfg *tls.Config
+		if tlsCAFile := os.Getenv("KMS_FAILPOINT_CA_FILE"); tlsCAFile == "" {
+			require.Fail(mt, "failed to load CA file")
+		} else {
+			var err error
+			clientAndCATlsMap := map[string]interface{}{
+				"tlsCAFile": tlsCAFile,
+			}
+			tlsCfg, err = options.BuildTLSConfig(clientAndCATlsMap)
+			require.Nil(mt, err, "BuildTLSConfig error: %v", err)
 		}
 
 		setFailPoint := func(failure string, count int) error {
@@ -3012,39 +3022,109 @@ func TestClientSideEncryptionProse(t *testing.T) {
 			return res.Body.Close()
 		}
 
-		keyVaultClient, err := mongo.Connect(options.Client().ApplyURI(mtest.ClusterURI()))
-		require.NoError(mt, err, "error on Connect: %v", err)
-
-		ceo := options.ClientEncryption().
-			SetKeyVaultNamespace("keyvault.datakeys").
-			SetKmsProviders(fullKmsProvidersMap).
-			SetTLSConfig(map[string]*tls.Config{"aws": tlsCfg})
-		clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
-		require.NoError(mt, err, "error on NewClientEncryption: %v", err)
-
-		err = setFailPoint("network", 1)
-		require.NoError(mt, err, "mock server error: %v", err)
+		kmsProviders := map[string]map[string]interface{}{
+			"aws": {
+				"accessKeyId":     awsAccessKeyID,
+				"secretAccessKey": awsSecretAccessKey,
+			},
+			"azure": {
+				"tenantId":                 azureTenantID,
+				"clientId":                 azureClientID,
+				"clientSecret":             azureClientSecret,
+				"identityPlatformEndpoint": "127.0.0.1:9003",
+			},
+			"gcp": {
+				"email":      gcpEmail,
+				"privateKey": gcpPrivateKey,
+				"endpoint":   "127.0.0.1:9003",
+			},
+		}
 
-		dkOpts := options.DataKey().SetMasterKey(
-			bson.D{
+		dataKeys := []struct {
+			provider  string
+			masterKey interface{}
+		}{
+			{"aws", bson.D{
 				{"region", "foo"},
 				{"key", "bar"},
 				{"endpoint", "127.0.0.1:9003"},
-			},
-		)
-		var keyID bson.Binary
-		keyID, err = clientEncryption.CreateDataKey(context.Background(), "aws", dkOpts)
-		require.NoError(mt, err, "error in CreateDataKey: %v", err)
+			}},
+			{"azure", bson.D{
+				{"keyVaultEndpoint", "127.0.0.1:9003"},
+				{"keyName", "foo"},
+			}},
+			{"gcp", bson.D{
+				{"projectId", "foo"},
+				{"location", "bar"},
+				{"keyRing", "baz"},
+				{"keyName", "qux"},
+				{"endpoint", "127.0.0.1:9003"},
+			}},
+		}
 
-		err = setFailPoint("network", 1)
-		require.NoError(mt, err, "mock server error: %v", err)
+		testCases := []struct {
+			name    string
+			failure string
+		}{
+			{"Case 1: createDataKey and encrypt with TCP retry", "network"},
+			{"Case 2: createDataKey and encrypt with HTTP retry", "http"},
+		}
 
-		testVal := bson.RawValue{Type: bson.TypeInt32, Value: bsoncore.AppendInt32(nil, 123)}
-		eo := options.Encrypt().
-			SetKeyID(keyID).
-			SetAlgorithm("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic")
-		_, err = clientEncryption.Encrypt(context.Background(), testVal, eo)
-		assert.NoError(mt, err, "error in Encrypt: %v", err)
+		for _, tc := range testCases {
+			for _, dataKey := range dataKeys {
+				mt.Run(fmt.Sprintf("%s_%s", tc.name, dataKey.provider), func(mt *mtest.T) {
+					keyVaultClient, err := mongo.Connect(options.Client().ApplyURI(mtest.ClusterURI()))
+					require.NoError(mt, err, "error on Connect: %v", err)
+
+					ceo := options.ClientEncryption().
+						SetKeyVaultNamespace(kvNamespace).
+						SetKmsProviders(kmsProviders).
+						SetTLSConfig(map[string]*tls.Config{dataKey.provider: tlsCfg})
+					clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
+					require.NoError(mt, err, "error on NewClientEncryption: %v", err)
+
+					err = setFailPoint(tc.failure, 1)
+					require.NoError(mt, err, "mock server error: %v", err)
+
+					dkOpts := options.DataKey().SetMasterKey(dataKey.masterKey)
+					var keyID bson.Binary
+					keyID, err = clientEncryption.CreateDataKey(context.Background(), dataKey.provider, dkOpts)
+					require.NoError(mt, err, "error in CreateDataKey: %v", err)
+
+					err = setFailPoint(tc.failure, 1)
+					require.NoError(mt, err, "mock server error: %v", err)
+
+					testVal := bson.RawValue{Type: bson.TypeInt32, Value: bsoncore.AppendInt32(nil, 123)}
+					eo := options.Encrypt().
+						SetKeyID(keyID).
+						SetAlgorithm("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic")
+					_, err = clientEncryption.Encrypt(context.Background(), testVal, eo)
+					assert.NoError(mt, err, "error in Encrypt: %v", err)
+				})
+			}
+		}
+
+		for _, dataKey := range dataKeys {
+			mt.Run(fmt.Sprintf("Case 3: createDataKey fails after too many retries_%s", dataKey.provider), func(mt *mtest.T) {
+				keyVaultClient, err := mongo.Connect(options.Client().ApplyURI(mtest.ClusterURI()))
+				require.NoError(mt, err, "error on Connect: %v", err)
+
+				ceo := options.ClientEncryption().
+					SetKeyVaultNamespace(kvNamespace).
+					SetKmsProviders(fullKmsProvidersMap).
+					SetTLSConfig(map[string]*tls.Config{dataKey.provider: tlsCfg})
+				clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
+				require.NoError(mt, err, "error on NewClientEncryption: %v", err)
+
+				err = setFailPoint("network", 4)
+				require.NoError(mt, err, "mock server error: %v", err)
+
+				dkOpts := options.DataKey().SetMasterKey(dataKey.masterKey)
+				_, err = clientEncryption.CreateDataKey(context.Background(), dataKey.provider, dkOpts)
+				require.Error(mt, err)
+				mt.Logf("CreateDataKey error: %v", err)
+			})
+		}
 	})
 }
 

From 614aba58275726279704fb05c11d48a29dc02b0c Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Thu, 5 Dec 2024 16:15:56 -0500
Subject: [PATCH 5/9] return a wrapped error

---
 .../integration/client_side_encryption_prose_test.go     | 5 ++---
 x/mongo/driver/crypt.go                                  | 5 +----
 x/mongo/driver/mongocrypt/mongocrypt_kms_context.go      | 9 ++++++---
 .../mongocrypt/mongocrypt_kms_context_not_enabled.go     | 4 ++--
 4 files changed, 11 insertions(+), 12 deletions(-)

diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index 29218284cb..71b6aba95d 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -3111,7 +3111,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 
 				ceo := options.ClientEncryption().
 					SetKeyVaultNamespace(kvNamespace).
-					SetKmsProviders(fullKmsProvidersMap).
+					SetKmsProviders(kmsProviders).
 					SetTLSConfig(map[string]*tls.Config{dataKey.provider: tlsCfg})
 				clientEncryption, err := mongo.NewClientEncryption(keyVaultClient, ceo)
 				require.NoError(mt, err, "error on NewClientEncryption: %v", err)
@@ -3121,8 +3121,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 
 				dkOpts := options.DataKey().SetMasterKey(dataKey.masterKey)
 				_, err = clientEncryption.CreateDataKey(context.Background(), dataKey.provider, dkOpts)
-				require.Error(mt, err)
-				mt.Logf("CreateDataKey error: %v", err)
+				require.ErrorContains(mt, err, "KMS request failed after 3 retries due to a network error")
 			})
 		}
 	})
diff --git a/x/mongo/driver/crypt.go b/x/mongo/driver/crypt.go
index 5a94d298cc..7a51b7a4b9 100644
--- a/x/mongo/driver/crypt.go
+++ b/x/mongo/driver/crypt.go
@@ -398,10 +398,7 @@ func (c *crypt) decryptKey(kmsCtx *mongocrypt.KmsContext) error {
 		res := make([]byte, bytesNeeded)
 		bytesRead, err := conn.Read(res)
 		if err != nil {
-			if kmsCtx.Fail() {
-				err = nil
-			}
-			return err
+			return kmsCtx.RequestError()
 		}
 
 		if err = kmsCtx.FeedResponse(res[:bytesRead]); err != nil {
diff --git a/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go b/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go
index d3b028f777..49baa37f2e 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go
+++ b/x/mongo/driver/mongocrypt/mongocrypt_kms_context.go
@@ -78,7 +78,10 @@ func (kc *KmsContext) createErrorFromStatus() error {
 	return errorFromStatus(status)
 }
 
-// Fail returns a boolean indicating whether the failed request may be retried.
-func (kc *KmsContext) Fail() bool {
-	return bool(C.mongocrypt_kms_ctx_fail(kc.wrapped))
+// RequestError returns the source of the network error for KMS requests.
+func (kc *KmsContext) RequestError() error {
+	if bool(C.mongocrypt_kms_ctx_fail(kc.wrapped)) {
+		return nil
+	}
+	return kc.createErrorFromStatus()
 }
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 608d3784f1..7968897648 100644
--- a/x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go
+++ b/x/mongo/driver/mongocrypt/mongocrypt_kms_context_not_enabled.go
@@ -38,7 +38,7 @@ 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 {
+// RequestError returns the source of the network error for KMS requests.
+func (kc *KmsContext) RequestError() error {
 	panic(cseNotSupportedMsg)
 }

From 893731f48d890d0ae7e698920d206421fb392bde Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Tue, 14 Jan 2025 14:59:24 -0500
Subject: [PATCH 6/9] Clarify test logic.

---
 .../client_side_encryption_prose_test.go        | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index 71b6aba95d..b516ddf077 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -2990,17 +2990,14 @@ func TestClientSideEncryptionProse(t *testing.T) {
 
 		mt.Parallel()
 
-		var tlsCfg *tls.Config
-		if tlsCAFile := os.Getenv("KMS_FAILPOINT_CA_FILE"); tlsCAFile == "" {
-			require.Fail(mt, "failed to load CA file")
-		} else {
-			var err error
-			clientAndCATlsMap := map[string]interface{}{
-				"tlsCAFile": tlsCAFile,
-			}
-			tlsCfg, err = options.BuildTLSConfig(clientAndCATlsMap)
-			require.Nil(mt, err, "BuildTLSConfig error: %v", err)
+		tlsCAFile := os.Getenv("KMS_FAILPOINT_CA_FILE")
+		require.NotEmpty(mt, tlsCAFile, "failed to load CA file")
+
+		clientAndCATlsMap := map[string]interface{}{
+			"tlsCAFile": tlsCAFile,
 		}
+		tlsCfg, err := options.BuildTLSConfig(clientAndCATlsMap)
+		require.NoError(mt, err, "BuildTLSConfig error: %v", err)
 
 		setFailPoint := func(failure string, count int) error {
 			url := fmt.Sprintf("https://localhost:9003/set_failpoint/%s", failure)

From f567dbfbac7fa5e2298c0fe7e56e1fdfe6702fea Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Tue, 14 Jan 2025 16:06:00 -0500
Subject: [PATCH 7/9] Clarify test logic.

---
 internal/integration/client_side_encryption_prose_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index b516ddf077..a938df9494 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -2991,7 +2991,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 		mt.Parallel()
 
 		tlsCAFile := os.Getenv("KMS_FAILPOINT_CA_FILE")
-		require.NotEmpty(mt, tlsCAFile, "failed to load CA file")
+		require.NotEqual(mt, tlsCAFile, "", "failed to load CA file")
 
 		clientAndCATlsMap := map[string]interface{}{
 			"tlsCAFile": tlsCAFile,

From 12650cd2fb067ecce3e8610b89825132ac1fcceb Mon Sep 17 00:00:00 2001
From: Qingyang Hu <103950869+qingyang-hu@users.noreply.github.com>
Date: Wed, 15 Jan 2025 15:19:09 -0500
Subject: [PATCH 8/9] Update
 internal/integration/client_side_encryption_prose_test.go

Co-authored-by: Preston Vasquez <prestonvasquez@icloud.com>
---
 internal/integration/client_side_encryption_prose_test.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/internal/integration/client_side_encryption_prose_test.go b/internal/integration/client_side_encryption_prose_test.go
index a938df9494..81011255f4 100644
--- a/internal/integration/client_side_encryption_prose_test.go
+++ b/internal/integration/client_side_encryption_prose_test.go
@@ -3096,7 +3096,7 @@ func TestClientSideEncryptionProse(t *testing.T) {
 						SetKeyID(keyID).
 						SetAlgorithm("AEAD_AES_256_CBC_HMAC_SHA_512-Deterministic")
 					_, err = clientEncryption.Encrypt(context.Background(), testVal, eo)
-					assert.NoError(mt, err, "error in Encrypt: %v", err)
+					require.NoError(mt, err, "error in Encrypt: %v", err)
 				})
 			}
 		}

From f5ccc05b5fcb3ffa01f1d97630118aabdc80f1da Mon Sep 17 00:00:00 2001
From: Qingyang Hu <qingyang.hu@mongodb.com>
Date: Fri, 17 Jan 2025 21:00:01 -0500
Subject: [PATCH 9/9] Update evergreen config.

---
 .evergreen/config.yml | 38 +++++++++-----------------------------
 1 file changed, 9 insertions(+), 29 deletions(-)

diff --git a/.evergreen/config.yml b/.evergreen/config.yml
index 9fe2ff5490..bf51417028 100644
--- a/.evergreen/config.yml
+++ b/.evergreen/config.yml
@@ -553,21 +553,6 @@ functions:
           KMS_MOCK_SERVERS_RUNNING: "true"
         args: [*task-runner, evg-test-kmip]
 
-  start-kms-failpoint-server:
-    - command: subprocess.exec
-      params:
-        working_dir: src/go.mongodb.org/mongo-driver
-        binary: bash
-        background: true
-        include_expansions_in_env: ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SESSION_TOKEN", "DRIVERS_TOOLS"]
-        # This cannot use task because it will hang on Windows.
-        args: [etc/setup-encryption.sh]
-    - command: subprocess.exec
-      params:
-        binary: python3
-        background: true
-        args: ["-u", "${DRIVERS_TOOLS}/.evergreen/csfle/kms_failpoint_server.py", "--port", "9003"]
-
   run-retry-kms-requests:
     - command: subprocess.exec
       type: test
@@ -575,7 +560,8 @@ functions:
         binary: "bash"
         env: 
           GO_BUILD_TAGS: cse
-        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, MONGO_GO_DRIVER_COMPRESSOR]
+        include_expansions_in_env: [AUTH, SSL, MONGODB_URI, TOPOLOGY, 
+          MONGO_GO_DRIVER_COMPRESSOR]
         args: [*task-runner, setup-test]
     - command: subprocess.exec
       type: test
@@ -1473,7 +1459,7 @@ tasks:
           SSL: "nossl"
 
   - name: "test-kms-tls-invalid-cert"
-    tags: ["kms-tls"]
+    tags: ["kms-test"]
     commands:
       - func: bootstrap-mongo-orchestration
         vars:
@@ -1489,7 +1475,7 @@ tasks:
           SSL: "nossl"
 
   - name: "test-kms-tls-invalid-hostname"
-    tags: ["kms-tls"]
+    tags: ["kms-test"]
     commands:
       - func: bootstrap-mongo-orchestration
         vars:
@@ -1520,14 +1506,14 @@ tasks:
           SSL: "nossl"
 
   - name: "test-retry-kms-requests"
-    tags: ["retry-kms-requests"]
+    tags: ["kms-test"]
     commands:
       - func: bootstrap-mongo-orchestration
         vars:
           TOPOLOGY: "server"
           AUTH: "noauth"
           SSL: "nossl"
-      - func: start-kms-failpoint-server
+      - func: start-cse-servers
       - func: run-retry-kms-requests
 
   - name: "test-serverless"
@@ -2207,11 +2193,11 @@ buildvariants:
     tasks:
       - name: ".versioned-api"
 
-  - matrix_name: "kms-tls-test"
+  - matrix_name: "kms-test"
     matrix_spec: { version: ["7.0"], os-ssl-40: ["rhel87-64"] }
-    display_name: "KMS TLS ${os-ssl-40}"
+    display_name: "KMS TEST ${os-ssl-40}"
     tasks:
-      - name: ".kms-tls"
+      - name: ".kms-test"
 
   - matrix_name: "load-balancer-test"
     tags: ["pullrequest"]
@@ -2239,12 +2225,6 @@ buildvariants:
     tasks:
       - name: ".kms-kmip"
 
-  - matrix_name: "retry-kms-requests-test"
-    matrix_spec: { version: ["7.0"], os-ssl-40: ["rhel87-64"] }
-    display_name: "Retry KMS Requests ${os-ssl-40}"
-    tasks:
-      - name: ".retry-kms-requests"
-
   - matrix_name: "fuzz-test"
     matrix_spec: { version: ["5.0"], os-ssl-40: ["rhel87-64"] }
     display_name: "Fuzz ${version} ${os-ssl-40}"