/*
 * 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 "mongocrypt-key-broker-private.h"
#include "mongocrypt-key-private.h"
#include "test-mongocrypt.h"

/* Given a string, populate a bson_value_t for that string */
static void
_bson_value_from_string (char *string, bson_value_t *value)
{
   bson_t *bson;
   bson_iter_t iter;

   bson = BCON_NEW ("key", string);
   BSON_ASSERT (bson_iter_init_find (&iter, bson, "key"));
   bson_value_copy (bson_iter_value (&iter), value);

   bson_destroy (bson);
}

static void
_key_broker_add_name (_mongocrypt_key_broker_t *kb, char *string)
{
   bson_value_t key_name;

   _bson_value_from_string (string, &key_name);
   ASSERT_OK (_mongocrypt_key_broker_request_name (kb, (void *) &key_name), kb);
   bson_value_destroy (&key_name);
}

/* Create an example 16 byte UUID. Use first_byte to distinguish. */
static void
_gen_uuid (uint8_t first_byte, _mongocrypt_buffer_t *out)
{
   _mongocrypt_tester_fill_buffer (out, 16);
   out->subtype = BSON_SUBTYPE_UUID;
   out->data[0] = first_byte;
}


/* Create an example 16 byte UUID and a corresponding key document with the same
 * _id as the UUID. */
static void
_gen_uuid_and_key_and_altname (_mongocrypt_tester_t *tester,
                               char *altname,
                               uint8_t first_byte,
                               _mongocrypt_buffer_t *id,
                               _mongocrypt_buffer_t *doc)
{
   bson_t as_bson, copied;

   _gen_uuid (first_byte, id);
   BSON_ASSERT (_mongocrypt_binary_to_bson (
      TEST_FILE ("./test/example/key-document.json"), &as_bson));
   bson_init (&copied);
   bson_copy_to_excluding_noinit (
      &as_bson, &copied, "_id", "keyAltNames", NULL);
   BSON_ASSERT (_mongocrypt_buffer_append (id, &copied, "_id", 3));
   if (altname) {
      bson_t child;
      bson_append_array_begin (&copied, "keyAltNames", -1, &child);
      bson_append_utf8 (&child, "0", -1, altname, -1);
      bson_append_array_end (&copied, &child);
   }
   _mongocrypt_buffer_steal_from_bson (doc, &copied);
}

static void
_gen_uuid_and_key (_mongocrypt_tester_t *tester,
                   uint8_t first_byte,
                   _mongocrypt_buffer_t *id,
                   _mongocrypt_buffer_t *doc)
{
   _gen_uuid_and_key_and_altname (tester, NULL, first_byte, id, doc);
}


static uint32_t
_key_broker_num_satisfied (_mongocrypt_key_broker_t *kb)
{
   key_request_t *req;
   uint32_t count;

   count = 0;
   for (req = kb->key_requests; NULL != req; req = req->next) {
      if (req->satisfied) {
         count++;
      }
   }
   return count;
}

static void
_test_key_broker_get_key_filter (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t key_id1, key_id2;
   mongocrypt_binary_t *filter;
   _mongocrypt_key_broker_t key_broker;
   bson_t as_bson;
   bson_t *expected;

   status = mongocrypt_status_new ();
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   _gen_uuid (1, &key_id1);
   _gen_uuid (2, &key_id2);

   /* Multiple different key ids. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id2),
              &key_broker);
   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_filter (&key_broker, filter), &key_broker);
   BSON_ASSERT (_mongocrypt_binary_to_bson (filter, &as_bson));

   expected = BCON_NEW ("$or",
                        "[",
                        "{",
                        "_id",
                        "{",
                        "$in",
                        "[",
                        BCON_BIN (BSON_SUBTYPE_UUID, key_id2.data, key_id2.len),
                        BCON_BIN (BSON_SUBTYPE_UUID, key_id1.data, key_id1.len),
                        "]",
                        "}",
                        "}",
                        "{",
                        "keyAltNames",
                        "{",
                        "$in",
                        "[",
                        "]",
                        "}",
                        "}",
                        "]");

   BSON_ASSERT (0 == bson_compare (expected, &as_bson));
   bson_destroy (expected);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_binary_destroy (filter);

   /* Duplicate key ids. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_filter (&key_broker, filter), &key_broker);
   BSON_ASSERT (_mongocrypt_binary_to_bson (filter, &as_bson));

   expected = BCON_NEW ("$or",
                        "[",
                        "{",
                        "_id",
                        "{",
                        "$in",
                        "[",
                        BCON_BIN (BSON_SUBTYPE_UUID, key_id1.data, key_id1.len),
                        "]",
                        "}",
                        "}",
                        "{",
                        "keyAltNames",
                        "{",
                        "$in",
                        "[",
                        "]",
                        "}",
                        "}",
                        "]");

   BSON_ASSERT (0 == bson_compare (expected, &as_bson));
   bson_destroy (expected);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_binary_destroy (filter);

   /* No key requests made. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   filter = mongocrypt_binary_new ();
   ASSERT_FAILS (_mongocrypt_key_broker_filter (&key_broker, filter),
                 &key_broker,
                 "in wrong state");
   mongocrypt_binary_destroy (filter);
   _mongocrypt_key_broker_cleanup (&key_broker);

   /* Both key ids and keyAltName */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   _key_broker_add_name (&key_broker, "Miriam");
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_filter (&key_broker, filter), &key_broker);
   BSON_ASSERT (_mongocrypt_binary_to_bson (filter, &as_bson));

   expected = BCON_NEW ("$or",
                        "[",
                        "{",
                        "_id",
                        "{",
                        "$in",
                        "[",
                        BCON_BIN (BSON_SUBTYPE_UUID, key_id1.data, key_id1.len),
                        "]",
                        "}",
                        "}",
                        "{",
                        "keyAltNames",
                        "{",
                        "$in",
                        "[",
                        BCON_UTF8 ("Miriam"),
                        "]",
                        "}",
                        "}",
                        "]");

   BSON_ASSERT (0 == bson_compare (expected, &as_bson));
   bson_destroy (expected);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_binary_destroy (filter);

   /* Keys with only keyAltName */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   _key_broker_add_name (&key_broker, "Sharlene");
   _key_broker_add_name (&key_broker, "Emily");
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_filter (&key_broker, filter), &key_broker);
   BSON_ASSERT (_mongocrypt_binary_to_bson (filter, &as_bson));

   expected = BCON_NEW ("$or",
                        "[",
                        "{",
                        "_id",
                        "{",
                        "$in",
                        "[",
                        "]",
                        "}",
                        "}",
                        "{",
                        "keyAltNames",
                        "{",
                        "$in",
                        "[",
                        BCON_UTF8 ("Emily"),
                        BCON_UTF8 ("Sharlene"),
                        "]",
                        "}",
                        "}",
                        "]");

   BSON_ASSERT (0 == bson_compare (expected, &as_bson));
   bson_destroy (expected);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_binary_destroy (filter);

   /* Duplicate alt names */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   _key_broker_add_name (&key_broker, "Jackie");
   _key_broker_add_name (&key_broker, "Jackie");
   _key_broker_add_name (&key_broker, "Jackie");
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_filter (&key_broker, filter), &key_broker);
   BSON_ASSERT (_mongocrypt_binary_to_bson (filter, &as_bson));

   expected = BCON_NEW ("$or",
                        "[",
                        "{",
                        "_id",
                        "{",
                        "$in",
                        "[",
                        "]",
                        "}",
                        "}",
                        "{",
                        "keyAltNames",
                        "{",
                        "$in",
                        "[",
                        BCON_UTF8 ("Jackie"),
                        "]",
                        "}",
                        "}",
                        "]");

   BSON_ASSERT (0 == bson_compare (expected, &as_bson));
   bson_destroy (expected);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_binary_destroy (filter);

   _mongocrypt_buffer_cleanup (&key_id1);
   _mongocrypt_buffer_cleanup (&key_id2);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
}


