/*
 * 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-config.h"
#include "mongocrypt-private.h"
#include "mongocrypt-crypto-private.h"

#include "test-mongocrypt.h"

#define IV_HEX "1F572A1B84EC8F99B7915AA2A2AEA2F4"
#define HMAC_HEX                                                      \
   "60676DE9FD305FD2C0815763C422687270DA2416D94A917B276E9DCBB13F412F" \
   "92FA403AA8AE172BD2E4729ED352793795EE588A2977C9C1F218D2AAD779C997"
/* only the first 32 bytes are appended. */
#define HMAC_HEX_TAG \
   "60676DE9FD305FD2C0815763C422687270DA2416D94A917B276E9DCBB13F412F"

#define HMAC_KEY_HEX \
   "CCD3836C8F24AC5FAAFAAA630C5C6C5D210FD03934EA1440CD67E0DCDE3F8EA6"
#define ENCRYPTION_KEY_HEX \
   "E1D1727BAF970E01181C0868CB9D3E574B47AC09771FF30FE2D093B0950C7DAF"
#define IV_KEY_HEX \
   "0A9328FCB6405ABDF5B4BFEC243FE9CF503CD4F24360872B75F08A2A3961802B"
/* full 96 byte key consists of three "sub" keys */
#define KEY_HEX HMAC_KEY_HEX ENCRYPTION_KEY_HEX IV_KEY_HEX
#define HASH_HEX \
   "489EC3238378DC624C74B8CC4598ACED2B7EA5DE5C5F7602D8761BAE92FD8ABE"
#define RANDOM_HEX                                                             \
   "670ACBB44D4E04A279CC0B95D217493205A038C50F537F452C59EFF6541D0026670ACBB44" \
   "D4E04A279CC0B95D217493205A038C50F537F452C59EFF6541D0026670ACBB44D4E04A279" \
   "CC0B95D217493205A038C50F537F452C59EFF6541D0026"

/* a document containing the history of calls */
static bson_string_t *call_history;

static void
_append_bin (const char *name, mongocrypt_binary_t *bin)
{
   _mongocrypt_buffer_t tmp;
   char *hex;

   _mongocrypt_buffer_from_binary (&tmp, bin);
   hex = _mongocrypt_buffer_to_hex (&tmp);
   bson_string_append_printf (call_history, "%s:%s\n", name, hex);
   bson_free (hex);
   _mongocrypt_buffer_cleanup (&tmp);
}

