Apply suggestions from code review

Co-authored-by: Kevin Albertson <[email protected]>
mongodb · Jan 13, 2025 · 0fa5257 · 0fa5257
1 parent 293ced9
commit 0fa5257
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c
@@ -142,6 +142,7 @@ static mongocrypt_kms_ctx_t *_next_kms_ctx_encrypt(mongocrypt_ctx_t *ctx) {
         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;

diff --git a/test/test-mongocrypt-ctx-rewrap-many-datakey.c b/test/test-mongocrypt-ctx-rewrap-many-datakey.c
@@ -590,8 +590,12 @@ static void _test_rewrap_many_datakey_need_kms_retry(_mongocrypt_tester_t *teste
     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((kms = mongocrypt_ctx_next_kms_ctx(ctx))); // Expect KMS request to retry.
+    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);
 }