static void
_test_key_broker_add_key (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t key_id1, key_id2, key_doc1, key_doc2, malformed_buf,
      key_buf_x, key_buf_y, key_doc_names;
   _mongocrypt_buffer_t *id_x;
   _mongocrypt_key_doc_t *key_x;
   bson_t key_bson_x;
   bson_t *malformed;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   _mongocrypt_key_broker_t key_broker;

   status = mongocrypt_status_new ();
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _gen_uuid_and_key (tester, 1, &key_id1, &key_doc1);
   _gen_uuid_and_key (tester, 2, &key_id2, &key_doc2);

   /* Valid key documents. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id2),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_OK (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc2),
      &key_broker);
   ASSERT_OK (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc1),
      &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&key_broker), &key_broker);
   _mongocrypt_key_broker_cleanup (&key_broker);

   /* Valid document with a key name. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   _key_broker_add_name (&key_broker, "Kasey");
   _mongocrypt_buffer_from_binary (
      &key_doc_names,
      TEST_FILE ("./test/data/key-document-with-alt-name.json"));
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (
                 &key_broker, kms_providers, &key_doc_names),
              &key_broker);
   _mongocrypt_key_broker_cleanup (&key_broker);

   /* Malformed key document. */
   malformed = BCON_NEW ("abc", BCON_INT32 (123));
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   _mongocrypt_buffer_from_bson (&malformed_buf, malformed);
   ASSERT_FAILS (_mongocrypt_key_broker_add_doc (
                    &key_broker, kms_providers, &malformed_buf),
                 &key_broker,
                 "unrecognized field");
   _mongocrypt_key_broker_cleanup (&key_broker);
   bson_destroy (malformed);

   /* NULL key document. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   BSON_ASSERT (_mongocrypt_key_broker_request_id (&key_broker, &key_id1));
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_FAILS (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, NULL),
      &key_broker,
      "invalid key");
   _mongocrypt_key_broker_cleanup (&key_broker);

   /* Unmatched key document. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_FAILS (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc2),
      &key_broker,
      "unexpected key returned");
   _mongocrypt_key_broker_cleanup (&key_broker);

   /* Two key documents with the same keyAltName and
      different key ids. In order to test this, we need
      to add a name X and an id Y to the broker, then
      add a document with name X and id Z (succeeds) and
      afterwards add a doc with name X and id Y (fails). */
   key_x = _mongocrypt_key_new ();
   _mongocrypt_key_broker_init (&key_broker, crypt);

   _mongocrypt_buffer_from_binary (
      &key_buf_x, TEST_FILE ("./test/data/key-document-with-alt-name.json"));
   _mongocrypt_buffer_from_binary (
      &key_buf_y,
      TEST_FILE ("./test/data/key-document-with-alt-name-duplicate-id.json"));

   BSON_ASSERT (_mongocrypt_buffer_to_bson (&key_buf_x, &key_bson_x));
   ASSERT_OR_PRINT (_mongocrypt_key_parse_owned (&key_bson_x, key_x, status),
                    status);
   id_x = &key_x->id;

   /* Configure the key broker so it contains:
      - { id : X }
      - { name : "Sharlene" } */
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, id_x),
              &key_broker);
   _key_broker_add_name (&key_broker, "Sharlene");
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);

   /* Add { id : Y, name : "Sharlene" }, should pass. */
   ASSERT_OK (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_buf_y),
      &key_broker);

   /* Add { id : X, name : "Sharlene" }, should fail, it shares an alt name. */
   ASSERT_FAILS (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_buf_x),
      &key_broker,
      "duplicate keyAltNames");

   _mongocrypt_key_broker_cleanup (&key_broker);

   /* Calling done before supplying all keys. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   _mongocrypt_key_broker_cleanup (&key_broker);

   bson_destroy (&key_bson_x);
   _mongocrypt_key_destroy (key_x);
   _mongocrypt_buffer_cleanup (&key_doc_names);
   _mongocrypt_buffer_cleanup (&key_id1);
   _mongocrypt_buffer_cleanup (&key_id2);
   _mongocrypt_buffer_cleanup (&key_doc1);
   _mongocrypt_buffer_cleanup (&key_doc2);
   _mongocrypt_buffer_cleanup (&key_buf_x);
   _mongocrypt_buffer_cleanup (&key_buf_y);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
}

static void
_test_key_broker_add_decrypted_key (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t key_id1, key_id2, key_doc1, key_doc2, key_doc_names,
      key_id_names;
   _mongocrypt_key_broker_t key_broker;
   mongocrypt_kms_ctx_t *kms;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   bson_iter_t iter;
   bson_t key_doc_names_bson;

   status = mongocrypt_status_new ();
   _gen_uuid_and_key (tester, 1, &key_id1, &key_doc1);
   _gen_uuid_and_key (tester, 2, &key_id2, &key_doc2);

   /* Success. With key ids. */
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _mongocrypt_key_broker_init (&key_broker, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id1),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id2),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_OK (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc2),
      &key_broker);
   ASSERT_OK (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc1),
      &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&key_broker), &key_broker);
   kms = _mongocrypt_key_broker_next_kms (&key_broker);
   BSON_ASSERT (kms);
   _mongocrypt_tester_satisfy_kms (tester, kms);
   kms = _mongocrypt_key_broker_next_kms (&key_broker);
   BSON_ASSERT (kms);
   _mongocrypt_tester_satisfy_kms (tester, kms);
   BSON_ASSERT (!_mongocrypt_key_broker_next_kms (&key_broker));
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&key_broker, kms_providers),
              &key_broker);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_destroy (crypt); /* destroy crypt to reset cache. */

   /* Success. With key alt names. */
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _mongocrypt_key_broker_init (&key_broker, crypt);
   _key_broker_add_name (&key_broker, "Sharlene");
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);

   _mongocrypt_buffer_from_binary (
      &key_doc_names,
      TEST_FILE ("./test/data/key-document-with-alt-name.json"));
   ASSERT_OK (_mongocrypt_key_broker_add_doc (
                 &key_broker, kms_providers, &key_doc_names),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&key_broker), &key_broker);
   kms = _mongocrypt_key_broker_next_kms (&key_broker);
   BSON_ASSERT (kms);

   _mongocrypt_tester_satisfy_kms (tester, kms);
   BSON_ASSERT (!_mongocrypt_key_broker_next_kms (&key_broker));
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&key_broker, kms_providers),
              &key_broker);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_destroy (crypt); /* destroy crypt to reset cache. */

   /* With both key ids and key alt names, some referring to the same key */
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _mongocrypt_key_broker_init (&key_broker, crypt);
   BSON_ASSERT (
      _mongocrypt_buffer_to_bson (&key_doc_names, &key_doc_names_bson));
   BSON_ASSERT (bson_iter_init_find (&iter, &key_doc_names_bson, "_id"));
   BSON_ASSERT (_mongocrypt_buffer_from_binary_iter (&key_id_names, &iter));
   BSON_ASSERT (key_id_names.subtype == BSON_SUBTYPE_UUID);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&key_broker, &key_id_names),
              &key_broker);
   _key_broker_add_name (&key_broker, "Sharlene");
   _key_broker_add_name (&key_broker, "Kasey");
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (
                 &key_broker, kms_providers, &key_doc_names),
              &key_broker);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&key_broker), &key_broker);
   kms = _mongocrypt_key_broker_next_kms (&key_broker);
   BSON_ASSERT (kms);
   _mongocrypt_tester_satisfy_kms (tester, kms);
   BSON_ASSERT (!_mongocrypt_key_broker_next_kms (&key_broker));
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&key_broker, kms_providers),
              &key_broker);
   _mongocrypt_key_broker_cleanup (&key_broker);

   bson_destroy (&key_doc_names_bson);
   _mongocrypt_buffer_cleanup (&key_id_names);
   _mongocrypt_buffer_cleanup (&key_id1);
   _mongocrypt_buffer_cleanup (&key_id2);
   _mongocrypt_buffer_cleanup (&key_doc1);
   _mongocrypt_buffer_cleanup (&key_doc2);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
}


