KMS Retry support

.agents/references/testing-guide.md

@@ -28,6 +28,10 @@ description: Testing frameworks, conventions, and commands for the MongoDB Java
 - Unit tests must not require a running MongoDB instance
 - Descriptive method names: `shouldReturnEmptyListWhenNoDocumentsMatch()` not `test1()`
 - Use `@DisplayName` for human-readable names
+- JUnit 5 does not require `public` visibility: declare test classes and their `@Test`,
+  `@ParameterizedTest`, `@BeforeEach`/`@AfterEach`/`@BeforeAll`/`@AfterAll` methods as
+  package-private (no modifier). Reserve `public`/`protected` for members a subclass or another
+  package genuinely needs to access (e.g. an abstract base test's `protected` hook methods)
 - Clean up test data in `@AfterEach` / `cleanup()` to prevent pollution
 - Import types at the top of the file — do not use inline fully-qualified names
   (`com.mongodb.connection.AsyncCompletionHandler`, `java.util.List`) in signatures or bodies

.evergreen/.evg.yml

@@ -764,6 +764,16 @@ functions:
           set +o xtrace
           MONGODB_URI="${MONGODB_URI}" KMS_TLS_ERROR_TYPE=${KMS_TLS_ERROR_TYPE} .evergreen/run-kms-tls-tests.sh
 
+  "run-kms-retry-test":
+    - command: shell.exec
+      type: "test"
+      params:
+        working_dir: "src"
+        script: |
+          ${PREPARE_SHELL}
+          set +o xtrace
+          MONGODB_URI="${MONGODB_URI}" .evergreen/run-kms-retry-tests.sh
+
   "run-csfle-aws-from-environment-test":
     - command: shell.exec
       type: "test"
@@ -1632,6 +1642,17 @@ tasks:
           AUTH: "noauth"
           SSL: "nossl"
 
+  - name: "test-kms-retry-task"
+    tags: [ "kms-retry" ]
+    commands:
+      - func: "start-mongo-orchestration"
+        vars:
+          TOPOLOGY: "server"
+          AUTH: "noauth"
+          SSL: "nossl"
+      - func: "start-csfle-servers"
+      - func: "run-kms-retry-test"
+
   - name: "test-csfle-aws-from-environment-task"
     tags: [ "csfle-aws-from-environment" ]
     commands:
@@ -2553,6 +2574,12 @@ buildvariants:
     tasks:
       - name: ".kms-tls"
 
+  - matrix_name: "kms-retry-test"
+    matrix_spec: { os: "linux", version: [ "5.0" ], topology: [ "standalone" ] }
+    display_name: "CSFLE KMS Retry"
+    tasks:
+      - name: ".kms-retry"
+
   - matrix_name: "csfle-aws-from-environment-test"
     matrix_spec: { os: "linux", version: [ "5.0" ], topology: [ "standalone" ] }
     display_name: "CSFLE AWS From Environment"

.evergreen/run-kms-retry-tests.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+
+set -o errexit  # Exit the script with error if any of the commands fail
+
+# Supported/used environment variables:
+#       MONGODB_URI                 Set the suggested connection MONGODB_URI (including credentials and topology info)
+
+############################################
+#            Main Program                  #
+############################################
+RELATIVE_DIR_PATH="$(dirname "${BASH_SOURCE:-$0}")"
+. "${RELATIVE_DIR_PATH}/setup-env.bash"
+echo "Running KMS Retry tests"
+
+cp ${JAVA_HOME}/lib/security/cacerts mongo-truststore
+${JAVA_HOME}/bin/keytool -importcert -trustcacerts -file ${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem -keystore mongo-truststore -storepass changeit -storetype JKS -noprompt
+
+export GRADLE_EXTRA_VARS="-Pssl.enabled=true -Pssl.trustStoreType=jks -Pssl.trustStore=`pwd`/mongo-truststore -Pssl.trustStorePassword=changeit"
+
+./gradlew -version
+
+# Disable errexit so both suites run and their exit codes can be captured below.
+set +o errexit
+
+./gradlew --stacktrace --info ${GRADLE_EXTRA_VARS} -Dorg.mongodb.test.uri=${MONGODB_URI} \
+    -Dorg.mongodb.test.kms.retry.run="true" \
+    driver-sync:cleanTest driver-sync:test --tests ClientSideEncryptionKmsRetryProseTest
+first=$?
+echo "sync exit code: $first"
+
+./gradlew --stacktrace --info ${GRADLE_EXTRA_VARS} -Dorg.mongodb.test.uri=${MONGODB_URI} \
+    -Dorg.mongodb.test.kms.retry.run="true" \
+    driver-reactive-streams:cleanTest driver-reactive-streams:test --tests ClientSideEncryptionKmsRetryProseTest
+second=$?
+echo "reactive exit code: $second"
+
+if [ $first -ne 0 ]; then
+   exit $first
+elif [ $second -ne 0 ]; then
+   exit $second
+else
+   exit 0
+fi

driver-core/src/main/com/mongodb/internal/capi/MongoCryptHelper.java

@@ -23,11 +23,14 @@
 import com.mongodb.MongoClientException;
 import com.mongodb.MongoClientSettings;
 import com.mongodb.MongoConfigurationException;
+import com.mongodb.MongoOperationTimeoutException;
 import com.mongodb.client.model.vault.RewrapManyDataKeyOptions;
+import com.mongodb.internal.TimeoutContext;
 import com.mongodb.internal.authentication.AwsCredentialHelper;
 import com.mongodb.internal.authentication.AzureCredentialHelper;
 import com.mongodb.internal.authentication.GcpCredentialHelper;
 import com.mongodb.internal.crypt.capi.MongoCryptOptions;
+import com.mongodb.internal.time.Timeout;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
 import org.bson.BsonDocumentWrapper;
@@ -52,6 +55,24 @@
  */
 public final class MongoCryptHelper {
 
+    public static final String KMS_TIMEOUT_ERROR_MESSAGE = "KMS key decryption exceeded the timeout limit.";
+
+    /**
+     * Throws a {@link MongoOperationTimeoutException} if the operation timeout has expired or the
+     * KMS retry backoff would exceed the remaining operation time.
+     *
+     * @param operationTimeout the operation timeout, or null if none
+     * @param backoffMicros the backoff to sleep before the next KMS attempt, in microseconds
+     */
+    public static void checkKmsRetryBackoff(@Nullable final Timeout operationTimeout, final long backoffMicros) {
+        if (operationTimeout == null) {
+            return;
+        }
+        operationTimeout.shortenBy(backoffMicros, TimeUnit.MICROSECONDS).onExpired(() -> {
+            throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
+        });
+    }
+
     public static MongoCryptOptions createMongoCryptOptions(final ClientEncryptionSettings settings) {
         return createMongoCryptOptions(settings.getKmsProviders(), false, emptyList(), emptyMap(), null, null,
                 settings.getKeyExpiration(TimeUnit.MILLISECONDS));