/*
 * Copyright 2019-present MongoDB, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <mongocrypt.h>

#include "kms_message/kms_b64.h"
#include "mongocrypt-crypto-private.h"
#include "mongocrypt-private.h"
#include "test-mongocrypt.h"

static void _init_buffer_with_count(_mongocrypt_buffer_t *out, uint32_t count) {
    out->len = count;
    out->data = bson_malloc0(out->len);
    BSON_ASSERT(out->data);

    out->owned = true;
}

static void _test_random_generator(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    _mongocrypt_buffer_t out;
    mongocrypt_status_t *status;
#define TEST_COUNT 32
    int mid = TEST_COUNT / 2;
    char zero[TEST_COUNT];

    crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);

    /* _mongocrypt_random handles the case where the count size is greater
     * than the buffer by throwing an error. Because of that, no additional tests
     * for this case is needed here. */

    memset(zero, 0, TEST_COUNT);
    status = mongocrypt_status_new();
    _init_buffer_with_count(&out, TEST_COUNT);

    BSON_ASSERT(_mongocrypt_random(crypt->crypto, &out, TEST_COUNT, status));
    BSON_ASSERT(0 != memcmp(zero, out.data, TEST_COUNT)); /* initialized */

    mongocrypt_status_destroy(status);
    _mongocrypt_buffer_cleanup(&out);

    status = mongocrypt_status_new();
    _init_buffer_with_count(&out, TEST_COUNT);

    ASSERT_FAILS_STATUS(_mongocrypt_random(crypt->crypto, &out, mid, status), status, "out should have length 16");

    mongocrypt_status_destroy(status);
    _mongocrypt_buffer_cleanup(&out);
    mongocrypt_destroy(crypt);
}