static void
_test_key_broker_wrong_subtype (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t key_id, key_doc;
   _mongocrypt_key_broker_t key_broker;

   status = mongocrypt_status_new ();
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   _gen_uuid_and_key (tester, 1, &key_id, &key_doc);

   /* Valid key documents. */
   _mongocrypt_key_broker_init (&key_broker, crypt);
   key_id.subtype = 0;
   ASSERT_FAILS (_mongocrypt_key_broker_request_id (&key_broker, &key_id),
                 &key_broker,
                 "expected UUID");

   _mongocrypt_buffer_cleanup (&key_id);
   _mongocrypt_buffer_cleanup (&key_doc);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
}


static void
_test_key_broker_multi_match (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   _mongocrypt_key_broker_t key_broker;
   status = mongocrypt_status_new ();
   _mongocrypt_buffer_t key_id1, key_id2, key_doc1, key_doc2;

   _gen_uuid_and_key_and_altname (tester, "alt1", 1, &key_id1, &key_doc1);
   _gen_uuid_and_key_and_altname (tester, "alt2", 2, &key_id2, &key_doc2);


   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _mongocrypt_key_broker_init (&key_broker, crypt);

   /* Add two ids and two alt names */
   BSON_ASSERT (_mongocrypt_key_broker_request_id (&key_broker, &key_id1));
   _key_broker_add_name (&key_broker, "alt1");
   BSON_ASSERT (_mongocrypt_key_broker_request_id (&key_broker, &key_id2));
   _key_broker_add_name (&key_broker, "alt2");
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&key_broker), &key_broker);

   /* Should be zero satisfied */
   BSON_ASSERT (0 == _key_broker_num_satisfied (&key_broker));

   /* Add one doc, should satisfy two requests. */
   BSON_ASSERT (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc1));
   BSON_ASSERT (2 == _key_broker_num_satisfied (&key_broker));

   /* Add other doc, should satisfy all. */
   BSON_ASSERT (
      _mongocrypt_key_broker_add_doc (&key_broker, kms_providers, &key_doc2));
   BSON_ASSERT (4 == _key_broker_num_satisfied (&key_broker));

   _mongocrypt_buffer_cleanup (&key_id1);
   _mongocrypt_buffer_cleanup (&key_doc1);
   _mongocrypt_buffer_cleanup (&key_id2);
   _mongocrypt_buffer_cleanup (&key_doc2);
   _mongocrypt_key_broker_cleanup (&key_broker);
   mongocrypt_destroy (crypt);
   mongocrypt_status_destroy (status);
}


