From 46213b655b3f0f7a4b3eca26feb7efd9d630e8bf Mon Sep 17 00:00:00 2001 From: Adrian Dole Date: Wed, 11 Dec 2024 16:28:03 -0500 Subject: [PATCH 1/3] encrypt retry test --- .../test-mongocrypt-ctx-rewrap-many-datakey.c | 22 +++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/test/test-mongocrypt-ctx-rewrap-many-datakey.c b/test/test-mongocrypt-ctx-rewrap-many-datakey.c index 69381db0b..482c558ab 100644 --- a/test/test-mongocrypt-ctx-rewrap-many-datakey.c +++ b/test/test-mongocrypt-ctx-rewrap-many-datakey.c @@ -570,6 +570,28 @@ 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. + mongocrypt_ctx_destroy(ctx); mongocrypt_destroy(crypt); } From 293ced9c7702dfa134d60fc2e5ae3049bbc220e6 Mon Sep 17 00:00:00 2001 From: Adrian Dole Date: Wed, 11 Dec 2024 16:28:09 -0500 Subject: [PATCH 2/3] encrypt retry fix --- src/mongocrypt-ctx-rewrap-many-datakey.c | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c index 32fc6f793..1b5f7dd45 100644 --- a/src/mongocrypt-ctx-rewrap-many-datakey.c +++ b/src/mongocrypt-ctx-rewrap-many-datakey.c @@ -136,6 +136,17 @@ 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) { + return &dkctx->kms; + } + it = it->next; + } + } /* No more datakey contexts requiring KMS. */ if (!rmdctx->datakeys_iter) { From 0fa52571abe47c4495c5db4e9856937d9ee54e16 Mon Sep 17 00:00:00 2001 From: Adrian Dole Date: Mon, 13 Jan 2025 12:04:43 -0800 Subject: [PATCH 3/3] Apply suggestions from code review Co-authored-by: Kevin Albertson --- src/mongocrypt-ctx-rewrap-many-datakey.c | 1 + test/test-mongocrypt-ctx-rewrap-many-datakey.c | 8 ++++++-- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c index 1b5f7dd45..b3025a7f1 100644 --- 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 index 482c558ab..3ee83662c 100644 --- 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); }