#include "crypto/crypto_kem.h"

#if OPENSSL_VERSION_MAJOR >= 3

#include "async_wrap-inl.h"

#include "base_object-inl.h"

#include "crypto/crypto_keys.h"

#include "crypto/crypto_util.h"

#include "env-inl.h"

#include "memory_tracker-inl.h"

#include "node_buffer.h"

#include "threadpoolwork-inl.h"

#include "v8.h"

namespace node {

using ncrypto::EVPKeyPointer;

using v8::Array;

using v8::FunctionCallbackInfo;

using v8::Local;

using v8::Maybe;

using v8::MaybeLocal;

using v8::Nothing;

using v8::Object;

using v8::Value;

namespace crypto {

KEMConfiguration::KEMConfiguration(KEMConfiguration&& other) noexcept

    : job_mode(other.job_mode),

      mode(other.mode),

      key(std::move(other.key)),

      ciphertext(std::move(other.ciphertext)) {}

KEMConfiguration& KEMConfiguration::operator=(

    KEMConfiguration&& other) noexcept {

  if (&other == this) return *this;

  this->~KEMConfiguration();

  return *new (this) KEMConfiguration(std::move(other));

}

void KEMConfiguration::MemoryInfo(MemoryTracker* tracker) const {

  tracker->TrackField("key", key);

  if (job_mode == kCryptoJobAsync) {

    tracker->TrackFieldWithSize("ciphertext", ciphertext.size());

  }

}

namespace {

bool DoKEMEncapsulate(Environment* env,

                      const EVPKeyPointer& public_key,

                      ByteSource* out,

                      CryptoJobMode mode) {

  auto result = ncrypto::KEM::Encapsulate(public_key);

  if (!result) {

    if (mode == kCryptoJobSync) {

      THROW_ERR_CRYPTO_OPERATION_FAILED(env, "Failed to perform encapsulation");

    }

    return false;

  }

  // Pack the result: [ciphertext_len][shared_key_len][ciphertext][shared_key]

  size_t ciphertext_len = result->ciphertext.size();

  size_t shared_key_len = result->shared_key.size();

  size_t total_len =

      sizeof(uint32_t) + sizeof(uint32_t) + ciphertext_len + shared_key_len;

  auto data = ncrypto::DataPointer::Alloc(total_len);

  if (!data) {

    if (mode == kCryptoJobSync) {

      THROW_ERR_CRYPTO_OPERATION_FAILED(env,

                                        "Failed to allocate output buffer");

    }

    return false;

  }

  unsigned char* ptr = static_cast<unsigned char*>(data.get());

  // Write size headers

  *reinterpret_cast<uint32_t*>(ptr) = static_cast<uint32_t>(ciphertext_len);

  *reinterpret_cast<uint32_t*>(ptr + sizeof(uint32_t)) =

      static_cast<uint32_t>(shared_key_len);

  // Write ciphertext and shared key data

  unsigned char* ciphertext_ptr = ptr + 2 * sizeof(uint32_t);

  unsigned char* shared_key_ptr = ciphertext_ptr + ciphertext_len;

  std::memcpy(ciphertext_ptr, result->ciphertext.get(), ciphertext_len);

  std::memcpy(shared_key_ptr, result->shared_key.get(), shared_key_len);

  *out = ByteSource::Allocated(data.release());

  return true;

}

bool DoKEMDecapsulate(Environment* env,

                      const EVPKeyPointer& private_key,

                      const ByteSource& ciphertext,

                      ByteSource* out,

                      CryptoJobMode mode) {

  ncrypto::Buffer<const void> ciphertext_buf{ciphertext.data(),

                                             ciphertext.size()};

  auto shared_key = ncrypto::KEM::Decapsulate(private_key, ciphertext_buf);

  if (!shared_key) {

    if (mode == kCryptoJobSync) {

      THROW_ERR_CRYPTO_OPERATION_FAILED(env, "Failed to perform decapsulation");

    }

    return false;

  }

  *out = ByteSource::Allocated(shared_key.release());

  return true;

}

}  // anonymous namespace

// KEMEncapsulateTraits implementation

Maybe<void> KEMEncapsulateTraits::AdditionalConfig(

    CryptoJobMode mode,

    const FunctionCallbackInfo<Value>& args,

    unsigned int offset,

    KEMConfiguration* params) {

  params->job_mode = mode;

  params->mode = KEMMode::Encapsulate;

  unsigned int key_offset = offset;

  auto public_key_data =

      KeyObjectData::GetPublicOrPrivateKeyFromJs(args, &key_offset);

  if (!public_key_data) {

    return Nothing<void>();

  }

  params->key = std::move(public_key_data);

  return v8::JustVoid();

}

bool KEMEncapsulateTraits::DeriveBits(Environment* env,

                                      const KEMConfiguration& params,

                                      ByteSource* out,

                                      CryptoJobMode mode) {

  Mutex::ScopedLock lock(params.key.mutex());

  const auto& public_key = params.key.GetAsymmetricKey();

  return DoKEMEncapsulate(env, public_key, out, mode);

}

MaybeLocal<Value> KEMEncapsulateTraits::EncodeOutput(

    Environment* env, const KEMConfiguration& params, ByteSource* out) {

  // The output contains:

  // [ciphertext_len][shared_key_len][ciphertext][shared_key]

  const unsigned char* data = out->data<unsigned char>();

  uint32_t ciphertext_len = *reinterpret_cast<const uint32_t*>(data);

  uint32_t shared_key_len =

      *reinterpret_cast<const uint32_t*>(data + sizeof(uint32_t));

  const unsigned char* ciphertext_ptr = data + 2 * sizeof(uint32_t);

  const unsigned char* shared_key_ptr = ciphertext_ptr + ciphertext_len;

  MaybeLocal<Object> ciphertext_buf =

      node::Buffer::Copy(env->isolate(),

                         reinterpret_cast<const char*>(ciphertext_ptr),

                         ciphertext_len);

  MaybeLocal<Object> shared_key_buf =

      node::Buffer::Copy(env->isolate(),

                         reinterpret_cast<const char*>(shared_key_ptr),

                         shared_key_len);

  Local<Object> ciphertext_obj;

  Local<Object> shared_key_obj;

  if (!ciphertext_buf.ToLocal(&ciphertext_obj) ||

      !shared_key_buf.ToLocal(&shared_key_obj)) {

    return MaybeLocal<Value>();

  }

  // Return an array [sharedKey, ciphertext].

  Local<Array> result = Array::New(env->isolate(), 2);

  if (result->Set(env->context(), 0, shared_key_obj).IsNothing() ||

      result->Set(env->context(), 1, ciphertext_obj).IsNothing()) {

    return MaybeLocal<Value>();

  }

  return result;

}

// KEMDecapsulateTraits implementation

Maybe<void> KEMDecapsulateTraits::AdditionalConfig(

    CryptoJobMode mode,

    const FunctionCallbackInfo<Value>& args,

    unsigned int offset,

    KEMConfiguration* params) {

  Environment* env = Environment::GetCurrent(args);

  params->job_mode = mode;

  params->mode = KEMMode::Decapsulate;

  unsigned int key_offset = offset;

  auto private_key_data =

      KeyObjectData::GetPrivateKeyFromJs(args, &key_offset, true);

  if (!private_key_data) {

    return Nothing<void>();

  }

  params->key = std::move(private_key_data);

  ArrayBufferOrViewContents<unsigned char> ciphertext(args[key_offset]);

  if (!ciphertext.CheckSizeInt32()) {

    THROW_ERR_OUT_OF_RANGE(env, "ciphertext is too big");

    return Nothing<void>();

  }

  params->ciphertext =

      mode == kCryptoJobAsync ? ciphertext.ToCopy() : ciphertext.ToByteSource();

  return v8::JustVoid();

}

bool KEMDecapsulateTraits::DeriveBits(Environment* env,

                                      const KEMConfiguration& params,

                                      ByteSource* out,

                                      CryptoJobMode mode) {

  Mutex::ScopedLock lock(params.key.mutex());

  const auto& private_key = params.key.GetAsymmetricKey();

  return DoKEMDecapsulate(env, private_key, params.ciphertext, out, mode);

}

MaybeLocal<Value> KEMDecapsulateTraits::EncodeOutput(

    Environment* env, const KEMConfiguration& params, ByteSource* out) {

  return out->ToBuffer(env);

}

void InitializeKEM(Environment* env, Local<Object> target) {

  KEMEncapsulateJob::Initialize(env, target);

  KEMDecapsulateJob::Initialize(env, target);

  constexpr int kKEMEncapsulate = static_cast<int>(KEMMode::Encapsulate);

  constexpr int kKEMDecapsulate = static_cast<int>(KEMMode::Decapsulate);

  NODE_DEFINE_CONSTANT(target, kKEMEncapsulate);

  NODE_DEFINE_CONSTANT(target, kKEMDecapsulate);

}

void RegisterKEMExternalReferences(ExternalReferenceRegistry* registry) {

  KEMEncapsulateJob::RegisterExternalReferences(registry);

  KEMDecapsulateJob::RegisterExternalReferences(registry);

}

namespace KEM {

void Initialize(Environment* env, Local<Object> target) {

  InitializeKEM(env, target);

}

void RegisterExternalReferences(ExternalReferenceRegistry* registry) {

  RegisterKEMExternalReferences(registry);

}

}  // namespace KEM

}  // namespace crypto

}  // namespace node

#endif