#include "crypto/crypto_aes.h"
#include "async_wrap-inl.h"
#include "base_object-inl.h"
#include "crypto/crypto_cipher.h"
#include "crypto/crypto_keys.h"
#include "crypto/crypto_util.h"
#include "env-inl.h"
#include "memory_tracker-inl.h"
#include "threadpoolwork-inl.h"
#include "v8.h"

#include <openssl/bn.h>
#include <openssl/aes.h>

#include <vector>

namespace node {

using ncrypto::BignumPointer;
using ncrypto::Cipher;
using ncrypto::CipherCtxPointer;
using ncrypto::DataPointer;
using v8::FunctionCallbackInfo;
using v8::Just;
using v8::JustVoid;
using v8::Local;
using v8::Maybe;
using v8::Nothing;
using v8::Object;
using v8::Uint32;
using v8::Value;

namespace crypto {
namespace {
constexpr size_t kAesBlockSize = 16;
constexpr const char* kDefaultWrapIV = "\xa6\xa6\xa6\xa6\xa6\xa6\xa6\xa6";

// Implements general AES encryption and decryption for CBC
// The key_data must be a secret key.
// On success, this function sets out to a new ByteSource
// instance containing the results and returns WebCryptoCipherStatus::OK.
WebCryptoCipherStatus AES_Cipher(Environment* env,
                                 const KeyObjectData& key_data,
                                 WebCryptoCipherMode cipher_mode,
                                 const AESCipherConfig& params,
                                 const ByteSource& in,
                                 ByteSource* out) {
  CHECK_EQ(key_data.GetKeyType(), kKeyTypeSecret);

  const bool encrypt = cipher_mode == kWebCryptoCipherEncrypt;
  // AES-KW requires at least two 64-bit plaintext blocks, plus the
  // 64-bit integrity check value when unwrapping.
  if (params.cipher.isWrapMode() && in.size() < (encrypt ? 16u : 24u)) {
    return WebCryptoCipherStatus::FAILED;
  }

  auto ctx = CipherCtxPointer::New();
  if (!ctx) {
    return WebCryptoCipherStatus::FAILED;
  }

  if (params.cipher.isWrapMode()) {
    ctx.setAllowWrap();
  }

  if (!ctx.init(params.cipher, encrypt)) {
    // Cipher init failed
    return WebCryptoCipherStatus::FAILED;
  }

  if ((params.cipher.isGcmMode() || params.cipher.isOcbMode()) &&
      !ctx.setIvLength(params.iv.size())) {
    return WebCryptoCipherStatus::FAILED;
  }

  if (!ctx.setKeyLength(key_data.GetSymmetricKeySize()) ||
      !ctx.init(
          Cipher(),
          encrypt,
          reinterpret_cast<const unsigned char*>(key_data.GetSymmetricKey()),
          params.iv.data<unsigned char>())) {
    return WebCryptoCipherStatus::FAILED;
  }

  size_t tag_len = 0;
  size_t data_len = in.size();

  if (params.cipher.isGcmMode() || params.cipher.isOcbMode()) {
    tag_len = params.length;
    switch (cipher_mode) {
      case kWebCryptoCipherDecrypt: {
        // In decrypt mode, the auth tag is appended to the end of the
        // ciphertext. Split it off and set it on the cipher context.
        if (data_len < tag_len) {
          return WebCryptoCipherStatus::FAILED;
        }
        data_len -= tag_len;

        if (params.cipher.isOcbMode()) {
          if (!ctx.setAeadTagLength(tag_len)) {
            return WebCryptoCipherStatus::FAILED;
          }
        }

        ncrypto::Buffer<const char> buffer = {
            .data = in.data<char>() + data_len,
            .len = tag_len,
        };
        if (!ctx.setAeadTag(buffer)) {
          return WebCryptoCipherStatus::FAILED;
        }
        break;
      }
      case kWebCryptoCipherEncrypt: {
        if (params.cipher.isOcbMode()) {
          if (!ctx.setAeadTagLength(tag_len)) {
            return WebCryptoCipherStatus::FAILED;
          }
        }
        break;
      }
      default:
        UNREACHABLE();
    }
  }

  size_t total = 0;
  const int block_size = ctx.getBlockSize();
  if (block_size < 0) {
    return WebCryptoCipherStatus::FAILED;
  }
  int buf_len;
  if (!TryGetIntCipherOutputLength(
          data_len,
          static_cast<size_t>(block_size) + (encrypt ? tag_len : 0),
          &buf_len)) {
    return WebCryptoCipherStatus::FAILED;
  }
  int out_len;

  ncrypto::Buffer<const unsigned char> buffer = {
      .data = params.additional_data.data<unsigned char>(),
      .len = params.additional_data.size(),
  };
  if ((params.cipher.isGcmMode() || params.cipher.isOcbMode()) &&
      params.additional_data.size() && !ctx.update(buffer, nullptr, &out_len)) {
    return WebCryptoCipherStatus::FAILED;
  }

  auto buf = DataPointer::Alloc(buf_len);
  auto ptr = static_cast<unsigned char*>(buf.get());

  // Some shared OpenSSL builds fail when the input size is zero. Keep handling
  // zero-length input in Node.js to avoid relying on backend-specific behavior.
  //
  // Refs:
  // https://github.com/openssl/openssl/commit/420cb707b880e4fb649094241371701013eeb15f
  // Refs: https://github.com/nodejs/node/pull/38913#issuecomment-866505244
  buffer = {
      .data = in.data<unsigned char>(),
      .len = data_len,
  };
  if (data_len == 0) {
    out_len = 0;
  } else if (!ctx.update(buffer, ptr, &out_len)) {
    return WebCryptoCipherStatus::FAILED;
  }

  total += out_len;
  CHECK_LE(out_len, buf_len);
  out_len = block_size;
  if (!ctx.update({}, ptr + total, &out_len, true)) {
    return WebCryptoCipherStatus::FAILED;
  }
  total += out_len;

  // If using AES_GCM or AES_OCB, grab the generated auth tag and append
  // it to the end of the ciphertext.
  if (encrypt && (params.cipher.isGcmMode() || params.cipher.isOcbMode())) {
    if (!ctx.getAeadTag(tag_len, ptr + total)) {
      return WebCryptoCipherStatus::FAILED;
    }
    total += tag_len;
  }

  if (total == 0) {
    *out = ByteSource::Allocated(nullptr, 0);
    return WebCryptoCipherStatus::OK;
  }

  // It's possible that we haven't used the full allocated space. Size down.
  buf = buf.resize(total);
  *out = ByteSource::Allocated(buf.release());

  return WebCryptoCipherStatus::OK;
}

#ifdef OPENSSL_IS_BORINGSSL
// AES Key Wrap using BoringSSL's low-level AES_wrap_key / AES_unwrap_key.
// BoringSSL does not expose EVP_aes_*_wrap via the
// EVP_CIPHER registry, so the EVP-based AES_Cipher path is unusable for
// AES-KW. This matches Chromium's WebCrypto AES-KW implementation.
WebCryptoCipherStatus AES_KW_Cipher(Environment* env,
                                    const KeyObjectData& key_data,
                                    WebCryptoCipherMode cipher_mode,
                                    const AESCipherConfig& params,
                                    const ByteSource& in,
                                    ByteSource* out) {
  CHECK_EQ(key_data.GetKeyType(), kKeyTypeSecret);

  const unsigned key_bits =
      static_cast<unsigned>(key_data.GetSymmetricKeySize()) * 8;
  const auto key_bytes =
      reinterpret_cast<const unsigned char*>(key_data.GetSymmetricKey());
  const bool encrypt = cipher_mode == kWebCryptoCipherEncrypt;

  AES_KEY aes_key;
  if (encrypt) {
    // Input must be a multiple of 8 bytes and at least 16 bytes.
    if (in.size() < 16 || in.size() % 8 != 0) {
      return WebCryptoCipherStatus::FAILED;
    }
    if (AES_set_encrypt_key(key_bytes, key_bits, &aes_key) != 0) {
      return WebCryptoCipherStatus::FAILED;
    }
    auto buf = DataPointer::Alloc(in.size() + 8);
    int len = AES_wrap_key(&aes_key,
                           nullptr,
                           static_cast<unsigned char*>(buf.get()),
                           in.data<unsigned char>(),
                           in.size());
    if (len < 0 || static_cast<size_t>(len) != in.size() + 8) {
      return WebCryptoCipherStatus::FAILED;
    }
    *out = ByteSource::Allocated(buf.release());
  } else {
    // Input must be a multiple of 8 bytes and at least 24 bytes.
    if (in.size() < 24 || in.size() % 8 != 0) {
      return WebCryptoCipherStatus::FAILED;
    }
    if (AES_set_decrypt_key(key_bytes, key_bits, &aes_key) != 0) {
      return WebCryptoCipherStatus::FAILED;
    }
    auto buf = DataPointer::Alloc(in.size() - 8);
    int len = AES_unwrap_key(&aes_key,
                             nullptr,
                             static_cast<unsigned char*>(buf.get()),
                             in.data<unsigned char>(),
                             in.size());
    if (len < 0 || static_cast<size_t>(len) != in.size() - 8) {
      return WebCryptoCipherStatus::FAILED;
    }
    *out = ByteSource::Allocated(buf.release());
  }

  return WebCryptoCipherStatus::OK;
}
#endif  // OPENSSL_IS_BORINGSSL

// The AES_CTR implementation here takes it's inspiration from the chromium
// implementation here:
// https://github.com/chromium/chromium/blob/7af6cfd/components/webcrypto/algorithms/aes_ctr.cc

template <std::integral T>
T CeilDiv(T a, T b) {
  return a == 0 ? 0 : 1 + (a - 1) / b;
}

BignumPointer GetCounter(const AESCipherConfig& params) {
  unsigned int remainder = (params.length % CHAR_BIT);
  const unsigned char* data = params.iv.data<unsigned char>();

  if (remainder == 0) {
    unsigned int byte_length = params.length / CHAR_BIT;
    return BignumPointer(data + params.iv.size() - byte_length, byte_length);
  }

  unsigned int byte_length =
      CeilDiv(params.length, static_cast<size_t>(CHAR_BIT));

  std::vector<unsigned char> counter(
      data + params.iv.size() - byte_length,
      data + params.iv.size());
  counter[0] &= ~(0xFF << remainder);

  return BignumPointer(counter.data(), counter.size());
}

std::vector<unsigned char> BlockWithZeroedCounter(
    const AESCipherConfig& params) {
  unsigned int length_bytes = params.length / CHAR_BIT;
  unsigned int remainder = params.length % CHAR_BIT;

  const unsigned char* data = params.iv.data<unsigned char>();

  std::vector<unsigned char> new_counter_block(data, data + params.iv.size());

  size_t index = new_counter_block.size() - length_bytes;
  memset(&new_counter_block.front() + index, 0, length_bytes);

  if (remainder)
    new_counter_block[index - 1] &= 0xFF << remainder;

  return new_counter_block;
}

WebCryptoCipherStatus AES_CTR_Cipher2(const KeyObjectData& key_data,
                                      WebCryptoCipherMode cipher_mode,
                                      const AESCipherConfig& params,
                                      const ByteSource& in,
                                      unsigned const char* counter,
                                      unsigned char* out) {
  auto ctx = CipherCtxPointer::New();
  if (!ctx) {
    return WebCryptoCipherStatus::FAILED;
  }

  if (!ctx.init(
          params.cipher,
          cipher_mode == kWebCryptoCipherEncrypt,
          reinterpret_cast<const unsigned char*>(key_data.GetSymmetricKey()),
          counter)) {
    // Cipher init failed
    return WebCryptoCipherStatus::FAILED;
  }

  int out_len = 0;
  int final_len = 0;
  ncrypto::Buffer<const unsigned char> buffer = {
      .data = in.data<unsigned char>(),
      .len = in.size(),
  };
  if (!ctx.update(buffer, out, &out_len) ||
      !ctx.update({}, out + out_len, &final_len, true)) {
    return WebCryptoCipherStatus::FAILED;
  }

  return static_cast<unsigned>(out_len + final_len) != in.size()
             ? WebCryptoCipherStatus::FAILED
             : WebCryptoCipherStatus::OK;
}

WebCryptoCipherStatus AES_CTR_Cipher(Environment* env,
                                     const KeyObjectData& key_data,
                                     WebCryptoCipherMode cipher_mode,
                                     const AESCipherConfig& params,
                                     const ByteSource& in,
                                     ByteSource* out) {
  auto num_counters = BignumPointer::NewLShift(params.length);
  if (!num_counters) return WebCryptoCipherStatus::FAILED;

  BignumPointer current_counter = GetCounter(params);

  auto num_output = BignumPointer::New();

  if (!num_output.setWord(CeilDiv(in.size(), kAesBlockSize))) {
    return WebCryptoCipherStatus::FAILED;
  }

  // Just like in chromium's implementation, if the counter will
  // be incremented more than there are counter values, we fail.
  if (num_output > num_counters) return WebCryptoCipherStatus::FAILED;

  auto remaining_until_reset =
      BignumPointer::NewSub(num_counters, current_counter);
  if (!remaining_until_reset) {
    return WebCryptoCipherStatus::FAILED;
  }

  // Output size is identical to the input size.
  auto buf = DataPointer::Alloc(in.size());

  // Also just like in chromium's implementation, if we can process
  // the input without wrapping the counter, we'll do it as a single
  // call here. If we can't, we'll fallback to the a two-step approach
  if (remaining_until_reset >= num_output) {
    auto status = AES_CTR_Cipher2(key_data,
                                  cipher_mode,
                                  params,
                                  in,
                                  params.iv.data<unsigned char>(),
                                  static_cast<unsigned char*>(buf.get()));
    if (status == WebCryptoCipherStatus::OK) {
      *out = ByteSource::Allocated(buf.release());
    }
    return status;
  }

  std::optional<BN_ULONG> remaining_blocks = remaining_until_reset.getWord();
  CHECK(remaining_blocks.has_value());
  BN_ULONG input_size_part1 = remaining_blocks.value() * kAesBlockSize;

  // Encrypt the first part...
  auto status =
      AES_CTR_Cipher2(key_data,
                      cipher_mode,
                      params,
                      ByteSource::Foreign(in.data<char>(), input_size_part1),
                      params.iv.data<unsigned char>(),
                      static_cast<unsigned char*>(buf.get()));

  if (status != WebCryptoCipherStatus::OK) {
    return status;
  }

  // Wrap the counter around to zero
  std::vector<unsigned char> new_counter_block = BlockWithZeroedCounter(params);

  auto ptr = static_cast<unsigned char*>(buf.get()) + input_size_part1;
  // Encrypt the second part...
  status =
      AES_CTR_Cipher2(key_data,
                      cipher_mode,
                      params,
                      ByteSource::Foreign(in.data<char>() + input_size_part1,
                                          in.size() - input_size_part1),
                      new_counter_block.data(),
                      ptr);

  if (status == WebCryptoCipherStatus::OK) {
    *out = ByteSource::Allocated(buf.release());
  }

  return status;
}

bool ValidateIV(
    Environment* env,
    CryptoJobMode mode,
    Local<Value> value,
    AESCipherConfig* params) {
  ArrayBufferOrViewContents<char> iv(value);
  if (!iv.CheckSizeInt32()) [[unlikely]] {
    THROW_ERR_OUT_OF_RANGE(env, "iv is too big");
    return false;
  }
  params->iv = (IsCryptoJobAsync(mode)) ? iv.ToCopy() : iv.ToByteSource();
  return true;
}

bool ValidateCounter(
  Environment* env,
  Local<Value> value,
  AESCipherConfig* params) {
  CHECK(value->IsUint32());  // Length
  params->length = value.As<Uint32>()->Value();
  if (params->iv.size() != 16 ||
      params->length == 0 ||
      params->length > 128) {
    THROW_ERR_CRYPTO_INVALID_COUNTER(env);
    return false;
  }
  return true;
}

bool ValidateAuthTag(Environment* env,
                     Local<Value> value,
                     AESCipherConfig* params) {
  if (!value->IsUint32()) {
    THROW_ERR_CRYPTO_INVALID_TAG_LENGTH(env);
    return false;
  }
  params->length = value.As<Uint32>()->Value();
  if (params->length > 128) {
    THROW_ERR_CRYPTO_INVALID_TAG_LENGTH(env);
    return false;
  }
  return true;
}

bool ValidateAdditionalData(
    Environment* env,
    CryptoJobMode mode,
    Local<Value> value,
    AESCipherConfig* params) {
  // Additional Data
  if (IsAnyBufferSource(value)) {
    ArrayBufferOrViewContents<char> additional(value);
    if (!additional.CheckSizeInt32()) [[unlikely]] {
      THROW_ERR_OUT_OF_RANGE(env, "additionalData is too big");
      return false;
    }
    params->additional_data = IsCryptoJobAsync(mode)
                                  ? additional.ToCopy()
                                  : additional.ToByteSource();
  }
  return true;
}

void UseDefaultIV(AESCipherConfig* params) {
  params->iv = ByteSource::Foreign(kDefaultWrapIV, strlen(kDefaultWrapIV));
}
}  // namespace

AESCipherConfig::AESCipherConfig(AESCipherConfig&& other) noexcept
    : variant(other.variant),
      cipher(other.cipher),
      length(other.length),
      iv(std::move(other.iv)),
      additional_data(std::move(other.additional_data)) {}

AESCipherConfig& AESCipherConfig::operator=(AESCipherConfig&& other) noexcept {
  if (&other == this) return *this;
  this->~AESCipherConfig();
  return *new (this) AESCipherConfig(std::move(other));
}

void AESCipherConfig::MemoryInfo(MemoryTracker* tracker) const {
  tracker->TraitTrackInline(iv, "iv");
  tracker->TraitTrackInline(additional_data, "additional_data");
}

Maybe<void> AESCipherTraits::AdditionalConfig(
    CryptoJobMode mode,
    const FunctionCallbackInfo<Value>& args,
    unsigned int offset,
    WebCryptoCipherMode cipher_mode,
    AESCipherConfig* params) {
  Environment* env = Environment::GetCurrent(args);

  CHECK(args[offset]->IsUint32());  // Key Variant
  params->variant =
      static_cast<AESKeyVariant>(args[offset].As<Uint32>()->Value());

#define V(name, _, nid)                                                        \
  case AESKeyVariant::name: {                                                  \
    params->cipher = nid;                                                      \
    break;                                                                     \
  }
  switch (params->variant) {
    VARIANTS(V)
    default:
      UNREACHABLE();
  }
#undef V

#ifdef OPENSSL_IS_BORINGSSL
  // On BoringSSL the KW variants have no backing EVP_CIPHER; they use
  // low-level AES_wrap_key / AES_unwrap_key instead.
  const bool is_kw = params->variant == AESKeyVariant::KW_128 ||
                     params->variant == AESKeyVariant::KW_192 ||
                     params->variant == AESKeyVariant::KW_256;

  if (is_kw) {
    UseDefaultIV(params);
    return JustVoid();
  }
#endif

  if (!params->cipher) {
    THROW_ERR_CRYPTO_UNKNOWN_CIPHER(env);
    return Nothing<void>();
  }

  if (!params->cipher.isWrapMode()) {
    if (!ValidateIV(env, mode, args[offset + 1], params)) {
      return Nothing<void>();
    }
    if (params->cipher.isCtrMode()) {
      if (!ValidateCounter(env, args[offset + 2], params)) {
        return Nothing<void>();
      }
    } else if (params->cipher.isGcmMode() || params->cipher.isOcbMode()) {
      if (!ValidateAuthTag(env, args[offset + 2], params) ||
          !ValidateAdditionalData(env, mode, args[offset + 3], params)) {
        return Nothing<void>();
      }
    }
  } else {
    UseDefaultIV(params);
  }

  if (params->cipher.isGcmMode()) {
    if (params->iv.size() == 0) {
      THROW_ERR_CRYPTO_INVALID_IV(env);
      return Nothing<void>();
    }
  } else if (params->cipher.isOcbMode()) {
    // For OCB mode, allow variable IV lengths (1-15 bytes).
    if (params->iv.size() == 0 || params->iv.size() > 15) {
      THROW_ERR_CRYPTO_INVALID_IV(env);
      return Nothing<void>();
    }
  } else {
    // For other modes, check against the cipher's expected IV length
    if (params->iv.size() < static_cast<size_t>(params->cipher.getIvLength())) {
      THROW_ERR_CRYPTO_INVALID_IV(env);
      return Nothing<void>();
    }
  }

  return JustVoid();
}

WebCryptoCipherStatus AESCipherTraits::DoCipher(Environment* env,
                                                const KeyObjectData& key_data,
                                                WebCryptoCipherMode cipher_mode,
                                                const AESCipherConfig& params,
                                                const ByteSource& in,
                                                ByteSource* out) {
#define V(name, fn, _)                                                         \
  case AESKeyVariant::name:                                                    \
    return fn(env, key_data, cipher_mode, params, in, out);
  switch (params.variant) {
    VARIANTS(V)
    default:
      UNREACHABLE();
  }
#undef V
}

void AES::Initialize(Environment* env, Local<Object> target) {
  AESCryptoJob::Initialize(env, target);

#define V(name, _, __)                                                         \
  constexpr static auto kKeyVariantAES_##name =                                \
      static_cast<int>(AESKeyVariant::name);                                   \
  NODE_DEFINE_CONSTANT(target, kKeyVariantAES_##name);

  VARIANTS(V)

#undef V
}

void AES::RegisterExternalReferences(ExternalReferenceRegistry* registry) {
  AESCryptoJob::RegisterExternalReferences(registry);
}

}  // namespace crypto
}  // namespace node