/*
<RequestMessage tag="0x420078" type="Structure">
 <RequestHeader tag="0x420077" type="Structure">
  <ProtocolVersion tag="0x420069" type="Structure">
   <ProtocolVersionMajor tag="0x42006a" type="Integer" value="1"/>
   <ProtocolVersionMinor tag="0x42006b" type="Integer" value="0"/>
  </ProtocolVersion>
  <BatchCount tag="0x42000d" type="Integer" value="1"/>
 </RequestHeader>
 <BatchItem tag="0x42000f" type="Structure">
  <Operation tag="0x42005c" type="Enumeration" value="10"/>
  <RequestPayload tag="0x420079" type="Structure">
   <UniqueIdentifier tag="0x420094" type="TextString"
value="ywxrSj5TLjswd1G4oGFJ6hwWgtTsQip0"/>
  </RequestPayload>
 </BatchItem>
</RequestMessage>
*/
static const uint8_t EXPECTED_GET_REQUEST[] = {
   0x42, 0x00, 0x78, 0x01, 0x00, 0x00, 0x00, 0x88, 0x42, 0x00, 0x77, 0x01,
   0x00, 0x00, 0x00, 0x38, 0x42, 0x00, 0x69, 0x01, 0x00, 0x00, 0x00, 0x20,
   0x42, 0x00, 0x6a, 0x02, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01,
   0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x6b, 0x02, 0x00, 0x00, 0x00, 0x04,
   0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x0d, 0x02,
   0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
   0x42, 0x00, 0x0f, 0x01, 0x00, 0x00, 0x00, 0x40, 0x42, 0x00, 0x5c, 0x05,
   0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x00,
   0x42, 0x00, 0x79, 0x01, 0x00, 0x00, 0x00, 0x28, 0x42, 0x00, 0x94, 0x07,
   0x00, 0x00, 0x00, 0x20, 0x79, 0x77, 0x78, 0x72, 0x53, 0x6a, 0x35, 0x54,
   0x4c, 0x6a, 0x73, 0x77, 0x64, 0x31, 0x47, 0x34, 0x6f, 0x47, 0x46, 0x4a,
   0x36, 0x68, 0x77, 0x57, 0x67, 0x74, 0x54, 0x73, 0x51, 0x69, 0x70, 0x30};