static void _test_create_data_key_with_provider(_mongocrypt_tester_t *tester,
                                                _mongocrypt_kms_provider_t provider,
                                                bool with_alt_name) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms;
    mongocrypt_binary_t *bin;
    bson_t as_bson;
    bson_iter_t iter;
    _mongocrypt_buffer_t buf;
    int64_t created_date;
    const int64_t current_epoch_time_ms = 1565097975532ll; /* the time this code was written */
    const int64_t one_hundred_years_ms = (int64_t)1000ll * 60ll * 60ll * 24ll * 365ll * 100ll;

    crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
    ctx = mongocrypt_ctx_new(crypt);
    if (provider == MONGOCRYPT_KMS_PROVIDER_AWS) {
        ASSERT_OK(mongocrypt_ctx_setopt_masterkey_aws(ctx, "region", -1, "cmk", -1), ctx);
    } else {
        ASSERT_OK(mongocrypt_ctx_setopt_masterkey_local(ctx), ctx);
    }

    if (with_alt_name) {
        ASSERT_OK(mongocrypt_ctx_setopt_key_alt_name(ctx, TEST_BSON("{'keyAltName': 'b'}")), ctx);
        ASSERT_OK(mongocrypt_ctx_setopt_key_alt_name(ctx, TEST_BSON("{'keyAltName': 'a'}")), ctx);
    }

    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
    if (provider == MONGOCRYPT_KMS_PROVIDER_AWS) {
        BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_NEED_KMS);
        kms = mongocrypt_ctx_next_kms_ctx(ctx);
        BSON_ASSERT(kms);
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms);
        BSON_ASSERT(0 == mongocrypt_kms_ctx_bytes_needed(kms));
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);
    }
    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_READY);
    bin = mongocrypt_binary_new();
    ASSERT_OK(mongocrypt_ctx_finalize(ctx, bin), ctx);
    /* Check the BSON document created. */
    BSON_ASSERT(_mongocrypt_binary_to_bson(bin, &as_bson));
    CRYPT_TRACEF(&crypt->log, "created data key: %s\n", tmp_json(&as_bson));
    /* _id is a UUID */
    BSON_ASSERT(bson_iter_init_find(&iter, &as_bson, "_id"));
    BSON_ASSERT(_mongocrypt_buffer_from_binary_iter(&buf, &iter));
    BSON_ASSERT(buf.subtype == BSON_SUBTYPE_UUID);
    /* keyMaterial is a binary blob of >= KEYMATERIAL_LEN bytes. */
    BSON_ASSERT(bson_iter_init_find(&iter, &as_bson, "keyMaterial"));
    BSON_ASSERT(_mongocrypt_buffer_from_binary_iter(&buf, &iter));
    BSON_ASSERT(buf.subtype == BSON_SUBTYPE_BINARY);
    BSON_ASSERT(buf.len >= MONGOCRYPT_KEY_LEN);
    /* creationDate and updatedDate exist and have the same value. */
    BSON_ASSERT(bson_iter_init_find(&iter, &as_bson, "creationDate"));
    BSON_ASSERT(BSON_ITER_HOLDS_DATE_TIME(&iter));
    created_date = bson_iter_date_time(&iter);
    BSON_ASSERT(created_date > current_epoch_time_ms && created_date < current_epoch_time_ms + one_hundred_years_ms);
    BSON_ASSERT(bson_iter_init_find(&iter, &as_bson, "updateDate"));
    BSON_ASSERT(BSON_ITER_HOLDS_DATE_TIME(&iter));
    BSON_ASSERT(created_date == bson_iter_date_time(&iter));
    if (with_alt_name) {
        BSON_ASSERT(bson_iter_init(&iter, &as_bson));
        BSON_ASSERT(bson_iter_find_descendant(&iter, "keyAltNames.0", &iter));
        BSON_ASSERT(BSON_ITER_HOLDS_UTF8(&iter));
        BSON_ASSERT(0 == strcmp(bson_iter_utf8(&iter, NULL), "a"));
        BSON_ASSERT(bson_iter_init(&iter, &as_bson));
        BSON_ASSERT(bson_iter_find_descendant(&iter, "keyAltNames.1", &iter));
        BSON_ASSERT(BSON_ITER_HOLDS_UTF8(&iter));
        BSON_ASSERT(0 == strcmp(bson_iter_utf8(&iter, NULL), "b"));
        BSON_ASSERT(BSON_ITER_HOLDS_UTF8(&iter));
    } else {
        BSON_ASSERT(!bson_iter_init_find(&iter, &as_bson, "keyAltNames"));
    }

    /* masterKey matches set options. */
    BSON_ASSERT(bson_iter_init(&iter, &as_bson));
    BSON_ASSERT(bson_iter_find_descendant(&iter, "masterKey.provider", &iter));
    BSON_ASSERT(BSON_ITER_HOLDS_UTF8(&iter));
    if (provider == MONGOCRYPT_KMS_PROVIDER_AWS) {
        BSON_ASSERT(0 == strcmp("aws", bson_iter_utf8(&iter, NULL)));
        BSON_ASSERT(bson_iter_init(&iter, &as_bson));
        BSON_ASSERT(bson_iter_find_descendant(&iter, "masterKey.region", &iter));
        BSON_ASSERT(BSON_ITER_HOLDS_UTF8(&iter));
        BSON_ASSERT(0 == strcmp("region", bson_iter_utf8(&iter, NULL)));
        BSON_ASSERT(bson_iter_init(&iter, &as_bson));
        BSON_ASSERT(bson_iter_find_descendant(&iter, "masterKey.key", &iter));
        BSON_ASSERT(BSON_ITER_HOLDS_UTF8(&iter));
        BSON_ASSERT(0 == strcmp("cmk", bson_iter_utf8(&iter, NULL)));
    } else {
        BSON_ASSERT(0 == strcmp("local", bson_iter_utf8(&iter, NULL)));
    }
    mongocrypt_binary_destroy(bin);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

static void _test_datakey_custom_endpoint(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms_ctx;
    mongocrypt_binary_t *bin;
    const char *endpoint;
    bson_t key_bson;
    bson_iter_t iter;

    /* Success. */
    crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_setopt_masterkey_aws(ctx, "region", -1, "cmk", -1), ctx);
    ASSERT_OK(mongocrypt_ctx_setopt_masterkey_aws_endpoint(ctx, "example.com", -1), ctx);
    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_NEED_KMS);
    kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
    BSON_ASSERT(kms_ctx);
    ASSERT_OK(mongocrypt_kms_ctx_endpoint(kms_ctx, &endpoint), ctx);
    BSON_ASSERT(0 == strcmp("example.com:443", endpoint));
    bin = mongocrypt_binary_new();
    ASSERT_OK(mongocrypt_kms_ctx_message(kms_ctx, bin), ctx);
    BSON_ASSERT(NULL != strstr((char *)bin->data, "Host:example.com"));
    ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx);
    BSON_ASSERT(0 == mongocrypt_kms_ctx_bytes_needed(kms_ctx));
    ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);

    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_READY);
    ASSERT_OK(mongocrypt_ctx_finalize(ctx, bin), ctx);
    /* Check the BSON document created. */
    BSON_ASSERT(_mongocrypt_binary_to_bson(bin, &key_bson));
    BSON_ASSERT(bson_iter_init(&iter, &key_bson));
    BSON_ASSERT(bson_iter_find_descendant(&iter, "masterKey.endpoint", &iter));
    BSON_ASSERT(0 == strcmp(bson_iter_utf8(&iter, NULL), "example.com"));

    mongocrypt_binary_destroy(bin);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

