Skip to content

KMS Retry support #1999

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rozza wants to merge 8 commits into
base: main
Choose a base branch
from
Open

KMS Retry support #1999

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
df2a759
Updated CAPI.java javadocs to their 1.18.1 mongocrypt.h equivalents a…
rozza Jun 11, 2026
5bff2ae
Add libmongocrypt KMS retry bindings
rozza Jun 17, 2026
68115e8
Retry KMS requests on transient errors
rozza Jun 17, 2026
226b0f7
Harden cstring binding against NULL native returns
rozza Jun 17, 2026
bf0140f
Merge branch 'main' into JAVA-5391
rozza Jul 7, 2026
3f6a2ce
Merge branch 'main' into JAVA-5391
rozza Jul 8, 2026
4a5c864
Merge branch 'main' into JAVA-5391
rozza Jul 9, 2026
38a268f
Merge branch 'main' into JAVA-5391
rozza Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions .evergreen/.evg.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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"

@stIncMale stIncMale Jul 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[informational comment]

This function runs ${DRIVERS_TOOLS}/.evergreen/csfle/setup.sh, which starts the KMS failpoint server required by the prose test. One can see that ${DRIVERS_TOOLS}/.evergreen/csfle/setup.sh indeed starts it at https://github.com/mongodb-labs/drivers-evergreen-tools/blob/18aed4f176dab1ed505c9a2a53466ddf3f4796ec/.evergreen/csfle/start-servers.sh#L70-L74:

echo "Starting Failpoint Server..."
$COMMAND kms_failpoint_server.py --port 9003 --cert-file "$CSFLE_TLS_FAILPOINT_CERT_FILE" --ca-file "$CSFLE_TLS_FAILPOINT_CA_FILE" >failpoint.log 2>&1 &
echo "$!" >>kmip_pids.pid
echo "Starting Failpoint Server...done."
sleep 1

- func: "run-kms-retry-test"

- name: "test-csfle-aws-from-environment-task"
tags: [ "csfle-aws-from-environment" ]
commands:
Expand Down Expand Up @@ -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"
Expand Down
43 changes: 43 additions & 0 deletions .evergreen/run-kms-retry-tests.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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

@stIncMale stIncMale Jul 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem currently occurs in 3 places. It's better to extract it into a variable:

readonly KMS_TLS_CERT_PATH="${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem"

I implemented the proposed change in stIncMale@189a6d6, tested here.

@stIncMale stIncMale Jul 22, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[informational comment]

The file ${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem comes from https://github.com/mongodb-labs/drivers-evergreen-tools/blob/master/.evergreen/x509gen/ca.pem, which is copied by the prepare-resources function in the .evg.yml.


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.ca.path="${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem" \
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.ca.path="${DRIVERS_TOOLS}/.evergreen/x509gen/ca.pem" \
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
28 changes: 28 additions & 0 deletions driver-core/src/main/com/mongodb/internal/capi/MongoCryptHelper.java
View file
Open in desktop

@stIncMale stIncMale Jul 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changes in this file are all CSOT-related, and @dariakp confirmed that we should not do any more CSOT-related work (except for removing CSOT, or making sure the code still compiles and works). Thus, I think we should remove all the CSOT-related changes from the PR, including CSOT-related tests in MongoCryptTest, KeyManagementServiceKmsRetryTest.

Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,12 @@
import com.mongodb.MongoClientSettings;
import com.mongodb.MongoConfigurationException;
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;
Expand All @@ -52,6 +54,32 @@
*/
public final class MongoCryptHelper {

public static final String KMS_TIMEOUT_ERROR_MESSAGE = "KMS key decryption exceeded the timeout limit.";

/**
* Throws a {@code MongoOperationTimeoutException} if the operation timeout has expired or the

@stIncMale stIncMale Jul 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's @link to MongoOperationTimeoutException

* 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.run(TimeUnit.MICROSECONDS,
// infinite timeout: no CSOT budget to enforce; libmongocrypt's retry count is the only limit
() -> { },
remainingMicros -> {
if (remainingMicros < backoffMicros) {
throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
}
},
() -> {
throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
});
Comment on lines +70 to +80

@stIncMale stIncMale Jul 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
operationTimeout.run(TimeUnit.MICROSECONDS,
// infinite timeout: no CSOT budget to enforce; libmongocrypt's retry count is the only limit
() -> { },
remainingMicros -> {
if (remainingMicros < backoffMicros) {
throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
}
},
() -> {
throw TimeoutContext.createMongoTimeoutException(KMS_TIMEOUT_ERROR_MESSAGE);
});
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));
Expand Down
Loading