/*
<ResponseMessage tag="0x42007b" type="Structure">
 <ResponseHeader tag="0x42007a" type="Structure">
  <ProtocolVersion tag="0x420069" type="Structure">
   <ProtocolVersionMajor tag="0x42006a" type="Integer" value="1"/>
   <ProtocolVersionMinor tag="0x42006b" type="Integer" value="4"/>
  </ProtocolVersion>
  <TimeStamp tag="0x420092" type="DateTime" value="2021-10-03T17:39:52-0500"/>
  <BatchCount tag="0x42000d" type="Integer" value="1"/>
 </ResponseHeader>
 <BatchItem tag="0x42000f" type="Structure">
  <Operation tag="0x42005c" type="Enumeration" value="10"/>
  <ResultStatus tag="0x42007f" type="Enumeration" value="0"/>
  <ResponsePayload tag="0x42007c" type="Structure">
   <ObjectType tag="0x420057" type="Enumeration" value="7"/>
   <UniqueIdentifier tag="0x420094" type="TextString"
value="ywxrSj5TLjswd1G4oGFJ6hwWgtTsQip0"/>
   <SecretData tag="0x420085" type="Structure">
    <SecretDataType tag="0x420086" type="Enumeration" value="2"/>
    <KeyBlock tag="0x420040" type="Structure">
     <KeyFormatType tag="0x420042" type="Enumeration" value="1"/>
     <KeyValue tag="0x420045" type="Structure">
      <KeyMaterial tag="0x420043" type="ByteString"
value="0c2ea7297180f82a984b2fd47d6327ce226f62e9017b91dc6e5d6dfd98747d97e89f17bf0926cfcc0afb24e69b7c00121dda12d0158c4375c31084abf7f2e6044edc2f92802ba3f676d470d2cbc4e33a2a8e53dced7828dd8a35f268437ff141"/>
     </KeyValue>
    </KeyBlock>
   </SecretData>
  </ResponsePayload>
 </BatchItem>
</ResponseMessage>
*/
static const uint8_t SUCCESS_GET_RESPONSE[] = {
   0x42, 0x00, 0x7b, 0x01, 0x00, 0x00, 0x01, 0x58, 0x42, 0x00, 0x7a, 0x01, 0x00,
   0x00, 0x00, 0x48, 0x42, 0x00, 0x69, 0x01, 0x00, 0x00, 0x00, 0x20, 0x42, 0x00,
   0x6a, 0x02, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   0x00, 0x42, 0x00, 0x6b, 0x02, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04,
   0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x92, 0x09, 0x00, 0x00, 0x00, 0x08, 0x00,
   0x00, 0x00, 0x00, 0x61, 0x59, 0xea, 0xe8, 0x42, 0x00, 0x0d, 0x02, 0x00, 0x00,
   0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x0f,
   0x01, 0x00, 0x00, 0x01, 0x00, 0x42, 0x00, 0x5c, 0x05, 0x00, 0x00, 0x00, 0x04,
   0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x7f, 0x05, 0x00,
   0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00,
   0x7c, 0x01, 0x00, 0x00, 0x00, 0xd8, 0x42, 0x00, 0x57, 0x05, 0x00, 0x00, 0x00,
   0x04, 0x00, 0x00, 0x00, 0x07, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x94, 0x07,
   0x00, 0x00, 0x00, 0x20, 0x79, 0x77, 0x78, 0x72, 0x53, 0x6a, 0x35, 0x54, 0x4c,
   0x6a, 0x73, 0x77, 0x64, 0x31, 0x47, 0x34, 0x6f, 0x47, 0x46, 0x4a, 0x36, 0x68,
   0x77, 0x57, 0x67, 0x74, 0x54, 0x73, 0x51, 0x69, 0x70, 0x30, 0x42, 0x00, 0x85,
   0x01, 0x00, 0x00, 0x00, 0x98, 0x42, 0x00, 0x86, 0x05, 0x00, 0x00, 0x00, 0x04,
   0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x40, 0x01, 0x00,
   0x00, 0x00, 0x80, 0x42, 0x00, 0x42, 0x05, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00,
   0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x45, 0x01, 0x00, 0x00, 0x00,
   0x68, 0x42, 0x00, 0x43, 0x08, 0x00, 0x00, 0x00, 0x60, 0x0c, 0x2e, 0xa7, 0x29,
   0x71, 0x80, 0xf8, 0x2a, 0x98, 0x4b, 0x2f, 0xd4, 0x7d, 0x63, 0x27, 0xce, 0x22,
   0x6f, 0x62, 0xe9, 0x01, 0x7b, 0x91, 0xdc, 0x6e, 0x5d, 0x6d, 0xfd, 0x98, 0x74,
   0x7d, 0x97, 0xe8, 0x9f, 0x17, 0xbf, 0x09, 0x26, 0xcf, 0xcc, 0x0a, 0xfb, 0x24,
   0xe6, 0x9b, 0x7c, 0x00, 0x12, 0x1d, 0xda, 0x12, 0xd0, 0x15, 0x8c, 0x43, 0x75,
   0xc3, 0x10, 0x84, 0xab, 0xf7, 0xf2, 0xe6, 0x04, 0x4e, 0xdc, 0x2f, 0x92, 0x80,
   0x2b, 0xa3, 0xf6, 0x76, 0xd4, 0x70, 0xd2, 0xcb, 0xc4, 0xe3, 0x3a, 0x2a, 0x8e,
   0x53, 0xdc, 0xed, 0x78, 0x28, 0xdd, 0x8a, 0x35, 0xf2, 0x68, 0x43, 0x7f, 0xf1,
   0x41};

static const uint8_t EXPECTED_SECRETDATA[] = {
   0x94, 0x82, 0x4f, 0x44, 0xbe, 0xb2, 0x20, 0x73, 0x14, 0xad, 0x8b, 0x36,
   0x38, 0xaf, 0x01, 0x45, 0xa5, 0x13, 0x80, 0x84, 0x44, 0x57, 0xdf, 0xde,
   0x9f, 0xb6, 0x7b, 0xfb, 0xf9, 0x21, 0xf9, 0x00, 0xb2, 0x00, 0x9e, 0x07,
   0xcf, 0x04, 0xc3, 0x5b, 0x9a, 0x98, 0x3b, 0xa9, 0x22, 0x83, 0x3d, 0x7a,
   0x07, 0xc5, 0x90, 0x84, 0xe7, 0x63, 0xf0, 0x47, 0xf0, 0x1a, 0x4b, 0xfe,
   0x03, 0xbc, 0xe3, 0x82, 0x96, 0x95, 0x88, 0xb8, 0x18, 0x63, 0x33, 0x15,
   0x73, 0x95, 0xe2, 0xb1, 0x38, 0xde, 0x6c, 0x13, 0xf8, 0x98, 0x43, 0xbe,
   0x3f, 0x85, 0x83, 0xd0, 0x11, 0x88, 0xb8, 0x0f, 0xb5, 0x8c, 0x2a, 0x1c};