static void _test_datakey_kms_per_ctx_credentials(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_kms_ctx_t *kms_ctx;
    mongocrypt_binary_t *bin;
    const char *endpoint;
    bson_t key_bson;
    bson_iter_t iter;

    /* Success. */
    crypt = mongocrypt_new();
    mongocrypt_setopt_use_need_kms_credentials_state(crypt);
    ASSERT_OK(mongocrypt_setopt_kms_providers(crypt, TEST_BSON("{'aws': {}}")), crypt);
    ASSERT_OK(_mongocrypt_init_for_test(crypt), crypt);
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_setopt_masterkey_aws(ctx, "region", -1, "cmk", -1), ctx);
    ASSERT_OK(mongocrypt_ctx_setopt_masterkey_aws_endpoint(ctx, "example.com", -1), ctx);
    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS);
    ASSERT_OK(mongocrypt_ctx_provide_kms_providers(ctx,
                                                   TEST_BSON("{'aws':{'accessKeyId': 'example',"
                                                             "'secretAccessKey': 'example'}}")),
              ctx);
    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_NEED_KMS);
    kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
    BSON_ASSERT(kms_ctx);
    ASSERT_OK(mongocrypt_kms_ctx_endpoint(kms_ctx, &endpoint), ctx);
    BSON_ASSERT(0 == strcmp("example.com:443", endpoint));
    bin = mongocrypt_binary_new();
    ASSERT_OK(mongocrypt_kms_ctx_message(kms_ctx, bin), ctx);
    BSON_ASSERT(NULL != strstr((char *)bin->data, "Host:example.com"));
    ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-aws/encrypt-response.txt")), kms_ctx);
    BSON_ASSERT(0 == mongocrypt_kms_ctx_bytes_needed(kms_ctx));
    ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);

    BSON_ASSERT(mongocrypt_ctx_state(ctx) == MONGOCRYPT_CTX_READY);
    ASSERT_OK(mongocrypt_ctx_finalize(ctx, bin), ctx);
    /* Check the BSON document created. */
    BSON_ASSERT(_mongocrypt_binary_to_bson(bin, &key_bson));
    BSON_ASSERT(bson_iter_init(&iter, &key_bson));
    BSON_ASSERT(bson_iter_find_descendant(&iter, "masterKey.endpoint", &iter));
    BSON_ASSERT(0 == strcmp(bson_iter_utf8(&iter, NULL), "example.com"));

    mongocrypt_binary_destroy(bin);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

/* Test creating a data key with "azure" when only "aws" credentials are needed.
 * Expect the context not to enter the MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS
 * state. */
static void _test_datakey_kms_per_ctx_credentials_not_requested(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;

    crypt = mongocrypt_new();
    mongocrypt_setopt_use_need_kms_credentials_state(crypt);
    ASSERT_OK(mongocrypt_setopt_kms_providers(crypt,
                                              TEST_BSON("{'aws': {}, 'azure': {'tenantId': '', 'clientId': "
                                                        "'', 'clientSecret': '' }}")),
              crypt);
    ASSERT_OK(_mongocrypt_init_for_test(crypt), crypt);
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_setopt_key_encryption_key(ctx,
                                                       TEST_BSON("{'provider': 'azure', 'keyVaultEndpoint': "
                                                                 "'example.vault.azure.net', 'keyName': 'test'}")),
              ctx);
    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);

    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
}

/* Test creating a data key with "local" when "local" credentials are required.
 * Expect the context to enter the MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS
 * state. */
