Skip to content

KMS Retry support - #1999

Open
rozza wants to merge 11 commits into
mongodb:mainfrom
rozza:JAVA-5391
Open

KMS Retry support#1999
rozza wants to merge 11 commits into
mongodb:mainfrom
rozza:JAVA-5391

Conversation

@rozza

@rozza rozza commented Jun 11, 2026

Copy link
Copy Markdown
Member

Added KMS Retry support.

4 Commits:

  • Update CAPI javadoc and added rules to: ‎mongodb-crypt/AGENTS.md‎
  • Implement KMS Retry API in mongo-crypt
  • Implemented KMS Retries in sync / reactive streams
  • CAPI cstring NPE hardening - found by copilot

JAVA-5391 / JAVA-5772

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

This PR adds end-to-end support for libmongocrypt’s KMS retry behavior across the sync and reactive-streams drivers, and updates the vendored TLS channel to correctly surface close_notify when it arrives in the same network segment as application data (addressing the KMS HTTP response parsing edge case referenced by JAVA-5391 / JAVA-5411).

Changes:

  • Enable and surface libmongocrypt KMS retry primitives (retryable HTTP errors, network failures, backoff sleep) via new JNA bindings and MongoKeyDecryptor APIs.
  • Update sync/reactive KMS I/O loops to use retry-aware feeding (feedWithRetry), budget exhaustion (fail), and pre-request sleep (sleepMicroseconds), including CSOT-aware timeout behavior.
  • Add prose/integration tests and Evergreen wiring for KMS retry scenarios, plus a unit test suite for the TLS close_notify + data coalescing behavior.

Reviewed changes

Copilot reviewed 16 out of 17 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
mongodb-crypt/src/test/resources/kms-reply-429.txt Adds a canned HTTP 429 KMS response used to validate retryable HTTP error handling.
mongodb-crypt/src/test/java/com/mongodb/crypt/capi/MongoCryptTest.java Adds unit tests validating retry-on-HTTP-error, retry-on-network-error, and retry budget exhaustion.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptorImpl.java Implements new retry/backoff APIs by calling newly bound libmongocrypt KMS functions.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptor.java Extends the decryptor interface with sleep/backoff and retry signaling methods.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java Enables libmongocrypt KMS retry behavior during initialization.
mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java Adds new JNA native method declarations for KMS retry and expands/updates binding documentation.
mongodb-crypt/AGENTS.md Documents contribution rules and mapping guidelines for CAPI.java JNA bindings/javadoc.
driver-sync/src/test/functional/com/mongodb/client/ClientSideEncryptionKmsRetryProseTest.java Adds sync-driver entrypoint for the KMS retry prose tests.
driver-sync/src/test/functional/com/mongodb/client/AbstractClientSideEncryptionKmsRetryProseTest.java Implements KMS retry prose tests (network retry, HTTP retry, budget exhaustion, CSOT timing).
driver-sync/src/main/com/mongodb/client/internal/Crypt.java Integrates retry-aware KMS decryption into the sync encryption state machine, including backoff and timeout mapping.
driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/ClientSideEncryptionKmsRetryProseTest.java Adds reactive-streams entrypoint for the KMS retry prose tests.
driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java Adds reactive retry/backoff handling and retry-aware KMS response feeding.
driver-core/src/test/unit/com/mongodb/internal/connection/tlschannel/TlsChannelCloseNotifyTest.java Adds regression tests for TLS channel read semantics when close_notify and data coalesce.
driver-core/src/main/com/mongodb/internal/connection/tlschannel/impl/TlsChannelImpl.java Updates vendored TLS-channel read/unwrap behavior to return bytes produced alongside close_notify and eagerly unwrap buffered records.
.evergreen/run-kms-retry-tests.sh Adds a CI script to run the new sync + reactive KMS retry prose tests.
.evergreen/.evg.yml Wires a new Evergreen function/task/buildvariant to execute the KMS retry test script.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java
Comment thread mongodb-crypt/AGENTS.md Outdated
@rozza
rozza force-pushed the JAVA-5391 branch 3 times, most recently from d4e6fbf to 7cad542 Compare June 11, 2026 14:06
@rozza
rozza requested a review from Copilot June 11, 2026 14:11
@rozza rozza changed the title Java 5391 KMS Retry support Jun 11, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 19 out of 20 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

mongodb-crypt/src/test/resources/kms-reply-429.txt:7

  • Content-Length: 32 implies the HTTP body is exactly the JSON string, but this resource currently ends with an extra blank line. getHttpResourceAsByteBuffer joins lines with \r\n, so the trailing empty line appends an extra CRLF after the body, making the response longer than the declared Content-Length and potentially breaking/parsing the retry logic in a non-representative way.