static void
_test_key_broker_kmip (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_key_broker_t kb;
   bson_t keydoc_bson;
   bson_iter_t iter;
   _mongocrypt_buffer_t id;
   _mongocrypt_buffer_t keydoc;
   mongocrypt_kms_ctx_t *kms;
   mongocrypt_binary_t *msg;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   _mongocrypt_buffer_t secretdata;

   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   status = mongocrypt_status_new ();
   kms_providers = &crypt->opts.kms_providers;
   _mongocrypt_key_broker_init (&kb, crypt);
   _load_json_as_bson ("./test/data/key-document-kmip.json", &keydoc_bson);

   ASSERT_OR_PRINT_MSG (bson_iter_init_find (&iter, &keydoc_bson, "_id"),
                        "could not find _id in key-document-kmip.json");
   BSON_ASSERT (_mongocrypt_buffer_from_binary_iter (&id, &iter));
   ASSERT_OK (_mongocrypt_key_broker_request_id (&kb, &id), &kb);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&kb), &kb);

   /* Add the key document. */
   _mongocrypt_buffer_from_bson (&keydoc, &keydoc_bson);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &keydoc),
              &kb);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);

   /* There should be exactly one KMS request for KMIP. */
   kms = _mongocrypt_key_broker_next_kms (&kb);
   ASSERT_OR_PRINT_MSG (kms, "expected KMS context returned, got none");

   msg = mongocrypt_binary_new ();
   mongocrypt_kms_ctx_message (kms, msg);
   ASSERT_CMPBYTES (EXPECTED_GET_REQUEST,
                    sizeof (EXPECTED_GET_REQUEST),
                    mongocrypt_binary_data (msg),
                    mongocrypt_binary_len (msg));

   ASSERT_OK (kms_ctx_feed_all (
                 kms, SUCCESS_GET_RESPONSE, sizeof (SUCCESS_GET_RESPONSE)),
              kms);
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&kb, kms_providers), &kb);

   BSON_ASSERT (
      _mongocrypt_key_broker_decrypted_key_by_id (&kb, &id, &secretdata));
   ASSERT_CMPBYTES (secretdata.data,
                    secretdata.len,
                    EXPECTED_SECRETDATA,
                    sizeof (EXPECTED_SECRETDATA));

   _mongocrypt_buffer_cleanup (&secretdata);
   mongocrypt_binary_destroy (msg);
   _mongocrypt_buffer_cleanup (&keydoc);
   _mongocrypt_buffer_cleanup (&id);
   bson_destroy (&keydoc_bson);
   _mongocrypt_key_broker_cleanup (&kb);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
}

/*
<ResponseMessage tag="0x42007b" type="Structure">
 <ResponseHeader tag="0x42007a" type="Structure">
  <ProtocolVersion tag="0x420069" type="Structure">
   <ProtocolVersionMajor tag="0x42006a" type="Integer" value="1"/>
   <ProtocolVersionMinor tag="0x42006b" type="Integer" value="4"/>
  </ProtocolVersion>
  <TimeStamp tag="0x420092" type="DateTime" value="2021-10-01T14:43:13-0500"/>
  <BatchCount tag="0x42000d" type="Integer" value="1"/>
 </ResponseHeader>
 <BatchItem tag="0x42000f" type="Structure">
  <Operation tag="0x42005c" type="Enumeration" value="10"/>
  <ResultStatus tag="0x42007f" type="Enumeration" value="1"/>
  <ResultReason tag="0x42007e" type="Enumeration" value="1"/>
  <ResultMessage tag="0x42007d" type="TextString"
value="ResultReasonItemNotFound"/>
 </BatchItem>
</ResponseMessage>
*/
static const uint8_t ERROR_GET_RESPOSE_NOTFOUND[] = {
   0x42, 0x00, 0x7b, 0x01, 0x00, 0x00, 0x00, 0xa8, 0x42, 0x00, 0x7a, 0x01, 0x00,
   0x00, 0x00, 0x48, 0x42, 0x00, 0x69, 0x01, 0x00, 0x00, 0x00, 0x20, 0x42, 0x00,
   0x6a, 0x02, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   0x00, 0x42, 0x00, 0x6b, 0x02, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04,
   0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x92, 0x09, 0x00, 0x00, 0x00, 0x08, 0x00,
   0x00, 0x00, 0x00, 0x61, 0x57, 0x1e, 0x81, 0x42, 0x00, 0x0d, 0x02, 0x00, 0x00,
   0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x0f,
   0x01, 0x00, 0x00, 0x00, 0x50, 0x42, 0x00, 0x5c, 0x05, 0x00, 0x00, 0x00, 0x04,
   0x00, 0x00, 0x00, 0x0a, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00, 0x7f, 0x05, 0x00,
   0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x42, 0x00,
   0x7e, 0x05, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00,
   0x00, 0x42, 0x00, 0x7d, 0x07, 0x00, 0x00, 0x00, 0x18, 0x52, 0x65, 0x73, 0x75,
   0x6c, 0x74, 0x52, 0x65, 0x61, 0x73, 0x6f, 0x6e, 0x49, 0x74, 0x65, 0x6d, 0x4e,
   0x6f, 0x74, 0x46, 0x6f, 0x75, 0x6e, 0x64};

static void
_test_key_broker_kmip_notfound (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_key_broker_t kb;
   bson_t keydoc_bson;
   bson_iter_t iter;
   _mongocrypt_buffer_t id;
   _mongocrypt_buffer_t keydoc;
   mongocrypt_kms_ctx_t *kms;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   mongocrypt_binary_t *msg;

   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   status = mongocrypt_status_new ();
   kms_providers = &crypt->opts.kms_providers;
   _mongocrypt_key_broker_init (&kb, crypt);
   _load_json_as_bson ("./test/data/key-document-kmip.json", &keydoc_bson);

   ASSERT_OR_PRINT_MSG (bson_iter_init_find (&iter, &keydoc_bson, "_id"),
                        "could not find _id in key-document-kmip.json");
   BSON_ASSERT (_mongocrypt_buffer_from_binary_iter (&id, &iter));
   ASSERT_OK (_mongocrypt_key_broker_request_id (&kb, &id), &kb);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&kb), &kb);

   /* Add the key document. */
   _mongocrypt_buffer_from_bson (&keydoc, &keydoc_bson);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &keydoc),
              &kb);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);

   /* There should be exactly one KMS request for KMIP. */
   kms = _mongocrypt_key_broker_next_kms (&kb);
   ASSERT_OR_PRINT_MSG (kms, "expected KMS context returned, got none");

   msg = mongocrypt_binary_new ();
   mongocrypt_kms_ctx_message (kms, msg);
   ASSERT_CMPBYTES (EXPECTED_GET_REQUEST,
                    sizeof (EXPECTED_GET_REQUEST),
                    mongocrypt_binary_data (msg),
                    mongocrypt_binary_len (msg));

   ASSERT_FAILS (kms_ctx_feed_all (kms,
                                   ERROR_GET_RESPOSE_NOTFOUND,
                                   sizeof (ERROR_GET_RESPOSE_NOTFOUND)),
                 kms,
                 "ResultReasonItemNotFound");

   mongocrypt_binary_destroy (msg);
   _mongocrypt_buffer_cleanup (&keydoc);
   _mongocrypt_buffer_cleanup (&id);
   bson_destroy (&keydoc_bson);
   _mongocrypt_key_broker_cleanup (&kb);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
}