static void _test_datakey_kms_per_ctx_credentials_local(_mongocrypt_tester_t *tester) {
    mongocrypt_t *crypt;
    mongocrypt_ctx_t *ctx;
    mongocrypt_binary_t *bin;
    bson_t key_bson;
    bson_iter_t iter;
    uint8_t local_kek_raw[MONGOCRYPT_KEY_LEN] = {0};
    char *local_kek = kms_message_raw_to_b64(local_kek_raw, sizeof(local_kek_raw));

    crypt = mongocrypt_new();
    mongocrypt_setopt_use_need_kms_credentials_state(crypt);
    ASSERT_OK(mongocrypt_setopt_kms_providers(crypt, TEST_BSON("{'local': {}}")), crypt);
    ASSERT_OK(_mongocrypt_init_for_test(crypt), crypt);
    ctx = mongocrypt_ctx_new(crypt);
    ASSERT_OK(mongocrypt_ctx_setopt_key_encryption_key(ctx, TEST_BSON("{'provider': 'local' }")), ctx);
    ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);
    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS_CREDENTIALS);

    ASSERT_OK(mongocrypt_ctx_provide_kms_providers(ctx,
                                                   TEST_BSON("{'local':{'key': { '$binary': {'base64': '%s', "
                                                             "'subType': '00'}}}}",
                                                             local_kek)),
              ctx);

    ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_READY);
    bin = mongocrypt_binary_new();
    ASSERT_OK(mongocrypt_ctx_finalize(ctx, bin), ctx);
    /* Check the BSON document created. */
    BSON_ASSERT(_mongocrypt_binary_to_bson(bin, &key_bson));
    BSON_ASSERT(bson_iter_init(&iter, &key_bson));
    BSON_ASSERT(bson_iter_find_descendant(&iter, "masterKey.provider", &iter));
    BSON_ASSERT(0 == strcmp(bson_iter_utf8(&iter, NULL), "local"));

    mongocrypt_binary_destroy(bin);
    mongocrypt_ctx_destroy(ctx);
    mongocrypt_destroy(crypt);
    bson_free(local_kek);
}

static void _test_datakey_custom_key_material(_mongocrypt_tester_t *tester) {
    const uint8_t expected_dek[MONGOCRYPT_KEY_LEN] = "0123456789abcdef"
                                                     "0123456789abcdef"
                                                     "0123456789abcdef"
                                                     "0123456789abcdef"
                                                     "0123456789abcdef"
                                                     "0123456789abcdef";

    mongocrypt_t *const crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);

    _mongocrypt_buffer_t encrypted_dek_buf;

    {
        mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);

        /* Generate encrypted DEK using custom key material. */
        {
            mongocrypt_binary_t *const kek = TEST_BSON("{'provider': 'local'}");

            bson_t bson;
            _mongocrypt_buffer_t key_material_buf;
            mongocrypt_binary_t key_material;

            bson_init(&bson);
            ASSERT(BSON_APPEND_BINARY(&bson, "keyMaterial", BSON_SUBTYPE_BINARY, expected_dek, MONGOCRYPT_KEY_LEN));
            _mongocrypt_buffer_from_bson(&key_material_buf, &bson);
            _mongocrypt_buffer_to_binary(&key_material_buf, &key_material);

            ASSERT_OK(mongocrypt_ctx_setopt_key_encryption_key(ctx, kek), ctx);
            ASSERT_OK(mongocrypt_ctx_setopt_key_material(ctx, &key_material), ctx);
            ASSERT_OK(mongocrypt_ctx_datakey_init(ctx), ctx);

            bson_destroy(&bson);
        }

        /* Extract encrypted DEK from datakey. */
        {
            mongocrypt_binary_t *const datakey = mongocrypt_binary_new();

            bson_t bson;
            bson_iter_t iter;
            const uint8_t *binary;
            uint32_t len;
            bson_subtype_t subtype;

            _mongocrypt_tester_run_ctx_to(tester, ctx, MONGOCRYPT_CTX_READY);
            ASSERT_OK(mongocrypt_ctx_finalize(ctx, datakey), ctx);
            ASSERT(_mongocrypt_binary_to_bson(datakey, &bson));

            ASSERT(bson_iter_init_find(&iter, &bson, "keyMaterial"));
            ASSERT(BSON_ITER_HOLDS_BINARY(&iter));
            bson_iter_binary(&iter, &subtype, &len, &binary);
            ASSERT(_mongocrypt_buffer_copy_from_data_and_size(&encrypted_dek_buf, binary, len));

            mongocrypt_binary_destroy(datakey);
        }

        mongocrypt_ctx_destroy(ctx);
    }

    /* Decrypt key material and confirm it matches original DEK. */
    {
        _mongocrypt_buffer_t decrypted_dek_buf;
        mongocrypt_binary_t decrypted_dek;

        mc_kms_creds_t kc;
        ASSERT(_mongocrypt_opts_kms_providers_lookup(&crypt->opts.kms_providers, "local", &kc));

        ASSERT(_mongocrypt_unwrap_key(crypt->crypto,
                                      &kc.value.local.key,
                                      &encrypted_dek_buf,
                                      &decrypted_dek_buf,
                                      crypt->status));

        _mongocrypt_buffer_to_binary(&decrypted_dek_buf, &decrypted_dek);

        ASSERT_CMPBYTES(expected_dek, MONGOCRYPT_KEY_LEN, decrypted_dek.data, decrypted_dek.len);

        _mongocrypt_buffer_cleanup(&decrypted_dek_buf);
    }

    _mongocrypt_buffer_cleanup(&encrypted_dek_buf);
    mongocrypt_destroy(crypt);
}

