From dd2f8308b6c03f82a0d4f49119003219697ad853 Mon Sep 17 00:00:00 2001 From: mdb-ad Date: Wed, 11 Dec 2024 16:28:03 -0500 Subject: [PATCH 1/4] 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 a6f8cdfba933ed7f8622c66d9ad3f3e170f6d648 Mon Sep 17 00:00:00 2001 From: mdb-ad Date: Tue, 4 Mar 2025 22:31:36 -0800 Subject: [PATCH 2/4] 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 6c6fbd57be59bc46b8e7b78f6e75dc0aec9f7a0f Mon Sep 17 00:00:00 2001 From: mdb-ad Date: Tue, 11 Mar 2025 21:08:21 -0700 Subject: [PATCH 3/4] reset retry state --- src/mongocrypt-ctx-rewrap-many-datakey.c | 1 + 1 file changed, 1 insertion(+) diff --git a/src/mongocrypt-ctx-rewrap-many-datakey.c b/src/mongocrypt-ctx-rewrap-many-datakey.c index 1b5f7dd45..e24865c22 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; From 7b8a9f6b14d64bf9df18b339e94e48aa17090f59 Mon Sep 17 00:00:00 2001 From: mdb-ad Date: Mon, 17 Mar 2025 14:07:49 -0700 Subject: [PATCH 4/4] test update --- test/test-mongocrypt-ctx-rewrap-many-datakey.c | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/test-mongocrypt-ctx-rewrap-many-datakey.c b/test/test-mongocrypt-ctx-rewrap-many-datakey.c index 482c558ab..e725f152f 100644 --- a/test/test-mongocrypt-ctx-rewrap-many-datakey.c +++ b/test/test-mongocrypt-ctx-rewrap-many-datakey.c @@ -591,6 +591,11 @@ static void _test_rewrap_many_datakey_need_kms_retry(_mongocrypt_tester_t *teste 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);