Remove the trailing blank line so the file ends immediately after the JSON body (matching how kms-reply.txt is formatted).

Content-Length: 32

{"__type":"ThrottlingException"}

@rozza
rozza force-pushed the JAVA-5391 branch 3 times, most recently from 165faec to cac6f2a Compare June 15, 2026 11:03
@rozza
rozza requested a review from Copilot June 15, 2026 11:24

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

@rozza
rozza force-pushed the JAVA-5391 branch 2 times, most recently from 5ba252a to a285ac3 Compare June 16, 2026 09:29
@rozza
rozza requested a review from Copilot June 16, 2026 13:41

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 18 out of 18 changed files in this pull request and generated 2 comments.

Comment thread mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java
@rozza

rozza commented Jun 17, 2026

Copy link
Copy Markdown
Member Author

Errors in evergreen are related to: DRIVERS-3547

@rozza
rozza force-pushed the JAVA-5391 branch 2 times, most recently from 2dea3ed to f985b18 Compare June 17, 2026 14:11
Expose libmongocrypt's KMS retry primitives so transient KMS failures
can be retried with library-supplied backoff delays.

- Add CAPI bindings: mongocrypt_setopt_retry_kms,
  mongocrypt_kms_ctx_usleep, mongocrypt_kms_ctx_feed_with_retry and
  mongocrypt_kms_ctx_fail
- Expose them via sleepMicroseconds(), feedWithRetry() and fail() on
  MongoKeyDecryptor, and enable retry_kms unconditionally in
  MongoCryptImpl
- HTTP 429 responses are signalled by feedWithRetry and network errors
  by fail(); in both cases the attempt ends early and libmongocrypt
  re-presents the KMS context for retry with the next backoff delay
- Add mongodb-crypt unit tests for the new bindings and a 429 KMS
  reply fixture

JAVA-5391
Wire libmongocrypt's KMS retry behaviour through the sync and reactive
drivers. Transient KMS failures are retried with backoff delays
supplied by libmongocrypt; retry is enabled unconditionally.

- The sync Crypt sleeps the backoff before each attempt, surfaces
  MongoOperationTimeoutException only once the CSOT deadline has
  actually expired (a fixed-timeout connect or handshake stall with
  budget remaining stays retryable) and treats unexpected EOF as a
  retryable network error
- The reactive KeyManagementService applies the backoff via
  Mono.delay, checks the CSOT deadline before consuming retry budget,
  and releases buffers and closes streams on all completion paths;
  the shared backoff budget check lives in MongoCryptHelper
- Add the spec Section 24 KMS retry prose tests (sync and reactive)
  plus a CSOT-during-retry test; each test resets the failpoint
  server and control requests send Connection: close so the
  single-threaded server cannot stall KMS connections or leak armed
  failpoints between tests and suites
- Add a sync unit test covering the CSOT classification of KMS network
  failures, and an Evergreen task running the KMS retry suites

JAVA-5391

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated 1 comment.

Comment thread mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java
@rozza
rozza marked this pull request as ready for review June 18, 2026 13:36
@rozza
rozza requested a review from a team as a code owner June 18, 2026 13:36
@rozza
rozza requested a review from vbabanin June 18, 2026 13:36
Native functions declared to return a cstring (e.g. mongocrypt_status_message)
may return a NULL pointer. By default JNA marshals that to a null cstring
reference, so call sites that immediately call toString() would throw a
NullPointerException.

Override cstring.fromNative to return a cstring with a null pointer rather than
a null reference, and have toString normalize a null pointer to an empty string.
This makes the binding's documented NULL contract safe for existing callers
without changing each call site.
@stIncMale
stIncMale requested review from stIncMale and removed request for vbabanin July 16, 2026 00:27

@stIncMale stIncMale left a comment

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.

This is a partial review. I will come back to reviewing this PR once we are out of the woods with the backpressure project.

The previous PR was closed and replaced with this one. Let's not do the same to this PR (also, let's not force push to the PR branch), as that generally complicates the review process.