static void _test_create_datakey_with_retry(_mongocrypt_tester_t *tester) {
    // Test that an HTTP error is retried.
    {
        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);
        // Expect KMS request is returned again for a retry.
        kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OK(kms_ctx, 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 HTTP error is retried in-place.
    {
        mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
        mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
        bool should_retry;
        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_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt"),
                                                     &should_retry),
                  kms_ctx);
        // In-place retry is indicated.
        ASSERT(should_retry);
        // Feed another retryable HTTP error.
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt"),
                                                     &should_retry),
                  kms_ctx);
        // Expect some sleep is requested
        ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), >=, 0);
        // In-place retry is indicated.
        ASSERT(should_retry);
        ASSERT(kms_ctx->attempts == 2);

        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/kms-aws/encrypt-response.txt"),
                                                     &should_retry),
                  kms_ctx);
        ASSERT(!should_retry);
        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);
        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);
        // Expect KMS request is returned again for a retry.
        kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OK(kms_ctx, 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 in-place.
    {
        mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
        mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
        bool should_retry;
        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 partial response
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/kms-aws/encrypt-response-partial.txt"),
                                                     &should_retry),
                  kms_ctx);
        ASSERT(!should_retry);
        // Mark another network error.
        ASSERT_OK(mongocrypt_kms_ctx_fail(kms_ctx), kms_ctx);
        // Expect some sleep is requested
        ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), >=, 0);
        ASSERT(kms_ctx->attempts == 2);
        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/kms-aws/encrypt-response.txt"),
                                                     &should_retry),
                  kms_ctx);
        ASSERT(!should_retry);
        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 subsequent network and HTTP errors can be retried in-place
    {
        mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
        mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
        bool should_retry;
        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 retryable HTTP error.
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt"),
                                                     &should_retry),
                  kms_ctx);
        // In-place retry is indicated.
        ASSERT(should_retry);
        // Expect some sleep is requested
        ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), >=, 0);
        ASSERT(kms_ctx->attempts == 2);
        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/kms-aws/encrypt-response.txt"),
                                                     &should_retry),
                  kms_ctx);
        ASSERT(!should_retry);
        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 subsequent HTTP and network errors can be retried in-place
    {
        mongocrypt_t *crypt = _mongocrypt_tester_mongocrypt(TESTER_MONGOCRYPT_DEFAULT);
        mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
        bool should_retry;
        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_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/rmd/kms-decrypt-reply-429.txt"),
                                                     &should_retry),
                  kms_ctx);
        // In-place retry is indicated.
        ASSERT(should_retry);
        // Mark a network error.
        ASSERT_OK(mongocrypt_kms_ctx_fail(kms_ctx), kms_ctx);
        // Expect some sleep is requested
        ASSERT_CMPINT64(mongocrypt_kms_ctx_usleep(kms_ctx), >=, 0);
        ASSERT(kms_ctx->attempts == 2);
        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed_with_retry(kms_ctx,
                                                     TEST_FILE("./test/data/kms-aws/encrypt-response.txt"),
                                                     &should_retry),
                  kms_ctx);
        ASSERT(!should_retry);
        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);
        mongocrypt_ctx_t *ctx = mongocrypt_ctx_new(crypt);
        ASSERT_OK(mongocrypt_ctx_setopt_key_encryption_key(
                      ctx,
                      TEST_BSON("{'provider': 'azure', 'keyVaultEndpoint': 'example.com', 'keyName': 'foo' }")),
                  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 an oauth request.
        {
            mongocrypt_binary_t *bin = mongocrypt_binary_new();
            ASSERT_OK(mongocrypt_kms_ctx_message(kms_ctx, bin), kms_ctx);
            const char *str = (const char *)mongocrypt_binary_data(bin);
            ASSERT_STRCONTAINS(str, "oauth2");
            mongocrypt_binary_destroy(bin);
        }
        // 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);
        // Expect KMS request is returned again for a retry.
        kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OK(kms_ctx, ctx);
        // Expect an oauth request.
        {
            mongocrypt_binary_t *bin = mongocrypt_binary_new();
            ASSERT_OK(mongocrypt_kms_ctx_message(kms_ctx, bin), kms_ctx);
            const char *str = (const char *)mongocrypt_binary_data(bin);
            ASSERT_STRCONTAINS(str, "oauth2");
            mongocrypt_binary_destroy(bin);
        }
        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-azure/oauth-response.txt")), kms_ctx);
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);

        // Expect KMS request for encrypt.
        ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
        kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OK(kms_ctx, ctx);
        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-azure/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 an HTTP error.
    {
        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': 'azure', 'keyVaultEndpoint': 'example.com', 'keyName': 'foo' }")),
                  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 an oauth request.
        {
            mongocrypt_binary_t *bin = mongocrypt_binary_new();
            ASSERT_OK(mongocrypt_kms_ctx_message(kms_ctx, bin), kms_ctx);
            const char *str = (const char *)mongocrypt_binary_data(bin);
            ASSERT_STRCONTAINS(str, "oauth2");
            mongocrypt_binary_destroy(bin);
        }
        // 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);
        // Expect KMS request is returned again for a retry.
        kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OK(kms_ctx, ctx);
        // Expect an oauth request.
        {
            mongocrypt_binary_t *bin = mongocrypt_binary_new();
            ASSERT_OK(mongocrypt_kms_ctx_message(kms_ctx, bin), kms_ctx);
            const char *str = (const char *)mongocrypt_binary_data(bin);
            ASSERT_STRCONTAINS(str, "oauth2");
            mongocrypt_binary_destroy(bin);
        }

        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-azure/oauth-response.txt")), kms_ctx);
        ASSERT_OK(mongocrypt_ctx_kms_done(ctx), ctx);

        // Expect KMS request for encrypt.
        ASSERT_STATE_EQUAL(mongocrypt_ctx_state(ctx), MONGOCRYPT_CTX_NEED_KMS);
        kms_ctx = mongocrypt_ctx_next_kms_ctx(ctx);
        ASSERT_OK(kms_ctx, ctx);
        // Feed a successful response.
        ASSERT_OK(mongocrypt_kms_ctx_feed(kms_ctx, TEST_FILE("./test/data/kms-azure/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);
    }
}

static void _test_create_data_key(_mongocrypt_tester_t *tester) {
    _test_create_data_key_with_provider(tester, MONGOCRYPT_KMS_PROVIDER_AWS, false /* with_alt_name */);
    _test_create_data_key_with_provider(tester, MONGOCRYPT_KMS_PROVIDER_LOCAL, false /* with_alt_name */);
    _test_create_data_key_with_provider(tester, MONGOCRYPT_KMS_PROVIDER_AWS, true /* with_alt_name */);
    _test_create_data_key_with_provider(tester, MONGOCRYPT_KMS_PROVIDER_LOCAL, true /* with_alt_name */);
}

void _mongocrypt_tester_install_data_key(_mongocrypt_tester_t *tester) {
    INSTALL_TEST(_test_random_generator);
    INSTALL_TEST(_test_create_data_key);
    INSTALL_TEST(_test_datakey_custom_endpoint);
    INSTALL_TEST(_test_datakey_custom_key_material);
    INSTALL_TEST(_test_datakey_kms_per_ctx_credentials);
    INSTALL_TEST(_test_datakey_kms_per_ctx_credentials_not_requested);
    INSTALL_TEST(_test_datakey_kms_per_ctx_credentials_local);
    INSTALL_TEST(_test_create_datakey_with_retry);
}