#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/cipher.h"
#include "gcrypt/error.h"
#include "gcrypt/aes.h"

extern int force_no_aesni;

extern "C" int LLVMFuzzerTestOneInput(const uint8_t *data, size_t size) {
  FuzzedDataProvider fuzzed_data(data, size);
  int key_lens[] = { 128, 192, 256 };
  int key_len, iv_len, rc_e, rc_d, input_length;
  unsigned char *output, *decrypted;
  size_t output_size, output_size2, decrypted_size;
  mbedtls_cipher_type_t cipher_type;
  /* TODO: GCM. This code/fuzzer doesn't work with GCM ciphers. Not sure why.. :( */
  const char *cipher_names[] = { NULL, "", "AES-128-ECB", "AES-192-ECB", "AES-256-ECB",
                                 /* "AES-128-GCM", "AES-192-GCM", "AES-256-GCM" */ };
  const char *cipher_name;
  mbedtls_cipher_context_t *ctx_e, *ctx_d;

  /* No real memory allocations involved */

  if(fuzzed_data.remaining_bytes() < 512) /* Some data */
    return -1;

  posix_memalign((void **)&ctx_e, 8, sizeof(mbedtls_cipher_context_t));
  posix_memalign((void **)&ctx_d, 8, sizeof(mbedtls_cipher_context_t));

  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, MBEDTLS_MAX_IV_LENGTH + 1);
  std::vector<u_int8_t>iv = fuzzed_data.ConsumeBytes<uint8_t>(iv_len);
  input_length = fuzzed_data.ConsumeIntegralInRange(16, 17);
  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);

  mbedtls_cipher_list();
  /* Random iteration */
  cipher_type = static_cast<mbedtls_cipher_type_t>(fuzzed_data.ConsumeIntegralInRange(0, (int)MBEDTLS_CIPHER_AES_256_KWP) + 1);
  mbedtls_cipher_info_from_type(cipher_type);

  /* Real cipher used */
  cipher_name = cipher_names[fuzzed_data.ConsumeIntegralInRange(0, (int)(sizeof(cipher_names) / sizeof(char *) - 1))];
  mbedtls_cipher_init(ctx_e);
  mbedtls_cipher_init(ctx_d);
  ctx_e->cipher_info = mbedtls_cipher_info_from_string(cipher_name);
  ctx_d->cipher_info = ctx_e->cipher_info;

  mbedtls_cipher_info_get_mode(ctx_e->cipher_info);
  mbedtls_cipher_info_get_type(ctx_e->cipher_info);
  mbedtls_cipher_info_get_name(ctx_e->cipher_info);
  mbedtls_cipher_info_has_variable_key_bitlen(ctx_e->cipher_info);
  mbedtls_cipher_info_has_variable_iv_size(ctx_e->cipher_info);
  mbedtls_cipher_info_get_iv_size(ctx_e->cipher_info);
  mbedtls_cipher_info_get_block_size(ctx_e->cipher_info);
  mbedtls_cipher_get_cipher_mode(ctx_e);
  mbedtls_cipher_get_iv_size(ctx_e);
  mbedtls_cipher_get_type(ctx_e);
  mbedtls_cipher_get_name(ctx_e);
  mbedtls_cipher_get_key_bitlen(ctx_e);
  mbedtls_cipher_get_operation(ctx_e);
  mbedtls_cipher_info_get_key_bitlen(ctx_e->cipher_info);
  mbedtls_error_add(0, 0, NULL, 0);

  posix_memalign((void **)&ctx_e->cipher_ctx, 8, sizeof(mbedtls_aes_context));
  posix_memalign((void **)&ctx_d->cipher_ctx, 8, sizeof(mbedtls_aes_context));

  rc_e = mbedtls_cipher_setkey(ctx_e, key.data(), key.size() * 8, MBEDTLS_ENCRYPT);
  rc_d = mbedtls_cipher_setkey(ctx_d, key.data(), key.size() * 8, MBEDTLS_DECRYPT);
  if(rc_e == 0 && rc_d == 0) {

    if(fuzzed_data.ConsumeBool()) {
      rc_e = mbedtls_cipher_crypt(ctx_e, iv.data(), iv.size(),
                                  input.data(), input.size(), output, &output_size);
    } else {
      rc_e = mbedtls_cipher_set_iv(ctx_e, iv.data(), iv.size());
      rc_d = mbedtls_cipher_set_iv(ctx_d, iv.data(), iv.size());
      if(rc_e == 0 && rc_d == 0) {
        mbedtls_cipher_reset(ctx_e);
        mbedtls_cipher_reset(ctx_d);

        rc_e = mbedtls_cipher_update(ctx_e, input.data(), input.size(), output, &output_size);
        if(rc_e == 0) {
	  rc_e = mbedtls_cipher_finish(ctx_e, NULL, &output_size2);
          if(rc_e == 0) {

            rc_d = mbedtls_cipher_update(ctx_d, output, output_size, decrypted, &decrypted_size);
            if(rc_d == 0) {
              rc_d = mbedtls_cipher_finish(ctx_d, NULL, &output_size2);
              /* TODO: decryption doesn't work with no-aesni data path!
	         Note that with MASAN, aesni is always disabled */
#if 0
	      if(rc_d == 0) {
                assert(input.size() == decrypted_size);
                assert(memcmp(input.data(), decrypted, decrypted_size) == 0);
              }
#endif
	    }
          }
        }
      }
    }
  }

  free(output);
  free(decrypted);
  free(ctx_e->cipher_ctx);
  free(ctx_e);
  free(ctx_d->cipher_ctx);
  free(ctx_d);
  return 0;
}