Skip to content

Commit ce988fd

Browse files
mdb-admdb-ad
committed
implement
1 parent 4e42cc3 commit ce988fd

File tree

4 files changed

+103
-17
lines changed

4 files changed

+103
-17
lines changed

‎src/mongocrypt-kms-ctx.c

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1120,6 +1120,35 @@ static bool _ctx_done_kmip_decrypt(mongocrypt_kms_ctx_t *kms_ctx) {
11201120
return ret;
11211121
}
11221122

1123+
static bool _is_retryable_req(_kms_request_type_t req_type) {
1124+
// Check if request type is retryable. Some requests are non-idempotent and cannot be safely retried.
1125+
_kms_request_type_t retryable_types[] = {MONGOCRYPT_KMS_AZURE_OAUTH,
1126+
MONGOCRYPT_KMS_GCP_OAUTH,
1127+
MONGOCRYPT_KMS_AWS_ENCRYPT,
1128+
MONGOCRYPT_KMS_AWS_DECRYPT,
1129+
MONGOCRYPT_KMS_AZURE_WRAPKEY,
1130+
MONGOCRYPT_KMS_AZURE_UNWRAPKEY,
1131+
MONGOCRYPT_KMS_GCP_ENCRYPT,
1132+
MONGOCRYPT_KMS_GCP_DECRYPT};
1133+
for (size_t i = 0; i < sizeof(retryable_types) / sizeof(retryable_types[0]); i++) {
1134+
if (retryable_types[i] == req_type) {
1135+
return true;
1136+
}
1137+
}
1138+
return false;
1139+
}
1140+
1141+
bool mongocrypt_kms_ctx_should_retry_http(mongocrypt_kms_ctx_t *kms) {
1142+
return kms && kms->should_retry;
1143+
}
1144+
1145+
void mongocrypt_kms_ctx_reset(mongocrypt_kms_ctx_t *kms) {
1146+
if (kms->parser) {
1147+
kms_response_parser_reset(kms->parser);
1148+
}
1149+
kms->should_retry = false;
1150+
}
1151+
11231152
bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms) {
11241153
if (!kms) {
11251154
return false;
@@ -1138,23 +1167,7 @@ bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms) {
11381167
return false;
11391168
}
11401169

1141-
// Check if request type is retryable. Some requests are non-idempotent and cannot be safely retried.
1142-
_kms_request_type_t retryable_types[] = {MONGOCRYPT_KMS_AZURE_OAUTH,
1143-
MONGOCRYPT_KMS_GCP_OAUTH,
1144-
MONGOCRYPT_KMS_AWS_ENCRYPT,
1145-
MONGOCRYPT_KMS_AWS_DECRYPT,
1146-
MONGOCRYPT_KMS_AZURE_WRAPKEY,
1147-
MONGOCRYPT_KMS_AZURE_UNWRAPKEY,
1148-
MONGOCRYPT_KMS_GCP_ENCRYPT,
1149-
MONGOCRYPT_KMS_GCP_DECRYPT};
1150-
bool is_retryable = false;
1151-
for (size_t i = 0; i < sizeof(retryable_types) / sizeof(retryable_types[0]); i++) {
1152-
if (retryable_types[i] == kms->req_type) {
1153-
is_retryable = true;
1154-
break;
1155-
}
1156-
}
1157-
if (!is_retryable) {
1170+
if (!_is_retryable_req(kms->req_type)) {
11581171
CLIENT_ERR("KMS request failed due to network error");
11591172
return false;
11601173
}

‎src/mongocrypt.h

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1188,6 +1188,23 @@ bool mongocrypt_kms_ctx_feed(mongocrypt_kms_ctx_t *kms, mongocrypt_binary_t *byt
11881188
MONGOCRYPT_EXPORT
11891189
bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms);
11901190

1191+
/**
1192+
* Reset a KMS context allowing it to be retried.
1193+
*
1194+
* @param[in] kms The @ref mongocrypt_kms_ctx_t.
1195+
*/
1196+
MONGOCRYPT_EXPORT
1197+
void mongocrypt_kms_ctx_reset(mongocrypt_kms_ctx_t *kms);
1198+
1199+
/**
1200+
* Indicate if a KMS is completed but should be retried due to an HTTP error.
1201+
*
1202+
* @param[in] kms The @ref mongocrypt_kms_ctx_t.
1203+
* @return A boolean indicating whether the failed request is complete but should be retried.
1204+
*/
1205+
MONGOCRYPT_EXPORT
1206+
bool mongocrypt_kms_ctx_should_retry_http(mongocrypt_kms_ctx_t *kms);
1207+
11911208
/**
11921209
* Get the status associated with a @ref mongocrypt_kms_ctx_t object.
11931210
*

‎test/data/kms-aws/encrypt-response-partial.txt

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
HTTP/1.1 200 OK
2+
x-amzn-RequestId: deeb35e5-4ecb-4bf1-9af5-84a54ff0af0e
3+
Content-Type: appli

‎test/test-mongocrypt-datakey.c

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -427,6 +427,33 @@ static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) {
427427
mongocrypt_destroy(crypt);
428428
}
429429

430+
// Test that an HTTP error is retried in-place.
431+
{
432+
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
433+
mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
434+
ASSERT_OK(
435+
mongocrypt_ctx_setopt_key_encryption_key(ctx,
436+
TEST_BSON("{'provider': 'aws', 'key': 'foo', 'region': 'bar'}")),
437+
ctx);
438+
ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
439+
ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
440+
mongocrypt_kms_ctx_t *kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
441+
ASSERT_OK(kms_ctx, ctx);
442+
// Expect no sleep is requested before any error.
443+
ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), ==, 0);
444+
// Feed a retryable HTTP error.
445+
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt")), kms_ctx);
446+
// In-place retry is indicated.
447+
ASSERT(mongocrypt_kms_ctx_should_retry_http(kms_ctx));
448+
mongocrypt_kms_ctx_reset(kms_ctx);
449+
// Feed a successful response.
450+
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx);
451+
ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
452+
_mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE);
453+
mongocrypt_ctx_destroy(ctx);
454+
mongocrypt_destroy(crypt);
455+
}
456+
430457
// Test that a network error is retried.
431458
{
432459
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
@@ -454,6 +481,32 @@ static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) {
454481
mongocrypt_destroy(crypt);
455482
}
456483

484+
// Test that a network error is retried in-place.
485+
{
486+
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
487+
mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
488+
ASSERT_OK(
489+
mongocrypt_ctx_setopt_key_encryption_key(ctx,
490+
TEST_BSON("{'provider': 'aws', 'key': 'foo', 'region': 'bar'}")),
491+
ctx);
492+
ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
493+
ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
494+
mongocrypt_kms_ctx_t *kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
495+
ASSERT_OK(kms_ctx, ctx);
496+
// Expect no sleep is requested before any error.
497+
ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), ==, 0);
498+
// Feed part of a response
499+
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response-partial.txt")), kms_ctx);
500+
// Assume a network error and reset the context.
501+
mongocrypt_kms_ctx_reset(kms_ctx);
502+
// Feed a successful response.
503+
ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx);
504+
ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
505+
_mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE);
506+
mongocrypt_ctx_destroy(ctx);
507+
mongocrypt_destroy(crypt);
508+
}
509+
457510
// Test that an oauth request is retried for a network error.
458511
{
459512
mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);

0 commit comments

Comments
 (0)