#include "test/extensions/transport_sockets/tls/test_private_key_method_provider.h"

#include <memory>

#include "envoy/api/api.h"

#include "openssl/ssl.h"

namespace Envoy {
namespace Extensions {
namespace PrivateKeyMethodProvider {

void TestPrivateKeyConnection::delayedOp() {
  const std::chrono::milliseconds timeout_0ms{0};

  timer_ = dispatcher_.createTimer([this]() -> void {
    finished_ = true;
    this->cb_.onPrivateKeyMethodComplete();
  });
  timer_->enableTimer(timeout_0ms);
}

static int calculateDigest(const EVP_MD* md, const uint8_t* in, size_t in_len, unsigned char* hash,
                           unsigned int* hash_len) {
  bssl::ScopedEVP_MD_CTX ctx;

  // Calculate the message digest for signing.
  if (!EVP_DigestInit_ex(ctx.get(), md, nullptr) || !EVP_DigestUpdate(ctx.get(), in, in_len) ||
      !EVP_DigestFinal_ex(ctx.get(), hash, hash_len)) {
    return 0;
  }
  return 1;
}

static ssl_private_key_result_t ecdsaPrivateKeySign(SSL* ssl, uint8_t* out, size_t* out_len,
                                                    size_t max_out, uint16_t signature_algorithm,
                                                    const uint8_t* in, size_t in_len) {
  unsigned char hash[EVP_MAX_MD_SIZE];
  unsigned int hash_len;
  TestPrivateKeyConnection* ops = static_cast<TestPrivateKeyConnection*>(
      SSL_get_ex_data(ssl, TestPrivateKeyMethodProvider::ecdsaConnectionIndex()));
  unsigned int out_len_unsigned;

  if (!ops) {
    return ssl_private_key_failure;
  }

  if (ops->test_options_.method_error_) {
    // Have an artificial test failure.
    return ssl_private_key_failure;
  }

  if (!ops->test_options_.sign_expected_) {
    return ssl_private_key_failure;
  }

  const EVP_MD* md = SSL_get_signature_algorithm_digest(signature_algorithm);
  if (!md) {
    return ssl_private_key_failure;
  }

  if (!calculateDigest(md, in, in_len, hash, &hash_len)) {
    return ssl_private_key_failure;
  }

  bssl::UniquePtr<EC_KEY> ec_key(EVP_PKEY_get1_EC_KEY(ops->getPrivateKey()));
  if (!ec_key) {
    return ssl_private_key_failure;
  }

  // Borrow "out" because it has been already initialized to the max_out size.
  if (!ECDSA_sign(0, hash, hash_len, out, &out_len_unsigned, ec_key.get())) {
    return ssl_private_key_failure;
  }

  if (ops->test_options_.sync_mode_) {
    // Return immediately with the results.
    if (out_len_unsigned > max_out) {
      return ssl_private_key_failure;
    }
    *out_len = out_len_unsigned;
    return ssl_private_key_success;
  }

  ops->output_.assign(out, out + out_len_unsigned);
  // Tell SSL socket that the operation is ready to be called again.
  ops->delayedOp();

  return ssl_private_key_retry;
}

static ssl_private_key_result_t ecdsaPrivateKeyDecrypt(SSL*, uint8_t*, size_t*, size_t,
                                                       const uint8_t*, size_t) {
  return ssl_private_key_failure;
}

static ssl_private_key_result_t rsaPrivateKeySign(SSL* ssl, uint8_t* out, size_t* out_len,
                                                  size_t max_out, uint16_t signature_algorithm,
                                                  const uint8_t* in, size_t in_len) {
  TestPrivateKeyConnection* ops = static_cast<TestPrivateKeyConnection*>(
      SSL_get_ex_data(ssl, TestPrivateKeyMethodProvider::rsaConnectionIndex()));
  unsigned char hash[EVP_MAX_MD_SIZE] = {0};
  unsigned int hash_len = EVP_MAX_MD_SIZE;
  std::vector<uint8_t> in2;

  if (!ops) {
    return ssl_private_key_failure;
  }

  if (ops->test_options_.method_error_) {
    return ssl_private_key_failure;
  }

  if (!ops->test_options_.sign_expected_) {
    return ssl_private_key_failure;
  }

  const EVP_MD* md = SSL_get_signature_algorithm_digest(signature_algorithm);
  if (!md) {
    return ssl_private_key_failure;
  }

  in2.assign(in, in + in_len);

  // If crypto error is set, we'll modify the incoming token by flipping
  // the bits.
  if (ops->test_options_.crypto_error_) {
    for (size_t i = 0; i < in_len; i++) {
      in2[i] = ~in2[i];
    }
  }

  if (!calculateDigest(md, in2.data(), in_len, hash, &hash_len)) {
    return ssl_private_key_failure;
  }

  RSA* rsa = EVP_PKEY_get0_RSA(ops->getPrivateKey());
  if (rsa == nullptr) {
    return ssl_private_key_failure;
  }

  // Perform RSA signing.
  if (SSL_is_signature_algorithm_rsa_pss(signature_algorithm)) {
    if (!RSA_sign_pss_mgf1(rsa, out_len, out, max_out, hash, hash_len, md, nullptr, -1)) {
      return ssl_private_key_failure;
    }
  } else {
    unsigned int out_len_unsigned;
    if (!RSA_sign(EVP_MD_type(md), hash, hash_len, out, &out_len_unsigned, rsa)) {
      return ssl_private_key_failure;
    }
    if (out_len_unsigned > max_out) {
      return ssl_private_key_failure;
    }
    *out_len = out_len_unsigned;
  }

  if (ops->test_options_.sync_mode_) {
    return ssl_private_key_success;
  }

  ops->output_.assign(out, out + *out_len);
  ops->delayedOp();

  return ssl_private_key_retry;
}

static ssl_private_key_result_t rsaPrivateKeyDecrypt(SSL* ssl, uint8_t* out, size_t* out_len,
                                                     size_t max_out, const uint8_t* in,
                                                     size_t in_len) {
  TestPrivateKeyConnection* ops = static_cast<TestPrivateKeyConnection*>(
      SSL_get_ex_data(ssl, TestPrivateKeyMethodProvider::rsaConnectionIndex()));

  if (!ops) {
    return ssl_private_key_failure;
  }

  if (ops->test_options_.method_error_) {
    return ssl_private_key_failure;
  }

  if (!ops->test_options_.decrypt_expected_) {
    return ssl_private_key_failure;
  }

  RSA* rsa = EVP_PKEY_get0_RSA(ops->getPrivateKey());
  if (rsa == nullptr) {
    return ssl_private_key_failure;
  }

  if (!RSA_decrypt(rsa, out_len, out, max_out, in, in_len, RSA_NO_PADDING)) {
    return ssl_private_key_failure;
  }

  if (ops->test_options_.sync_mode_) {
    return ssl_private_key_success;
  }

  ops->output_.assign(out, out + *out_len);
  ops->delayedOp();

  return ssl_private_key_retry;
}

static ssl_private_key_result_t privateKeyComplete(SSL* ssl, uint8_t* out, size_t* out_len,
                                                   size_t max_out, int id) {
  TestPrivateKeyConnection* ops = static_cast<TestPrivateKeyConnection*>(SSL_get_ex_data(ssl, id));

  if (!ops->finished_) {
    // The operation didn't finish yet, retry.
    return ssl_private_key_retry;
  }

  if (ops->test_options_.async_method_error_) {
    return ssl_private_key_failure;
  }

  if (ops->output_.size() > max_out) {
    return ssl_private_key_failure;
  }

  std::copy(ops->output_.begin(), ops->output_.end(), out);
  *out_len = ops->output_.size();

  return ssl_private_key_success;
}

static ssl_private_key_result_t rsaPrivateKeyComplete(SSL* ssl, uint8_t* out, size_t* out_len,
                                                      size_t max_out) {
  return privateKeyComplete(ssl, out, out_len, max_out,
                            TestPrivateKeyMethodProvider::rsaConnectionIndex());
}

static ssl_private_key_result_t ecdsaPrivateKeyComplete(SSL* ssl, uint8_t* out, size_t* out_len,
                                                        size_t max_out) {
  return privateKeyComplete(ssl, out, out_len, max_out,
                            TestPrivateKeyMethodProvider::ecdsaConnectionIndex());
}

Ssl::BoringSslPrivateKeyMethodSharedPtr
TestPrivateKeyMethodProvider::getBoringSslPrivateKeyMethod() {
  return method_;
}

bool TestPrivateKeyMethodProvider::checkFips() {
  if (mode_ == "rsa") {
    RSA* rsa_private_key = EVP_PKEY_get0_RSA(pkey_.get());
    if (rsa_private_key == nullptr || !RSA_check_fips(rsa_private_key)) {
      return false;
    }
  } else { // if (mode_ == "ecdsa")
    const EC_KEY* ecdsa_private_key = EVP_PKEY_get0_EC_KEY(pkey_.get());
    if (ecdsa_private_key == nullptr || !EC_KEY_check_fips(ecdsa_private_key)) {
      return false;
    }
  }
  return true;
}

bool TestPrivateKeyMethodProvider::isAvailable() { return test_options_.is_available_; }

TestPrivateKeyConnection::TestPrivateKeyConnection(
    Ssl::PrivateKeyConnectionCallbacks& cb, Event::Dispatcher& dispatcher,
    bssl::UniquePtr<EVP_PKEY> pkey, TestPrivateKeyConnectionTestOptions& test_options)
    : test_options_(test_options), cb_(cb), dispatcher_(dispatcher), pkey_(std::move(pkey)) {}

void TestPrivateKeyMethodProvider::registerPrivateKeyMethod(SSL* ssl,
                                                            Ssl::PrivateKeyConnectionCallbacks& cb,
                                                            Event::Dispatcher& dispatcher) {
  TestPrivateKeyConnection* ops;
  // In multi-cert case, when the same provider is used in different modes with the same SSL object,
  // we need to keep both rsa and ecdsa connection objects in store because the test options for the
  // two certificates may be different. We need to be able to deduct in the signing, decryption, and
  // completion functions which options to use, so we associate the connection objects to the same
  // SSL object using different user data indexes.
  //
  // Another way to do this would be to store both test options in one connection object.
  int index = mode_ == "rsa" ? TestPrivateKeyMethodProvider::rsaConnectionIndex()
                             : TestPrivateKeyMethodProvider::ecdsaConnectionIndex();

  // Check if there is another certificate of the same mode associated with the context. This would
  // be an error.
  ops = static_cast<TestPrivateKeyConnection*>(SSL_get_ex_data(ssl, index));
  if (ops != nullptr) {
    throw EnvoyException(
        "Can't distinguish between two registered providers for the same SSL object.");
  }

  ops = new TestPrivateKeyConnection(cb, dispatcher, bssl::UpRef(pkey_), test_options_);
  SSL_set_ex_data(ssl, index, ops);
}

void TestPrivateKeyMethodProvider::unregisterPrivateKeyMethod(SSL* ssl) {
  int index = mode_ == "rsa" ? TestPrivateKeyMethodProvider::rsaConnectionIndex()
                             : TestPrivateKeyMethodProvider::ecdsaConnectionIndex();
  TestPrivateKeyConnection* ops =
      static_cast<TestPrivateKeyConnection*>(SSL_get_ex_data(ssl, index));
  SSL_set_ex_data(ssl, index, nullptr);
  delete ops;
}

static int createIndex() {
  int index = SSL_get_ex_new_index(0, nullptr, nullptr, nullptr, nullptr);
  RELEASE_ASSERT(index >= 0, "Failed to get SSL user data index.");
  return index;
}

int TestPrivateKeyMethodProvider::rsaConnectionIndex() {
  CONSTRUCT_ON_FIRST_USE(int, createIndex());
}

int TestPrivateKeyMethodProvider::ecdsaConnectionIndex() {
  CONSTRUCT_ON_FIRST_USE(int, createIndex());
}

TestPrivateKeyMethodProvider::TestPrivateKeyMethodProvider(
    const ProtobufWkt::Any& typed_config,
    Server::Configuration::TransportSocketFactoryContext& factory_context) {
  std::string private_key_path;

  auto config = MessageUtil::anyConvert<ProtobufWkt::Struct>(typed_config);

  for (auto& value_it : config.fields()) {
    auto& value = value_it.second;
    if (value_it.first == "private_key_file" &&
        value.kind_case() == ProtobufWkt::Value::kStringValue) {
      private_key_path = value.string_value();
    }
    if (value_it.first == "sync_mode" && value.kind_case() == ProtobufWkt::Value::kBoolValue) {
      test_options_.sync_mode_ = value.bool_value();
    }
    if (value_it.first == "crypto_error" && value.kind_case() == ProtobufWkt::Value::kBoolValue) {
      test_options_.crypto_error_ = value.bool_value();
    }
    if (value_it.first == "method_error" && value.kind_case() == ProtobufWkt::Value::kBoolValue) {
      test_options_.method_error_ = value.bool_value();
    }
    if (value_it.first == "is_available" && value.kind_case() == ProtobufWkt::Value::kBoolValue) {
      test_options_.is_available_ = value.bool_value();
    }
    if (value_it.first == "async_method_error" &&
        value.kind_case() == ProtobufWkt::Value::kBoolValue) {
      test_options_.async_method_error_ = value.bool_value();
    }
    if (value_it.first == "expected_operation" &&
        value.kind_case() == ProtobufWkt::Value::kStringValue) {
      if (value.string_value() == "decrypt") {
        test_options_.decrypt_expected_ = true;
      } else if (value.string_value() == "sign") {
        test_options_.sign_expected_ = true;
      }
    }
    if (value_it.first == "mode" && value.kind_case() == ProtobufWkt::Value::kStringValue) {
      mode_ = value.string_value();
    }
  }

  if (!test_options_.is_available_) {
    return;
  }

  std::string private_key = factory_context.serverFactoryContext()
                                .api()
                                .fileSystem()
                                .fileReadToEnd(private_key_path)
                                .value();
  bssl::UniquePtr<BIO> bio(
      BIO_new_mem_buf(const_cast<char*>(private_key.data()), private_key.size()));
  bssl::UniquePtr<EVP_PKEY> pkey(PEM_read_bio_PrivateKey(bio.get(), nullptr, nullptr, nullptr));
  if (pkey == nullptr) {
    throw EnvoyException("Failed to read private key from disk.");
  }

  method_ = std::make_shared<SSL_PRIVATE_KEY_METHOD>();

  // Have two modes, "rsa" and "ecdsa", for testing multi-cert use cases.
  if (mode_ == "rsa") {
    if (EVP_PKEY_id(pkey.get()) != EVP_PKEY_RSA) {
      throw EnvoyException("Private key is not RSA.");
    }
    method_->sign = rsaPrivateKeySign;
    method_->decrypt = rsaPrivateKeyDecrypt;
    method_->complete = rsaPrivateKeyComplete;
  } else if (mode_ == "ecdsa") {
    if (EVP_PKEY_id(pkey.get()) != EVP_PKEY_EC) {
      throw EnvoyException("Private key is not ECDSA.");
    }
    method_->sign = ecdsaPrivateKeySign;
    method_->decrypt = ecdsaPrivateKeyDecrypt;
    method_->complete = ecdsaPrivateKeyComplete;
  } else {
    throw EnvoyException("Unknown test provider mode, supported modes are \"rsa\" and \"ecdsa\".");
  }

  pkey_ = std::move(pkey);
}

} // namespace PrivateKeyMethodProvider
} // namespace Extensions
} // namespace Envoy