static bool
_mock_aes_256_xxx_encrypt (void *ctx,
                           mongocrypt_binary_t *key,
                           mongocrypt_binary_t *iv,
                           mongocrypt_binary_t *in,
                           mongocrypt_binary_t *out,
                           uint32_t *bytes_written,
                           mongocrypt_status_t *status)
{
   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   _append_bin ("key", key);
   if (NULL != iv) {
      _append_bin ("iv", iv);
   }
   _append_bin ("in", in);
   /* append it directly, don't encrypt. */
   memcpy (out->data + *bytes_written, in->data, in->len);
   *bytes_written += in->len;
   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);
   if (0 == strcmp ((char *) ctx, "error_on:aes_256_cbc_encrypt") ||
       0 == strcmp ((char *) ctx, "error_on:aes_256_ctr_encrypt") ||
       0 == strcmp ((char *) ctx, "error_on:aes_256_ecb_encrypt")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

static bool
_mock_aes_256_xxx_decrypt (void *ctx,
                           mongocrypt_binary_t *key,
                           mongocrypt_binary_t *iv,
                           mongocrypt_binary_t *in,
                           mongocrypt_binary_t *out,
                           uint32_t *bytes_written,
                           mongocrypt_status_t *status)
{
   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   _append_bin ("key", key);
   _append_bin ("iv", iv);
   _append_bin ("in", in);
   /* append it directly, don't decrypt. */
   memcpy (out->data + *bytes_written, in->data, in->len);
   *bytes_written += in->len;
   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);
   if (0 == strcmp ((char *) ctx, "error_on:aes_256_cbc_decrypt") ||
       0 == strcmp ((char *) ctx, "error_on:aes_256_ctr_decrypt")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

bool
_hmac_sha_512 (void *ctx,
               mongocrypt_binary_t *key,
               mongocrypt_binary_t *in,
               mongocrypt_binary_t *out,
               mongocrypt_status_t *status)
{
   _mongocrypt_buffer_t tmp;

   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   _append_bin ("key", key);
   _append_bin ("in", in);

   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);

   _mongocrypt_buffer_copy_from_hex (&tmp, HMAC_HEX);
   memcpy (out->data, tmp.data, tmp.len);
   _mongocrypt_buffer_cleanup (&tmp);
   if (0 == strcmp ((char *) ctx, "error_on:hmac_sha512")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

bool
_hmac_sha_256 (void *ctx,
               mongocrypt_binary_t *key,
               mongocrypt_binary_t *in,
               mongocrypt_binary_t *out,
               mongocrypt_status_t *status)
{
   _mongocrypt_buffer_t tmp;

   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   _append_bin ("key", key);
   _append_bin ("in", in);

   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);

   _mongocrypt_buffer_copy_from_hex (&tmp, HASH_HEX);
   memcpy (out->data, tmp.data, tmp.len);
   _mongocrypt_buffer_cleanup (&tmp);
   if (0 == strcmp ((char *) ctx, "error_on:hmac_sha256")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

bool
_sha_256 (void *ctx,
          mongocrypt_binary_t *in,
          mongocrypt_binary_t *out,
          mongocrypt_status_t *status)
{
   _mongocrypt_buffer_t tmp;

   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   _append_bin ("in", in);

   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);

   _mongocrypt_buffer_copy_from_hex (&tmp, HASH_HEX);
   memcpy (out->data, tmp.data, tmp.len);
   _mongocrypt_buffer_cleanup (&tmp);
   if (0 == strcmp ((char *) ctx, "error_on:sha256")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

bool
_random (void *ctx,
         mongocrypt_binary_t *out,
         uint32_t count,
         mongocrypt_status_t *status)
{
   /* only have 32 bytes of random test data. */
   BSON_ASSERT (count <= 96);

   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   bson_string_append_printf (call_history, "count:%d\n", (int) count);
   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);

   _mongocrypt_buffer_t tmp;
   _mongocrypt_buffer_copy_from_hex (&tmp, RANDOM_HEX);
   memcpy (out->data, tmp.data, count);
   _mongocrypt_buffer_cleanup (&tmp);
   if (0 == strcmp ((char *) ctx, "error_on:random")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

bool
_sign_rsaes_pkcs1_v1_5 (void *ctx,
                        mongocrypt_binary_t *key,
                        mongocrypt_binary_t *in,
                        mongocrypt_binary_t *out,
                        mongocrypt_status_t *status)
{
   _mongocrypt_buffer_t tmp;

   BSON_ASSERT (0 == strncmp ("error_on:", (char *) ctx, strlen ("error_on:")));
   bson_string_append_printf (call_history, "call:%s\n", BSON_FUNC);
   _append_bin ("key", key);
   _append_bin ("in", in);

   bson_string_append_printf (call_history, "ret:%s\n", BSON_FUNC);
   memset (out->data, 0, out->len);

   _mongocrypt_buffer_copy_from_hex (&tmp, HASH_HEX);
   memcpy (out->data, tmp.data, tmp.len);
   _mongocrypt_buffer_cleanup (&tmp);
   if (0 == strcmp ((char *) ctx, "error_on:sign_rsaes_pkcs1_v1_5")) {
      mongocrypt_status_set (
         status, MONGOCRYPT_STATUS_ERROR_CLIENT, 1, (char *) ctx, -1);
      return false;
   }
   return true;
}

static mongocrypt_t *
_create_mongocrypt_and_hooks (_mongocrypt_tester_t *tester,
                              const char *error_on,
                              bool ctr_hook,
                              bool ecb_hook)
{
   bool ret;

   mongocrypt_t *crypt = mongocrypt_new ();
   ASSERT_OK (
      mongocrypt_setopt_kms_provider_aws (crypt, "example", -1, "example", -1),
      crypt);
   ASSERT_OK (
      mongocrypt_setopt_kms_providers (
         crypt,
         TEST_BSON ("{'gcp': { 'email': 'test', 'privateKey': 'AAAA'}}")),
      crypt);
   ret = mongocrypt_setopt_crypto_hooks (crypt,
                                         _mock_aes_256_xxx_encrypt,
                                         _mock_aes_256_xxx_decrypt,
                                         _random,
                                         _hmac_sha_512,
                                         _hmac_sha_256,
                                         _sha_256,
                                         (void *) error_on);
   ASSERT_OK (ret, crypt);
   ret = mongocrypt_setopt_crypto_hook_sign_rsaes_pkcs1_v1_5 (
      crypt, _sign_rsaes_pkcs1_v1_5, (void *) error_on);
   ASSERT_OK (ret, crypt);
   if (ctr_hook) {
      ret = mongocrypt_setopt_aes_256_ctr (crypt,
                                           _mock_aes_256_xxx_encrypt,
                                           _mock_aes_256_xxx_decrypt,
                                           (void *) error_on);
      ASSERT_OK (ret, crypt);
   }
   if (ecb_hook) {
      ret = mongocrypt_setopt_aes_256_ecb (
         crypt, _mock_aes_256_xxx_encrypt, (void *) error_on);
      ASSERT_OK (ret, crypt);
   }
   ASSERT_OK (mongocrypt_init (crypt), crypt);
   return crypt;
}

static mongocrypt_t *
_create_mongocrypt (_mongocrypt_tester_t *tester, const char *error_on)
{
   return _create_mongocrypt_and_hooks (tester, error_on, false, false);
}

static void
_test_crypto_hooks_encryption_helper (_mongocrypt_tester_t *tester,
                                      const char *error_on,
                                      bool ctr_hook,
                                      bool ecb_hook)
{
   mongocrypt_t *crypt;
   bool ret;
   uint32_t bytes_written;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t iv, associated_data, key, plaintext, ciphertext;
   const char *expected_call_history =
      "call:_mock_aes_256_xxx_encrypt\n"
      "key:" ENCRYPTION_KEY_HEX "\n"
      "iv:" IV_HEX "\n"
      "in:BBBB0E0E0E0E0E0E0E0E0E0E0E0E0E0E\n"
      "ret:_mock_aes_256_xxx_encrypt\n"
      "call:_hmac_sha_512\n"
      "key:CCD3836C8F24AC5FAAFAAA630C5C6C5D210FD03934EA1440CD67E0DCDE3F8EA6\n"
      "in:AAAA" IV_HEX "BBBB0E0E0E0E0E0E0E0E0E0E0E0E0E0E0000000000000010\n"
      "ret:_hmac_sha_512\n";

   status = mongocrypt_status_new ();
   crypt = _create_mongocrypt_and_hooks (tester, error_on, ctr_hook, ecb_hook);

   _mongocrypt_buffer_copy_from_hex (&iv, IV_HEX);
   _mongocrypt_buffer_copy_from_hex (&associated_data, "AAAA");
   _mongocrypt_buffer_copy_from_hex (&plaintext, "BBBB");

   call_history = bson_string_new (NULL);

   if (ctr_hook || ecb_hook) {
      _mongocrypt_buffer_copy_from_hex (&key, ENCRYPTION_KEY_HEX);
      _mongocrypt_buffer_init (&ciphertext);
      _mongocrypt_buffer_resize (
         &ciphertext,
         _mongocrypt_fle2_calculate_ciphertext_len (plaintext.len, status));
      ret = _mongocrypt_fle2_do_encryption (crypt->crypto,
                                            &iv,
                                            &key,
                                            &plaintext,
                                            &ciphertext,
                                            &bytes_written,
                                            status);
   } else {
      _mongocrypt_buffer_copy_from_hex (&key, KEY_HEX);
      _mongocrypt_buffer_init (&ciphertext);
      _mongocrypt_buffer_resize (
         &ciphertext,
         _mongocrypt_calculate_ciphertext_len (plaintext.len, status));
      ret = _mongocrypt_do_encryption (crypt->crypto,
                                       &iv,
                                       &associated_data,
                                       &key,
                                       &plaintext,
                                       &ciphertext,
                                       &bytes_written,
                                       status);
   }

   if (0 == strcmp (error_on, "error_on:none")) {
      ASSERT_OK_STATUS (ret, status);
      ciphertext.len = bytes_written;

      /* Check the full trace. */
      ASSERT_STREQUAL (call_history->str, expected_call_history);

      /* Check the structure of the ciphertext */
      BSON_ASSERT (
         0 == _mongocrypt_buffer_cmp_hex (
                 &ciphertext,
                 IV_HEX "BBBB0E0E0E0E0E0E0E0E0E0E0E0E0E0E" /* the "encrypted"
                                                            block which is
                                                            really plaintext.
                                                            BBBB + padding. */
                 HMAC_HEX_TAG));
   } else {
      ASSERT_FAILS_STATUS (ret, status, error_on);
   }


   _mongocrypt_buffer_cleanup (&key);
   _mongocrypt_buffer_cleanup (&iv);
   _mongocrypt_buffer_cleanup (&associated_data);
   _mongocrypt_buffer_cleanup (&plaintext);
   _mongocrypt_buffer_cleanup (&ciphertext);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}


static void
_test_crypto_hooks_encryption (_mongocrypt_tester_t *tester)
{
   _test_crypto_hooks_encryption_helper (tester, "error_on:none", false, false);
   _test_crypto_hooks_encryption_helper (
      tester, "error_on:aes_256_cbc_encrypt", false, false);
   _test_crypto_hooks_encryption_helper (
      tester, "error_on:aes_256_ctr_encrypt", true, false);
   _test_crypto_hooks_encryption_helper (
      tester, "error_on:aes_256_ecb_encrypt", false, true);
   _test_crypto_hooks_encryption_helper (
      tester, "error_on:hmac_sha512", false, false);
}


static void
_test_crypto_hooks_decryption_helper (_mongocrypt_tester_t *tester,
                                      const char *error_on,
                                      bool ctr_hook,
                                      bool ecb_hook)
{
   mongocrypt_t *crypt;
   bool ret;
   uint32_t bytes_written;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t associated_data, key, plaintext, ciphertext;
   const char *expected_call_history =
      "call:_hmac_sha_512\n"
      "key:" HMAC_KEY_HEX "\n"
      "in:AAAA" IV_HEX "BBBB0E0E0E0E0E0E0E0E0E0E0E0E0E0E0000000000000010\n"
      "ret:_hmac_sha_512\n"
      "call:_mock_aes_256_xxx_decrypt\n"
      "key:" ENCRYPTION_KEY_HEX "\n"
      "iv:" IV_HEX "\n"
      "in:BBBB0E0E0E0E0E0E0E0E0E0E0E0E0E0E\n"
      "ret:_mock_aes_256_xxx_decrypt\n";

   status = mongocrypt_status_new ();
   crypt = _create_mongocrypt_and_hooks (tester, error_on, ctr_hook, ecb_hook);

   _mongocrypt_buffer_copy_from_hex (&associated_data, "AAAA");
   _mongocrypt_buffer_copy_from_hex (
      &ciphertext, IV_HEX "BBBB0E0E0E0E0E0E0E0E0E0E0E0E0E0E" HMAC_HEX_TAG);

   call_history = bson_string_new (NULL);

   if (ctr_hook || ecb_hook) {
      _mongocrypt_buffer_copy_from_hex (&key, ENCRYPTION_KEY_HEX);
      _mongocrypt_buffer_init (&plaintext);
      _mongocrypt_buffer_resize (
         &plaintext,
         _mongocrypt_fle2_calculate_plaintext_len (ciphertext.len, status));

      ret = _mongocrypt_fle2_do_decryption (
         crypt->crypto, &key, &ciphertext, &plaintext, &bytes_written, status);
   } else {
      _mongocrypt_buffer_copy_from_hex (&key, KEY_HEX);
      _mongocrypt_buffer_init (&plaintext);
      _mongocrypt_buffer_resize (
         &plaintext,
         _mongocrypt_calculate_plaintext_len (ciphertext.len, status));

      ret = _mongocrypt_do_decryption (crypt->crypto,
                                       &associated_data,
                                       &key,
                                       &ciphertext,
                                       &plaintext,
                                       &bytes_written,
                                       status);
   }

   if (0 == strcmp (error_on, "error_on:none")) {
      ASSERT_OK_STATUS (ret, status);
      plaintext.len = bytes_written;

      /* Check the full trace. */
      ASSERT_STREQUAL (call_history->str, expected_call_history);

      /* Check the resulting plaintext */
      BSON_ASSERT (0 == _mongocrypt_buffer_cmp_hex (&plaintext, "BBBB"));
   } else {
      ASSERT_FAILS_STATUS (ret, status, error_on);
   }

   _mongocrypt_buffer_cleanup (&key);
   _mongocrypt_buffer_cleanup (&associated_data);
   _mongocrypt_buffer_cleanup (&plaintext);
   _mongocrypt_buffer_cleanup (&ciphertext);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

static void
_test_crypto_hooks_decryption (_mongocrypt_tester_t *tester)
{
   _test_crypto_hooks_decryption_helper (tester, "error_on:none", false, false);
   _test_crypto_hooks_decryption_helper (
      tester, "error_on:aes_256_cbc_decrypt", false, false);
   _test_crypto_hooks_decryption_helper (
      tester, "error_on:aes_256_ctr_decrypt", true, false);
   _test_crypto_hooks_decryption_helper (
      tester, "error_on:aes_256_ecb_encrypt", false, true);
   _test_crypto_hooks_decryption_helper (
      tester, "error_on:hmac_sha512", false, false);
}

static void
_test_crypto_hooks_iv_gen_helper (_mongocrypt_tester_t *tester, char *error_on)
{
   mongocrypt_t *crypt;
   bool ret;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t associated_data, key, plaintext, iv;
   char *expected_iv = bson_strndup (
      HMAC_HEX_TAG, 16 * 2); /* only the first 16 bytes are used for IV. */
   const char *expected_call_history = "call:_hmac_sha_512\n"
                                       "key:" IV_KEY_HEX "\n"
                                       "in:AAAA0000000000000010BBBB\n"
                                       "ret:_hmac_sha_512\n";

   status = mongocrypt_status_new ();
   crypt = _create_mongocrypt (tester, error_on);

   _mongocrypt_buffer_copy_from_hex (&associated_data, "AAAA");
   _mongocrypt_buffer_copy_from_hex (&key, KEY_HEX);
   _mongocrypt_buffer_copy_from_hex (&plaintext, "BBBB");

   _mongocrypt_buffer_init (&iv);
   _mongocrypt_buffer_resize (&iv, MONGOCRYPT_IV_LEN);

   call_history = bson_string_new (NULL);

   ret = _mongocrypt_calculate_deterministic_iv (
      crypt->crypto, &key, &plaintext, &associated_data, &iv, status);

   if (0 == strcmp (error_on, "error_on:none")) {
      ASSERT_OK_STATUS (ret, status);

      /* Check the full trace. */
      ASSERT_STREQUAL (call_history->str, expected_call_history);

      /* Check the resulting iv */
      BSON_ASSERT (0 == _mongocrypt_buffer_cmp_hex (&iv, expected_iv));
   } else {
      ASSERT_FAILS_STATUS (ret, status, error_on);
   }

   bson_free (expected_iv);
   _mongocrypt_buffer_cleanup (&key);
   _mongocrypt_buffer_cleanup (&associated_data);
   _mongocrypt_buffer_cleanup (&plaintext);
   _mongocrypt_buffer_cleanup (&iv);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

static void
_test_crypto_hooks_iv_gen (_mongocrypt_tester_t *tester)
{
   _test_crypto_hooks_iv_gen_helper (tester, "error_on:none");
   _test_crypto_hooks_iv_gen_helper (tester, "error_on:hmac_sha512");
}


static void
_test_crypto_hooks_random_helper (_mongocrypt_tester_t *tester,
                                  const char *error_on)
{
   mongocrypt_t *crypt;
   bool ret;
   mongocrypt_status_t *status;
   _mongocrypt_buffer_t random;
   const char *expected_call_history = "call:_random\n"
                                       "count:96\n"
                                       "ret:_random\n";

   status = mongocrypt_status_new ();
   crypt = _create_mongocrypt (tester, error_on);

   _mongocrypt_buffer_init (&random);
   _mongocrypt_buffer_resize (&random, 96);

   call_history = bson_string_new (NULL);

   ret = _mongocrypt_random (crypt->crypto, &random, random.len, status);

   if (0 == strcmp (error_on, "error_on:none")) {
      ASSERT_OK_STATUS (ret, status);

      /* Check the full trace. */
      ASSERT_STREQUAL (call_history->str, expected_call_history);

      /* Check the resulting iv */
      BSON_ASSERT (0 == _mongocrypt_buffer_cmp_hex (&random, RANDOM_HEX));
   } else {
      ASSERT_FAILS_STATUS (ret, status, error_on);
   }

   _mongocrypt_buffer_cleanup (&random);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

static void
_test_crypto_hooks_random (_mongocrypt_tester_t *tester)
{
   _test_crypto_hooks_random_helper (tester, "error_on:none");
   _test_crypto_hooks_random_helper (tester, "error_on:random");
}

static void
_test_kms_request_helper (_mongocrypt_tester_t *tester, const char *error_on)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   mongocrypt_ctx_t *ctx;
   bool ret;

   status = mongocrypt_status_new ();
   crypt = _create_mongocrypt (tester, error_on);
   ctx = mongocrypt_ctx_new (crypt);

   call_history = bson_string_new (NULL);

   ASSERT_OK (
      mongocrypt_ctx_setopt_masterkey_aws (ctx, "us-east-1", -1, "cmk", -1),
      ctx);

   mongocrypt_ctx_datakey_init (ctx);
   ret = mongocrypt_ctx_status (ctx, status);

   if (0 == strcmp (error_on, "error_on:none")) {
      ASSERT_OK_STATUS (ret, status);

      /* The call history includes some random data, just assert we've called
       * our hooks. */
      BSON_ASSERT (strstr (call_history->str, "call:_hmac_sha_256"));
      BSON_ASSERT (strstr (call_history->str, "call:_sha_256"));
   } else {
      ASSERT_FAILS_STATUS (ret, status, error_on);
   }

   mongocrypt_ctx_destroy (ctx);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

static void
_test_kms_request (_mongocrypt_tester_t *tester)
{
   _test_kms_request_helper (tester, "error_on:none");
   _test_kms_request_helper (tester, "error_on:hmac_sha256");
   _test_kms_request_helper (tester, "error_on:sha256");
}


static void
_test_crypto_hooks_unset (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;

   crypt = mongocrypt_new ();
   mongocrypt_setopt_kms_provider_aws (crypt, "example", -1, "example", -1);
   ASSERT_OK (mongocrypt_init (crypt), crypt);
   mongocrypt_destroy (crypt);
}


/* test a bug fix, that an error on explicit encryption in the crypto hooks sets
 * the context state */
static void
_test_crypto_hooks_explicit_err (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_ctx_t *ctx;
   mongocrypt_binary_t *bin, *key_id;
   const char *deterministic = MONGOCRYPT_ALGORITHM_DETERMINISTIC_STR;

   call_history = bson_string_new (NULL);

   /* error on something during encryption. */
   crypt = _create_mongocrypt (tester, "error_on:hmac_sha512");

   ctx = mongocrypt_ctx_new (crypt);
   key_id = mongocrypt_binary_new_from_data (
      MONGOCRYPT_DATA_AND_LEN ("aaaaaaaaaaaaaaaa"));

   ASSERT_OK (mongocrypt_ctx_setopt_algorithm (ctx, deterministic, -1), ctx);
   ASSERT_OK (mongocrypt_ctx_setopt_key_id (ctx, key_id), ctx);
   ASSERT_OK (
      mongocrypt_ctx_explicit_encrypt_init (ctx, TEST_BSON ("{'v': 123}")),
      ctx);

   _mongocrypt_tester_run_ctx_to (tester, ctx, MONGOCRYPT_CTX_READY);
   bin = mongocrypt_binary_new ();
   ASSERT_FAILS (
      mongocrypt_ctx_finalize (ctx, bin), ctx, "error_on:hmac_sha512");
   BSON_ASSERT (MONGOCRYPT_CTX_ERROR == mongocrypt_ctx_state (ctx));
   mongocrypt_binary_destroy (bin);
   mongocrypt_binary_destroy (key_id);
   mongocrypt_ctx_destroy (ctx);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

/* validate that sha256 errors are handled correctly */
static void
_test_crypto_hooks_explicit_sha256_err (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_status_t *status;
   mongocrypt_ctx_t *ctx;

   status = mongocrypt_status_new ();
   crypt = _create_mongocrypt (tester, "error_on:sha256");
   ctx = mongocrypt_ctx_new (crypt);

   call_history = bson_string_new (NULL);

   ASSERT_OK (
      mongocrypt_ctx_setopt_masterkey_aws (ctx, "us-east-1", -1, "cmk", -1),
      ctx);
   ASSERT_FAILS (
      mongocrypt_ctx_datakey_init (ctx), ctx, "failed to create KMS message");

   mongocrypt_ctx_destroy (ctx);
   mongocrypt_status_destroy (status);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

static void
_test_crypto_hook_sign_rsaes_pkcs1_v1_5 (_mongocrypt_tester_t *tester)
{
   mongocrypt_t *crypt;
   mongocrypt_ctx_t *ctx;

   crypt = _create_mongocrypt (tester, "error_on:none");
   call_history = bson_string_new (NULL);

   ctx = mongocrypt_ctx_new (crypt);
   mongocrypt_ctx_setopt_key_encryption_key (
      ctx,
      TEST_BSON ("{'provider': 'gcp', 'projectId': 'test', 'location': "
                 "'global', 'keyRing': 'ring', 'keyName': 'key'}"));
   ASSERT_OK (mongocrypt_ctx_datakey_init (ctx), ctx);

   BSON_ASSERT (strstr (call_history->str, "call:_sign_rsaes_pkcs1_v1_5"));
   BSON_ASSERT (strstr (call_history->str, "key:000000"));

   mongocrypt_ctx_destroy (ctx);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);

   /* Test error when creating a data key. */
   crypt = _create_mongocrypt (tester, "error_on:sign_rsaes_pkcs1_v1_5");
   ctx = mongocrypt_ctx_new (crypt);
   call_history = bson_string_new (NULL);

   mongocrypt_ctx_setopt_key_encryption_key (
      ctx,
      TEST_BSON ("{'provider': 'gcp', 'projectId': 'test', 'location': "
                 "'global', 'keyRing': 'ring', 'keyName': 'key'}"));
   ASSERT_FAILS (
      mongocrypt_ctx_datakey_init (ctx), ctx, "error_on:sign_rsaes_pkcs1_v1_5");

   mongocrypt_ctx_destroy (ctx);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);

   /* Test error when encrypting. */
   crypt = _create_mongocrypt (tester, "error_on:sign_rsaes_pkcs1_v1_5");
   ctx = mongocrypt_ctx_new (crypt);
   call_history = bson_string_new (NULL);

   ASSERT_OK (mongocrypt_ctx_encrypt_init (
                 ctx, "test", -1, TEST_FILE ("./test/example/cmd.json")),
              ctx);
   _mongocrypt_tester_run_ctx_to (tester, ctx, MONGOCRYPT_CTX_NEED_MONGO_KEYS);
   ASSERT_FAILS (mongocrypt_ctx_mongo_feed (
                    ctx, TEST_FILE ("./test/data/key-document-gcp.json")),
                 ctx,
                 "error_on:sign_rsaes_pkcs1_v1_5");

   mongocrypt_ctx_destroy (ctx);
   mongocrypt_destroy (crypt);
   bson_string_free (call_history, true);
}

#ifdef MONGOCRYPT_ENABLE_CRYPTO_LIBCRYPTO
bool
_native_crypto_aes_256_ecb_encrypt (aes_256_args_t args);

static bool
_aes_256_ecb_encrypt (void *ctx,
                      mongocrypt_binary_t *key,
                      mongocrypt_binary_t *iv,
                      mongocrypt_binary_t *in,
                      mongocrypt_binary_t *out,
                      uint32_t *bytes_written,
                      mongocrypt_status_t *status)
{
   _mongocrypt_buffer_t key_buf;
   _mongocrypt_buffer_from_binary (&key_buf, key);
   if (iv) {
      CLIENT_ERR ("IV expected to be NULL in this mode");
      return false;
   }
   _mongocrypt_buffer_t in_buf;
   _mongocrypt_buffer_from_binary (&in_buf, in);
   _mongocrypt_buffer_t out_buf;
   _mongocrypt_buffer_from_binary (&out_buf, out);

   aes_256_args_t args = {
      &key_buf, NULL, &in_buf, &out_buf, bytes_written, status};

   return _native_crypto_aes_256_ecb_encrypt (args);
}


void
_test_fle2_crypto_via_ecb_hook (_mongocrypt_tester_t *tester)
{
   bool ret;
   _mongocrypt_buffer_t key;
   _mongocrypt_buffer_t iv;
   _mongocrypt_buffer_t plaintext;
   _mongocrypt_buffer_t ciphertext_reg;
   _mongocrypt_buffer_t ciphertext_ecb;
   _mongocrypt_buffer_t plaintext_ecb;
   uint32_t bytes_written;
   mongocrypt_status_t *status = mongocrypt_status_new ();

   _mongocrypt_buffer_copy_from_hex (&iv, IV_HEX);
   _mongocrypt_buffer_copy_from_hex (
      &plaintext,
      "4f6c64204d63446f6e616c64206861642061206661726d2e20456965696f0a");
   _mongocrypt_buffer_copy_from_hex (&key, ENCRYPTION_KEY_HEX);

   /* Encrypt data using native CTR and ECB-hook and compare */

   mongocrypt_t *crypt_reg = mongocrypt_new ();
   _mongocrypt_buffer_init (&ciphertext_reg);
   _mongocrypt_buffer_resize (
      &ciphertext_reg,
      _mongocrypt_fle2_calculate_ciphertext_len (plaintext.len, status));
   ret = _mongocrypt_fle2_do_encryption (crypt_reg->crypto,
                                         &iv,
                                         &key,
                                         &plaintext,
                                         &ciphertext_reg,
                                         &bytes_written,
                                         status);
   ASSERT_OK (ret, crypt_reg);

   mongocrypt_t *crypt_ecb = mongocrypt_new ();
   ret = mongocrypt_setopt_aes_256_ecb (crypt_ecb, _aes_256_ecb_encrypt, NULL);
   ASSERT_OK (ret, crypt_ecb);
   _mongocrypt_buffer_init (&ciphertext_ecb);
   _mongocrypt_buffer_resize (
      &ciphertext_ecb,
      _mongocrypt_fle2_calculate_ciphertext_len (plaintext.len, status));
   ret = _mongocrypt_fle2_do_encryption (crypt_ecb->crypto,
                                         &iv,
                                         &key,
                                         &plaintext,
                                         &ciphertext_ecb,
                                         &bytes_written,
                                         status);
   ASSERT_OK (ret, crypt_ecb);

   ASSERT (0 == _mongocrypt_buffer_cmp (&ciphertext_reg, &ciphertext_ecb));

   /* Decrypt data using ECB-hook and compare to original */

   _mongocrypt_buffer_init (&plaintext_ecb);
   _mongocrypt_buffer_resize (
      &plaintext_ecb,
      _mongocrypt_fle2_calculate_plaintext_len (ciphertext_ecb.len, status));
   ret = _mongocrypt_fle2_do_decryption (crypt_ecb->crypto,
                                         &key,
                                         &ciphertext_ecb,
                                         &plaintext_ecb,
                                         &bytes_written,
                                         status);
   ASSERT_OK (ret, crypt_ecb);

   ASSERT (0 == _mongocrypt_buffer_cmp (&plaintext, &plaintext_ecb));

   _mongocrypt_buffer_cleanup (&key);
   _mongocrypt_buffer_cleanup (&iv);
   _mongocrypt_buffer_cleanup (&plaintext);
   _mongocrypt_buffer_cleanup (&ciphertext_reg);
   _mongocrypt_buffer_cleanup (&ciphertext_ecb);
   _mongocrypt_buffer_cleanup (&plaintext_ecb);
   mongocrypt_destroy (crypt_reg);
   mongocrypt_destroy (crypt_ecb);
   mongocrypt_status_destroy (status);
}
#endif

void
_mongocrypt_tester_install_crypto_hooks (_mongocrypt_tester_t *tester)
{
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_encryption, CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_decryption, CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_iv_gen, CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_random, CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_kms_request, CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_unset, CRYPTO_PROHIBITED);
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_explicit_err, CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_crypto_hooks_explicit_sha256_err,
                        CRYPTO_OPTIONAL);
   INSTALL_TEST_CRYPTO (_test_crypto_hook_sign_rsaes_pkcs1_v1_5,
                        CRYPTO_OPTIONAL);
#ifdef MONGOCRYPT_ENABLE_CRYPTO_LIBCRYPTO
   INSTALL_TEST (_test_fle2_crypto_via_ecb_hook);
#endif
}