Review status of individual files as of 38a268fd7abd774a246a68023c5467ade21ce86c (✓ means the file was reviewed, ✗ means it wasn't):

  • .evergreen/.evg.yml
  • .evergreen/run-kms-retry-tests.sh
  • driver-core/src/main/com/mongodb/internal/capi/MongoCryptHelper.java
  • driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementService.java
  • driver-reactive-streams/src/test/functional/com/mongodb/reactivestreams/client/ClientSideEncryptionKmsRetryProseTest.java
  • driver-reactive-streams/src/test/unit/com/mongodb/reactivestreams/client/internal/crypt/KeyManagementServiceKmsRetryTest.java
  • driver-sync/src/main/com/mongodb/client/internal/Crypt.java
  • driver-sync/src/test/functional/com/mongodb/client/AbstractClientSideEncryptionKmsRetryProseTest.java
  • driver-sync/src/test/functional/com/mongodb/client/ClientSideEncryptionKmsRetryProseTest.java
  • driver-sync/src/test/unit/com/mongodb/client/internal/CryptKmsRetryTest.java
  • mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java - I did not review the unrelated API documentation changes, and do not plan to review them.
  • mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoCryptImpl.java
  • mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptor.java
  • mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/MongoKeyDecryptorImpl.java
  • mongodb-crypt/src/test/java/com/mongodb/crypt/capi/MongoCryptTest.java
  • mongodb-crypt/src/test/java/com/mongodb/internal/crypt/capi/CstringTest.java
  • mongodb-crypt/src/test/resources/kms-reply-429.txt
  • mongodb-crypt/AGENTS.md

}

// always enabled; backoff and budget management are the driver's responsibility
configure(() -> mongocrypt_setopt_retry_kms(wrapped, true));

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.

Is this comment useful? For me it brings only confusion.

  • "always enabled" - that is pretty obvious, as mongocrypt_setopt_retry_kms is called unconditionally. What useful information does this part add?
  • "backoff and budget management are the driver's responsibility" - as far as I understand, whether to retry and how long to sleep before retrying is managed by lybmongocrypt, not the driver, while the comment seem to imply the opposite.

@rozza rozza Jul 27, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I can see your confusion - its the responsibility of the driver to sleep. Removed.

