#include <stdlib.h>

#include <stdint.h>

#include <assert.h>

#include "fuzzer/FuzzedDataProvider.h"

#define MBEDTLS_CHECK_RETURN_TYPICAL

#define MBEDTLS_INTERNAL_VALIDATE_RET( cond, ret )  do { } while( 0 )

#include "gcrypt/aes.h"

#include "gcrypt/cipher.h"

#include "gcrypt/gcm.h"

extern int force_no_aesni;

extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {

  FuzzedDataProvider fuzzed_data(data, size);

  mbedtls_aes_context *aes_e_ctx, *aes_d_ctx;

  mbedtls_gcm_context *gcm_e_ctx, *gcm_d_ctx;

  int key_lens[] = { 128, 192, 256, 512 /* invalid */ };

  unsigned char *output, *decrypted;

  int key_len, rc_e, rc_d;

  mbedtls_cipher_id_t cipher;

  unsigned char *tag;

  int iv_len, tag_len, ad_len, input_length, force_auth_tag_error;

  /* No real memory allocations involved */

  if(fuzzed_data.remaining_bytes() < 1 + 4 + 512 / 8 +

             1 + 64 + /* iv */

             1 + /* tag_len */

             1 + 17 + /* ad */

             1 + 64 + /* input */

             1 + /* force_auth_tag_error */

             1 /* useless data: to be able to add the check with assert */)

    return -1;

  gcm_e_ctx = (mbedtls_gcm_context *)malloc(sizeof(mbedtls_gcm_context));

  gcm_d_ctx = (mbedtls_gcm_context *)malloc(sizeof(mbedtls_gcm_context));

  aes_e_ctx = (mbedtls_aes_context *)malloc(sizeof(mbedtls_aes_context));

  aes_d_ctx = (mbedtls_aes_context *)malloc(sizeof(mbedtls_aes_context));

  force_no_aesni = 0;

  if(fuzzed_data.ConsumeBool())

    force_no_aesni = 1;

  key_len = fuzzed_data.PickValueInArray(key_lens);

  std::vector<unsigned char>key = fuzzed_data.ConsumeBytes<u_int8_t>(key_len / 8);

  iv_len = fuzzed_data.ConsumeIntegralInRange(0, 64);

  std::vector<u_int8_t>iv = fuzzed_data.ConsumeBytes<uint8_t>(iv_len);

  tag_len = fuzzed_data.ConsumeIntegralInRange(0, 17);

  tag = (unsigned char *)malloc(tag_len);

  ad_len = fuzzed_data.ConsumeIntegralInRange(0, 17);

  std::vector<u_int8_t>ad = fuzzed_data.ConsumeBytes<uint8_t>(ad_len);

  input_length = fuzzed_data.ConsumeIntegralInRange(16, 64);

  std::vector<unsigned char>input = fuzzed_data.ConsumeBytes<u_int8_t>(input_length);

  output = (unsigned char *)malloc(input_length);

  decrypted = (unsigned char *)malloc(input_length);

  force_auth_tag_error = fuzzed_data.ConsumeBool();

  cipher = static_cast<mbedtls_cipher_id_t>(fuzzed_data.ConsumeIntegralInRange(0, (int)MBEDTLS_CIPHER_ID_CHACHA20));

  assert(fuzzed_data.remaining_bytes() > 0);

  mbedtls_gcm_init(gcm_e_ctx, aes_e_ctx);

  mbedtls_gcm_init(gcm_d_ctx, aes_d_ctx);

  rc_e = mbedtls_gcm_setkey(gcm_e_ctx, cipher, key.data(), key.size() * 8);

  rc_d = mbedtls_gcm_setkey(gcm_d_ctx, cipher, key.data(), key.size() * 8);

  if (rc_e == 0 && rc_d == 0) {

    rc_e = mbedtls_gcm_crypt_and_tag(gcm_e_ctx, MBEDTLS_GCM_ENCRYPT,

             input.size(),

             iv.data(), iv.size(),

             ad.data(), ad.size(),

             input.data(),

             output,

             tag_len, tag);

    if(rc_e == 0) {

      if(force_auth_tag_error && tag_len > 0 && tag[0] != 0) {

        tag[0] = 0;

      } else {

        force_auth_tag_error = 0;

      }

      rc_d = mbedtls_gcm_auth_decrypt(gcm_d_ctx,

              input.size(),

              iv.data(), iv.size(),

              ad.data(), ad.size(),

              tag, tag_len,

              output,

              decrypted);

      if(rc_d == 0)

        assert(memcmp(input.data(), decrypted, input.size()) == 0);

      if(force_auth_tag_error)

        assert(rc_d == MBEDTLS_ERR_GCM_AUTH_FAILED);

    }

  }

  mbedtls_gcm_free(gcm_e_ctx);

  mbedtls_gcm_free(gcm_d_ctx);

  free(tag);

  free(gcm_e_ctx);

  free(gcm_d_ctx);

  free(aes_e_ctx);

  free(aes_d_ctx);

  free(output);

  free(decrypted);

  return 0;

}