static void
_test_key_broker_request_any (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *const crypt =
      _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   _mongocrypt_key_broker_t kb;

   /* Can switch to any mode before any keys are requested. */
   _mongocrypt_key_broker_init (&kb, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_any (&kb), &kb);
   _mongocrypt_key_broker_cleanup (&kb);

   /* If keys have already been requested, cannot switch to any mode.*/
   _mongocrypt_key_broker_init (&kb, crypt);
   _key_broker_add_name (&kb, "test");
   ASSERT_FAILS (_mongocrypt_key_broker_request_any (&kb),
                 &kb,
                 "attempting to request any keys, but requests already made");
   _mongocrypt_key_broker_cleanup (&kb);

   mongocrypt_destroy (crypt);
}


static void
_test_key_broker_add_any (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *const crypt =
      _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   _mongocrypt_opts_kms_providers_t *const kms_providers =
      &crypt->opts.kms_providers;
   _mongocrypt_key_broker_t kb;
   _mongocrypt_buffer_t key_id, key_doc;

   _gen_uuid_and_key (tester, 1, &key_id, &key_doc);

   /* Can add valid documents. */
   _mongocrypt_key_broker_init (&kb, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_any (&kb), &kb);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc),
              &kb);
   ASSERT (_key_broker_num_satisfied (&kb) == 1);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);
   _mongocrypt_key_broker_cleanup (&kb);

   /* Still validates no duplicate/incompatible keys. */
   _mongocrypt_key_broker_init (&kb, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_any (&kb), &kb);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc),
              &kb);
   ASSERT (_key_broker_num_satisfied (&kb) == 1);
   ASSERT_FAILS (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc),
                 &kb,
                 "keys returned have duplicate keyAltNames or _id");
   ASSERT (_key_broker_num_satisfied (&kb) == 1);
   _mongocrypt_key_broker_cleanup (&kb);

   /* Still requests KMS as needed. */
   _mongocrypt_key_broker_init (&kb, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_any (&kb), &kb);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc),
              &kb);
   ASSERT (_key_broker_num_satisfied (&kb) == 1);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);
   {
      mongocrypt_kms_ctx_t *kms;
      kms = _mongocrypt_key_broker_next_kms (&kb);
      ASSERT_OK (kms, &kb);
      /* Key material should should be cached here. */
      _mongocrypt_tester_satisfy_kms (tester, kms);
      ASSERT_OK (!_mongocrypt_key_broker_next_kms (&kb), &kb);
      ASSERT_OK (_mongocrypt_key_broker_kms_done (&kb, kms_providers), &kb);
   }
   _mongocrypt_key_broker_cleanup (&kb);
   _mongocrypt_key_broker_init (&kb, crypt);
   ASSERT_OK (_mongocrypt_key_broker_request_any (&kb), &kb);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc),
              &kb);
   ASSERT (_key_broker_num_satisfied (&kb) == 1);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);
   /* Cached key material, should skip straight to done. */
   ASSERT (kb.state == KB_DONE);

   _mongocrypt_key_broker_cleanup (&kb);

   _mongocrypt_buffer_cleanup (&key_doc);
   _mongocrypt_buffer_cleanup (&key_id);
   mongocrypt_destroy (crypt);
}

/* assert_filter_requests_id asserts that filter contains exactly one request
 * for a key ID of expect. */
static void
assert_filter_requests_id (mongocrypt_binary_t *filter,
                           const _mongocrypt_buffer_t *expect)
{
   bson_t filter_bson;
   ASSERT (_mongocrypt_binary_to_bson (filter, &filter_bson));
   bson_iter_t iter;
   ASSERT (bson_iter_init (&iter, &filter_bson));
   ASSERT (bson_iter_find_descendant (&iter, "$or.0._id.$in", &iter));
   ASSERT (bson_iter_recurse (&iter, &iter));
   ASSERT (bson_iter_next (&iter));
   ASSERT (BSON_ITER_HOLDS_BINARY (&iter));

   _mongocrypt_buffer_t actual;
   ASSERT (_mongocrypt_buffer_from_binary_iter (&actual, &iter));
   ASSERT_CMPBUF (*expect, actual);

   /* Check that there are no additional _id's requested. */
   ASSERT (!bson_iter_next (&iter));

   /* Check that there are no keyAltName requests. */
   ASSERT (bson_iter_init (&iter, &filter_bson));
   ASSERT (bson_iter_find_descendant (&iter, "$or.1.keyAltNames.$in", &iter));
   ASSERT (bson_iter_recurse (&iter, &iter));
   ASSERT (!bson_iter_next (&iter));
}

/* Test that key requests can be added again after transitioning to DONE and
 * calling _mongocrypt_key_broker_restart. */
static void
_test_key_broker_restart (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t key_id1, key_id2, key_doc1, key_doc2;
   _mongocrypt_key_broker_t kb;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   mongocrypt_kms_ctx_t *kms;
   mongocrypt_binary_t *filter;

   status = mongocrypt_status_new ();
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _gen_uuid_and_key (tester, 1, &key_id1, &key_doc1);
   _gen_uuid_and_key (tester, 2, &key_id2, &key_doc2);
   _mongocrypt_key_broker_init (&kb, crypt);

   ASSERT (kb.state == KB_REQUESTING);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&kb, &key_id1), &kb);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&kb), &kb);

   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_filter (&kb, filter), &kb);
   assert_filter_requests_id (filter, &key_id1);
   mongocrypt_binary_destroy (filter);

   ASSERT (kb.state == KB_ADDING_DOCS);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc1),
              &kb);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);

   ASSERT (kb.state == KB_DECRYPTING_KEY_MATERIAL);
   kms = _mongocrypt_key_broker_next_kms (&kb);
   ASSERT (kms);
   _mongocrypt_tester_satisfy_kms (tester, kms);
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&kb, kms_providers), &kb);

   /* Restart, and add request for key_id2. */
   ASSERT (kb.state == KB_DONE);
   ASSERT_OK (_mongocrypt_key_broker_restart (&kb), &kb);

   ASSERT (kb.state == KB_REQUESTING);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&kb, &key_id2), &kb);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&kb), &kb);

   filter = mongocrypt_binary_new ();
   ASSERT_OK (_mongocrypt_key_broker_filter (&kb, filter), &kb);
   assert_filter_requests_id (filter, &key_id2);
   mongocrypt_binary_destroy (filter);

   ASSERT (kb.state == KB_ADDING_DOCS);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc2),
              &kb);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);

   ASSERT (kb.state == KB_DECRYPTING_KEY_MATERIAL);
   kms = _mongocrypt_key_broker_next_kms (&kb);
   ASSERT (kms);
   _mongocrypt_tester_satisfy_kms (tester, kms);
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&kb, kms_providers), &kb);
   ASSERT (kb.state == KB_DONE);

   _mongocrypt_key_broker_cleanup (&kb);
   _mongocrypt_buffer_cleanup (&key_doc2);
   _mongocrypt_buffer_cleanup (&key_id2);
   _mongocrypt_buffer_cleanup (&key_doc1);
   _mongocrypt_buffer_cleanup (&key_id1);
   mongocrypt_destroy (crypt);
   mongocrypt_status_destroy (status);
}

/* Test that a decrypted key can be returned while in the KB_REQUESTING state.
 */
static void
_test_key_broker_get_decrypted_key_while_requesting (
   _mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t key_id1, key_doc1, key_decrypted1, key_decrypted1_copy;
   _mongocrypt_key_broker_t kb;
   _mongocrypt_opts_kms_providers_t *kms_providers;
   mongocrypt_kms_ctx_t *kms;

   status = mongocrypt_status_new ();
   crypt = _mongocrypt_tester_mongocrypt (TESTER_MONGOCRYPT_DEFAULT);
   kms_providers = &crypt->opts.kms_providers;
   _gen_uuid_and_key (tester, 1, &key_id1, &key_doc1);
   _mongocrypt_key_broker_init (&kb, crypt);

   ASSERT (kb.state == KB_REQUESTING);
   ASSERT_OK (_mongocrypt_key_broker_request_id (&kb, &key_id1), &kb);
   ASSERT_OK (_mongocrypt_key_broker_requests_done (&kb), &kb);

   ASSERT (kb.state == KB_ADDING_DOCS);
   ASSERT_OK (_mongocrypt_key_broker_add_doc (&kb, kms_providers, &key_doc1),
              &kb);
   ASSERT_OK (_mongocrypt_key_broker_docs_done (&kb), &kb);

   ASSERT (kb.state == KB_DECRYPTING_KEY_MATERIAL);
   kms = _mongocrypt_key_broker_next_kms (&kb);
   ASSERT (kms);
   _mongocrypt_tester_satisfy_kms (tester, kms);
   ASSERT_OK (_mongocrypt_key_broker_kms_done (&kb, kms_providers), &kb);

   ASSERT (kb.state == KB_DONE);
   ASSERT_OK (_mongocrypt_key_broker_decrypted_key_by_id (
                 &kb, &key_id1, &key_decrypted1),
              &kb);

   /* Restart. */
   ASSERT_OK (_mongocrypt_key_broker_restart (&kb), &kb);
   ASSERT (kb.state == KB_REQUESTING);
   ASSERT_OK (_mongocrypt_key_broker_decrypted_key_by_id (
                 &kb, &key_id1, &key_decrypted1_copy),
              &kb);

   ASSERT_CMPBUF (key_decrypted1, key_decrypted1_copy);

   _mongocrypt_buffer_cleanup (&key_decrypted1_copy);
   _mongocrypt_buffer_cleanup (&key_decrypted1);
   _mongocrypt_key_broker_cleanup (&kb);
   _mongocrypt_buffer_cleanup (&key_doc1);
   _mongocrypt_buffer_cleanup (&key_id1);
   mongocrypt_destroy (crypt);
   mongocrypt_status_destroy (status);
}

void
_mongocrypt_tester_install_key_broker (_mongocrypt_tester_t *tester)
{
   INSTALL_TEST (_test_key_broker_get_key_filter);
   INSTALL_TEST (_test_key_broker_add_key);
   INSTALL_TEST (_test_key_broker_add_decrypted_key);
   INSTALL_TEST (_test_key_broker_wrong_subtype);
   INSTALL_TEST (_test_key_broker_multi_match);
   INSTALL_TEST (_test_key_broker_kmip);
   INSTALL_TEST (_test_key_broker_kmip_notfound);
   INSTALL_TEST (_test_key_broker_request_any);
   INSTALL_TEST (_test_key_broker_add_any);
   INSTALL_TEST (_test_key_broker_restart);
   INSTALL_TEST (_test_key_broker_get_decrypted_key_while_requesting);
}
