diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c index 32fc6f793..e24865c22 100644 --- a/src/mongocrypt-ctx-rewrap-many-datakey.c +++ b/src/mongocrypt-ctx-rewrap-many-datakey.c @@ -136,6 +136,18 @@ static mongocrypt_kms_ctx_t *_next_kms_ctx_encrypt(mongocrypt_ctx_t *ctx) { mongocrypt_ctx_t *dkctx = NULL; BSON_ASSERT_PARAM(ctx); + /* Check if any need retry */ + { + _mongocrypt_ctx_rmd_datakey_t *it = rmdctx->datakeys; + while (it != NULL) { + _mongocrypt_ctx_datakey_t *dkctx = (_mongocrypt_ctx_datakey_t *)it->dkctx; + if (dkctx->kms.should_retry) { + dkctx->kms.should_retry = false; // Reset retry state. + return &dkctx->kms; + } + it = it->next; + } + } /* No more datakey contexts requiring KMS. */ if (!rmdctx->datakeys_iter) { diff --git a/test/test-mongocrypt-ctx-rewrap-many-datakey.c b/test/test-mongocrypt-ctx-rewrap-many-datakey.c index 69381db0b..e725f152f 100644 --- a/test/test-mongocrypt-ctx-rewrap-many-datakey.c +++ b/test/test-mongocrypt-ctx-rewrap-many-datakey.c @@ -570,6 +570,33 @@ static void _test_rewrap_many_datakey_need_kms_retry(_mongocrypt_tester_t *teste ASSERT(mongocrypt_kms_ctx_bytes_needed(kms) == 0); ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx); ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS); // To encrypt. + mongocrypt_ctx_destroy(ctx); + + /* Clear key cache. */ + mongocrypt_destroy(crypt); + crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT); + + /* Ensure KMS encrypt requests retry for network errors */ + ctx = mongocrypt_ctx_new(crypt); + ASSERT_OK(mongocrypt_ctx_rewrap_many_datakey_init(ctx, filter), ctx); + ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_MONGO_KEYS); + ASSERT_OK(mongocrypt_ctx_mongo_feed(ctx, TEST_FILE("./test/data/rmd/key-document-a.json")), ctx); + ASSERT_OK(mongocrypt_ctx_mongo_done(ctx), ctx); + ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS); // To decrypt. + ASSERT((kms = mongocrypt_ctx_next_kms_ctx(ctx))); + ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/rmd/kms-decrypt-reply-a.txt")), kms); + ASSERT(mongocrypt_kms_ctx_bytes_needed(kms) == 0); + ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx); + ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS); // To encrypt. + ASSERT((kms = mongocrypt_ctx_next_kms_ctx(ctx))); + ASSERT(mongocrypt_kms_ctx_fail(kms)); // Simulate driver-side network failure for an encrypt request. + ASSERT((kms = mongocrypt_ctx_next_kms_ctx(ctx))); // Assert fails. Expected KMS request to retry but did not. + ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/rmd/kms-encrypt-reply-a.txt")), kms); + ASSERT(mongocrypt_kms_ctx_bytes_needed(kms) == 0); + ASSERT_OK(!mongocrypt_ctx_next_kms_ctx(ctx), ctx); + ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx); + ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY); + mongocrypt_ctx_destroy(ctx); mongocrypt_destroy(crypt); }