Comment on lines +48 to +53
// By default JNA returns a null cstring reference when a native function returns a NULL
// pointer (e.g. mongocrypt_status_message for a status with no message). Override fromNative
// to return a cstring with a null pointer instead, so callers never receive a null reference
// and toString below can normalize it to an empty string.
@Override
public Object fromNative(final Object nativeValue, final FromNativeContext context) {

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.

I did not find the JNA behavior mentioned in the comment anywhere in the JNA documentation. Let's add the java-native-access/jna@3fc6bc7 commit link to the comment as the supporting evidence.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I don't think we need to add a commit link the comment is clear enough.
Its all because of the getPointer method returns a null - https://java-native-access.github.io/jna/4.2.1/com/sun/jna/Pointer.html#getPointer-long-

Comment on lines +29 to +34
// A native function declared to return a cstring (e.g. mongocrypt_status_message) may return a
// NULL pointer. By default JNA would marshal that to a null cstring reference, causing callers
// that immediately call toString() to throw a NullPointerException. The fromNative override must
// return a non-null cstring instead.
@Test
void fromNativeReturnsEmptyStringBackedInstanceForNullPointer() {

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 use the documentation comment here, as the comment documents the test.

Comment on lines +29 to +34
// A native function declared to return a cstring (e.g. mongocrypt_status_message) may return a
// NULL pointer. By default JNA would marshal that to a null cstring reference, causing callers
// that immediately call toString() to throw a NullPointerException. The fromNative override must
// return a non-null cstring instead.
@Test
void fromNativeReturnsEmptyStringBackedInstanceForNullPointer() {

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.

This test method name is misleading. fromNative(null, null) does not return an "empty string–backed instance", as can be observed by result.getPointer().getString(0) throwing NullPointerException. That is, the returned cstring is backed by Pointer.NULL, but its toString method returns the empty String.

protected abstract ClientEncryption getClientEncryption(ClientEncryptionSettings settings);

@BeforeEach
public void setUp() {

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.

@BeforeEach, @AfterEach, @ParameterizedTest, @Test methods do not have to be public. Let's make them all package-access.

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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This is such a minor nit for test code - I wonder if its even worth doing and reorganising the commits for.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Done - and added a new rule to .agents/references/testing-guide.md seems the driver is consistently inconsistent here!

It would be good to have these nits automated away via spotless or some other tool in the future.

Comment on lines +335 to +342
keyDecryptor = reenterNeedKms(decryptor);
assertTrue(keyDecryptor.bytesNeeded() > 0);
assertTrue(keyDecryptor.getMessage().remaining() > 0);

// A network error is retryable while budget remains
assertTrue(keyDecryptor.fail());
keyDecryptor = reenterNeedKms(decryptor);

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 move the following two assertTrue invocations

assertTrue(keyDecryptor.bytesNeeded() > 0);
assertTrue(keyDecryptor.getMessage().remaining() > 0);

into reenterNeedKms? This way, we will neither miss them on the second quoted invocation of reenterNeedKms, nor will have to duplicate them there.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Done

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

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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Done

Comment on lines +331 to +337
assertTrue(keyDecryptor.feedWithRetry(getHttpResourceAsByteBuffer("kms-reply-429.txt")));
keyDecryptor.sleepMicroseconds();

// The context stays in NEED_KMS and the key decryptor is re-presented for the retry
keyDecryptor = reenterNeedKms(decryptor);
assertTrue(keyDecryptor.bytesNeeded() > 0);
assertTrue(keyDecryptor.getMessage().remaining() > 0);

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 PR description says that it implements JAVA-5772: Support retrying a KMS in-place (and indeed, the PR uses the MongoKeyDecryptor.feedWithRetry (mongocrypt_kms_ctx_feed_with_retry(mongocrypt_kms_ctx_t, ...)) method from that ticket, yet reenterNeedKms calls MongoCryptContext.nextKeyDecryptor (mongocrypt_ctx_next_kms_ctx(mongocrypt_ctx_t)) after feeding a KMS error response. This surprised me, as in an in-place retry "If a KMS context indicates retry, retry the KMS request and feed the new response to the same KMS context".

com.mongodb.client.internal.Crypt.decryptKeys also does not do an in-place retry (at the moment, I am unsure about com.mongodb.reactivestreams.client.internal.crypt.Crypt.decryptKeys) - it calls MongoCryptContext.nextKeyDecryptor each time it retries.

Interestingly, testKmsRetryExhaustsBudget does reuse the MongoKeyDecryptor (mongocrypt_kms_ctx_t), thus testing an in-place retry of a KMS network error.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Confirmed correct — the retry is in-place at the libmongocrypt level.

The contract is about the native mongocrypt_kms_ctx_t, not the Java wrapper:

  • feedWithRetry setting should_retry (HTTP 429/503) or fail() returning true (network error) marks that same kms_ctx for retry. libmongocrypt keeps it pending rather than done.
  • mongocrypt_ctx_next_kms_ctx then hands back the same kms_ctx.
  • nextKeyDecryptor() wraps it in a new MongoKeyDecryptorImpl, but wrapped is the same native context — so the retried request is fed to the same context, as integrating.md requires.

This is just libmongocrypt's documented sync loop:

next_kms_ctx → send → feed_with_retry / fail → next_kms_ctx → … → kms_done

The binding docs back this up: feed_with_retry says "retry in-place without calling mongocrypt_kms_ctx_fail", and fail says the ctx "may be reused".

Tests confirm it:

  • testKmsRetryOnHttpErrorAndNetworkErrorreenterNeedKms(decryptor) asserts a non-null decryptor with bytesNeeded() > 0 after both a 429 feed and a fail(). If next_kms_ctx didn't re-present retryable contexts, this would fail.
  • testKmsRetryExhaustsBudget — drives repeated in-place fail() retries on the same MongoKeyDecryptor.

Reactive Crypt.decryptKeys uses the same model (recursing via doOnSuccess), so it behaves identically.

The only non-in-place bit is the per-iteration wrapper allocation — a thin JNA wrapper over the same native pointer, no behavioral effect.

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

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I'm going to keep them in for now - again we haven't removed CSOT yet and the process of removing it is not set. Eg. deprecation for a major release before another major release or straight removal as its alpha

rozza and others added 3 commits July 27, 2026 09:31
…r.java

Co-authored-by: Valentin Kovalenko <valentin.male.kovalenko@gmail.com>
Co-authored-by: Valentin Kovalenko <valentin.male.kovalenko@gmail.com>
@rozza

rozza commented Jul 27, 2026

Copy link
Copy Markdown
Member Author

Thanks @stIncMale - I've actioned most of the comments.

Main open point at the moment is the CSOT implementation as its piped through. I'm the tech lead on DRIVERS-3580 and its yet to go through triage. Also the deprecation / removal process hasn't been determined. So its a cheap add at the moment and we may release this improvement before DRIVERS-3580 is implemented.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

mongodb-crypt/src/main/com/mongodb/internal/crypt/capi/CAPI.java:64

  • cstring.toString() only checks pointer == null, but CstringTest (and the Javadoc above) expects it to also treat an explicit Pointer.NULL as an empty string. With the current implementation, pointer.getString(0) will be invoked for Pointer.NULL, which is likely to return null or error and would fail toStringReturnsEmptyStringForExplicitNullPointer(). Consider normalizing NULL pointers in fromNative and checking for an address of 0 in toString().
        public String toString() {
            Pointer pointer = getPointer();
            return pointer == null ? "" : pointer.getString(0);
        }

@rozza

rozza commented Jul 27, 2026

Copy link
Copy Markdown
Member Author

@stIncMale note recent changes in server behaviour are causing some of the CSOT tests to fail. These fail against main as well and are not related to the changes in this PR. See: SERVER-132246

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants