diff --git a/integrating.md b/integrating.md index 1773e0ce5..230bdfd99 100644 --- a/integrating.md +++ b/integrating.md @@ -248,8 +248,8 @@ Ensure `mongocrypt_setopt_retry_kms` is called on the `mongocrypt_t` to enable r d. Feed the reply back with `mongocrypt_kms_ctx_feed`. Repeat > until `mongocrypt_kms_ctx_bytes_needed` returns 0. - If any step encounters a network error, call `mongocrypt_kms_ctx_fail`. - If `mongocrypt_kms_ctx_fail` returns true, continue to the next KMS context. + If any step encounters a network error or if `mongocrypt_kms_ctx_should_retry` returns true after feeding the reply, call `mongocrypt_kms_ctx_fail`. + If `mongocrypt_kms_ctx_fail` returns true, retry the request by continuing to the next KMS context or by feeding the new response into the same context. If `mongocrypt_kms_ctx_fail` returns false, abort and report an error. Consider wrapping the error reported in `mongocrypt_kms_ctx_status` to include the last network error. 2. When done feeding all replies, call `mongocrypt_ctx_kms_done`. diff --git a/src/mongocrypt-kms-ctx.c b/src/mongocrypt-kms-ctx.c index 6e18df16a..25ffa9f93 100644 --- a/src/mongocrypt-kms-ctx.c +++ b/src/mongocrypt-kms-ctx.c @@ -1120,6 +1120,28 @@ static bool _ctx_done_kmip_decrypt(mongocrypt_kms_ctx_t *kms_ctx) { return ret; } +static bool _is_retryable_req(_kms_request_type_t req_type) { + // Check if request type is retryable. Some requests are non-idempotent and cannot be safely retried. + _kms_request_type_t retryable_types[] = {MONGOCRYPT_KMS_AZURE_OAUTH, + MONGOCRYPT_KMS_GCP_OAUTH, + MONGOCRYPT_KMS_AWS_ENCRYPT, + MONGOCRYPT_KMS_AWS_DECRYPT, + MONGOCRYPT_KMS_AZURE_WRAPKEY, + MONGOCRYPT_KMS_AZURE_UNWRAPKEY, + MONGOCRYPT_KMS_GCP_ENCRYPT, + MONGOCRYPT_KMS_GCP_DECRYPT}; + for (size_t i = 0; i < sizeof(retryable_types) / sizeof(retryable_types[0]); i++) { + if (retryable_types[i] == req_type) { + return true; + } + } + return false; +} + +bool mongocrypt_kms_ctx_should_retry(mongocrypt_kms_ctx_t *kms) { + return kms && kms->should_retry; +} + bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms) { if (!kms) { return false; @@ -1138,23 +1160,7 @@ bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms) { return false; } - // Check if request type is retryable. Some requests are non-idempotent and cannot be safely retried. - _kms_request_type_t retryable_types[] = {MONGOCRYPT_KMS_AZURE_OAUTH, - MONGOCRYPT_KMS_GCP_OAUTH, - MONGOCRYPT_KMS_AWS_ENCRYPT, - MONGOCRYPT_KMS_AWS_DECRYPT, - MONGOCRYPT_KMS_AZURE_WRAPKEY, - MONGOCRYPT_KMS_AZURE_UNWRAPKEY, - MONGOCRYPT_KMS_GCP_ENCRYPT, - MONGOCRYPT_KMS_GCP_DECRYPT}; - bool is_retryable = false; - for (size_t i = 0; i < sizeof(retryable_types) / sizeof(retryable_types[0]); i++) { - if (retryable_types[i] == kms->req_type) { - is_retryable = true; - break; - } - } - if (!is_retryable) { + if (!_is_retryable_req(kms->req_type)) { CLIENT_ERR("KMS request failed due to network error"); return false; } @@ -1178,6 +1184,10 @@ bool mongocrypt_kms_ctx_feed(mongocrypt_kms_ctx_t *kms, mongocrypt_binary_t *byt if (!mongocrypt_status_ok(status)) { return false; } + if (kms->should_retry) { + // This happens when a KMS context is reused in-place + kms->should_retry = false; + } if (!bytes) { CLIENT_ERR("argument 'bytes' is required"); diff --git a/src/mongocrypt.h b/src/mongocrypt.h index 894d0b463..b8275f5b6 100644 --- a/src/mongocrypt.h +++ b/src/mongocrypt.h @@ -1180,7 +1180,8 @@ MONGOCRYPT_EXPORT bool mongocrypt_kms_ctx_feed(mongocrypt_kms_ctx_t *kms, mongocrypt_binary_t *bytes); /** - * Indicate a network-level failure. + * Indicate a failure. Discards all data fed to this KMS context with @ref mongocrypt_kms_ctx_feed. + * The @ref mongocrypt_kms_ctx_t may be reused. * * @param[in] kms The @ref mongocrypt_kms_ctx_t. * @return A boolean indicating whether the failed request may be retried. @@ -1188,6 +1189,15 @@ bool mongocrypt_kms_ctx_feed(mongocrypt_kms_ctx_t *kms, mongocrypt_binary_t *byt MONGOCRYPT_EXPORT bool mongocrypt_kms_ctx_fail(mongocrypt_kms_ctx_t *kms); +/** + * Indicate if a KMS context is completed but should be retried. + * + * @param[in] kms The @ref mongocrypt_kms_ctx_t. + * @return A boolean indicating whether the failed request should be retried. + */ +MONGOCRYPT_EXPORT +bool mongocrypt_kms_ctx_should_retry(mongocrypt_kms_ctx_t *kms); + /** * Get the status associated with a @ref mongocrypt_kms_ctx_t object. * diff --git a/test/data/kms-aws/encrypt-response-partial.txt b/test/data/kms-aws/encrypt-response-partial.txt new file mode 100644 index 000000000..65eaa01fe --- /dev/null +++ b/test/data/kms-aws/encrypt-response-partial.txt @@ -0,0 +1,7 @@ +HTTP/1.1 200 OK +x-amzn-RequestId: deeb35e5-4ecb-4bf1-9af5-84a54ff0af0e +Content-Type: application/x-amz-json-1.1 +Content-Length: 446 +Connection: close + +{"KeyId": "arn:aws:k \ No newline at end of file diff --git a/test/test-mongocrypt-datakey.c b/test/test-mongocrypt-datakey.c index 1d48940cb..e4d20b918 100644 --- a/test/test-mongocrypt-datakey.c +++ b/test/test-mongocrypt-datakey.c @@ -427,6 +427,33 @@ static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) { mongocrypt_destroy(crypt); } + // Test that an HTTP error is retried in-place. + { + mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT); + mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt); + ASSERT_OK( + mongocrypt_ctx_setopt_key_encryption_key(ctx, + TEST_BSON("{'provider': 'aws', 'key': 'foo', 'region': 'bar'}")), + ctx); + ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx); + ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS); + mongocrypt_kms_ctx_t *kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx); + ASSERT_OK(kms_ctx, ctx); + // Expect no sleep is requested before any error. + ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), ==, 0); + // Feed a retryable HTTP error. + ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt")), kms_ctx); + // In-place retry is indicated. + ASSERT(mongocrypt_kms_ctx_should_retry(kms_ctx)); + ASSERT(mongocrypt_kms_ctx_fail(kms_ctx)); + // Feed a successful response. + ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx); + ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx); + _mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE); + mongocrypt_ctx_destroy(ctx); + mongocrypt_destroy(crypt); + } + // Test that a network error is retried. { mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT); @@ -454,6 +481,30 @@ static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) { mongocrypt_destroy(crypt); } + // Test that a network error is retried in-place. + { + mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT); + mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt); + ASSERT_OK( + mongocrypt_ctx_setopt_key_encryption_key(ctx, + TEST_BSON("{'provider': 'aws', 'key': 'foo', 'region': 'bar'}")), + ctx); + ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx); + ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS); + mongocrypt_kms_ctx_t *kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx); + ASSERT_OK(kms_ctx, ctx); + // Expect no sleep is requested before any error. + ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), ==, 0); + // Mark a network error. + ASSERT_OK(mongocrypt_kms_ctx_fail(kms_ctx), kms_ctx); + // Feed a successful response. + ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx); + ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx); + _mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_DONE); + mongocrypt_ctx_destroy(ctx); + mongocrypt_destroy(crypt); + } + // Test that an oauth request is retried for a network error. { mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);