// Copyright 2020 Google LLC. All Rights Reserved.
#include "api/aead_whitebox.h"

#include <string.h>
#include <memory>
#include <vector>

#include "base/check_op.h"
#include "base/logging.h"
#include "crypto_utils/crypto_util.h"
#include "impl/reference/memory_util.h"
#include "third_party/boringssl/src/include/openssl/aead.h"
#include "third_party/boringssl/src/include/openssl/rand.h"

struct WB_Aead_Whitebox {
  const EVP_AEAD* algorithm;
  size_t nonce_size;

  EVP_AEAD_CTX context;
};

namespace {
// EVP_AEAD_CTX_init(), EVP_AEAD_CTX_open(), and EVP_AEAD_CTX_seal() return 1
// for success and 0 for failure.
const int kAeadSuccess = 1;

void InitAeadMetadata(WB_Aead_Whitebox* whitebox) {
  const EVP_AEAD* algorithm = EVP_aead_aes_128_gcm_siv();

  whitebox->algorithm = algorithm;
  whitebox->nonce_size = EVP_AEAD_nonce_length(algorithm);
}

std::vector<uint8_t> DeriveKey(const WB_Aead_Whitebox* whitebox,
                               const uint8_t* init_data,
                               size_t init_data_size,
                               const uint8_t* context,
                               size_t context_size) {
  // To derive a key, we need a label to introduce some entropy. The label only
  // needs to be some fixed, arbitrary non-trivial string.
  constexpr char kLabel[] = "Covfefe";

  const std::string init_data_view(init_data, init_data + init_data_size);
  const std::string context_view(context, context + context_size);

  // While we get the key size in bytes, DeriveKey() needs the key size in bits.
  const auto derived_key = widevine::crypto_util::DeriveKey(
      init_data_view, kLabel, context_view,
      EVP_AEAD_key_length(whitebox->algorithm) * 8);

  return std::vector<uint8_t>(derived_key.begin(), derived_key.end());
}
}  // namespace

WB_Result WB_Aead_Create(const uint8_t* whitebox_init_data,
                         size_t whitebox_init_data_size,
                         const uint8_t* context,
                         size_t context_size,
                         ::WB_Aead_Whitebox** whitebox) {
  if (!whitebox_init_data || !context || !whitebox) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (whitebox_init_data_size == 0 || context_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // Use a unique pointer internally so that we can return early and
  // automatically release the pointer. Should always be non-null on modern
  // compilers (https://isocpp.org/wiki/faq/freestore-mgmt).
  std::unique_ptr<WB_Aead_Whitebox> aead_whitebox(new WB_Aead_Whitebox());
  InitAeadMetadata(aead_whitebox.get());

  // Rather than using a fixed key, shared across all instances, derive a key
  // for each instance using the provided context. This will allow each instance
  // to have its own unique key (assuming they provide a unique context) while
  // still allowing a specific key to be recreated when needed (e.g.
  // device-locked key).
  const std::vector<uint8_t> key =
      DeriveKey(aead_whitebox.get(), whitebox_init_data,
                whitebox_init_data_size, context, context_size);

  const int result = EVP_AEAD_CTX_init(
      &aead_whitebox->context, aead_whitebox->algorithm, key.data(), key.size(),
      EVP_AEAD_DEFAULT_TAG_LENGTH, nullptr);

  if (result != kAeadSuccess) {
    DVLOG(1) << "Invalid parameter: invalid init data.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // release() will release ownership of the pointer and return the unmanaged
  // raw pointer.
  *whitebox = aead_whitebox.release();
  return WB_RESULT_OK;
}

void WB_Aead_Delete(WB_Aead_Whitebox* whitebox) {
  if (whitebox != nullptr) {
    EVP_AEAD_CTX_cleanup(&whitebox->context);
  }

  // Safe to delete nullptr (https://isocpp.org/wiki/faq/freestore-mgmt).
  delete whitebox;
}
WB_Result WB_Aead_Encrypt(const WB_Aead_Whitebox* whitebox,
                          const uint8_t* input_data,
                          size_t input_data_size,
                          uint8_t* output_data,
                          size_t* output_data_size) {
  if (!whitebox || !input_data || !output_data || !output_data_size) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  if (input_data_size == 0) {
    DVLOG(1) << "Invalid parameter: array size 0.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  std::vector<uint8_t> nonce(whitebox->nonce_size);
  RAND_bytes(nonce.data(), nonce.size());

  std::vector<uint8_t> output(input_data_size +
                              EVP_AEAD_max_overhead(whitebox->algorithm));
  size_t sealed_size;
  CHECK_EQ(EVP_AEAD_CTX_seal(&whitebox->context, output.data(), &sealed_size,
                             output.size(), nonce.data(), nonce.size(),
                             input_data, input_data_size, nullptr, 0),
           kAeadSuccess);
  output.resize(sealed_size);

  // At this point, |output| will have the encrypted data and authentication
  // tag, but in order to decrypt later, we will need the nonce, so append the
  // nonce to the output.
  output.insert(output.end(), nonce.begin(), nonce.end());

  if (!widevine::MemCopy(output.data(), output.size(), output_data,
                         *output_data_size)) {
    DVLOG(1) << "Buffer too small: output needs " << output.size() << ".";
    *output_data_size = output.size();
    return WB_RESULT_BUFFER_TOO_SMALL;
  }

  *output_data_size = output.size();
  return WB_RESULT_OK;
}

WB_Result WB_Aead_Decrypt(const WB_Aead_Whitebox* whitebox,
                          const uint8_t* input_data,
                          size_t input_data_size,
                          uint8_t* output_data,
                          size_t* output_data_size) {
  if (!whitebox || !input_data || !output_data || !output_data_size) {
    DVLOG(1) << "Invalid parameter: null pointer.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // If we were to encrypt nothing, we would still get the nonce, so the input
  // must at least be large enough to contain the nonce.
  if (input_data_size <= whitebox->nonce_size) {
    DVLOG(1) << "Invalid parameter: invalid input size.";
    return WB_RESULT_INVALID_PARAMETER;
  }

  // The input follows the structure:
  // |--------input---------|
  // |---payload---||-nonce-|
  //
  // The payload is the combination of authentication tag and ciphertext
  // created by our previous call to EVP_AEAD_CTX_seal().
  //
  // To use EVP_AEAD_CTX_open(), we need to separate the payload and nonce as
  // they must be passed to EVP_AEAD_CTX_open() separately. The result of
  // calling EVP_AEAD_CTX_open() will be the plaintext, which will be smaller
  // than the payload as the payload included the ciphertext and authentication
  // tag.

  const uint8_t* payload_start = input_data;
  const size_t payload_size = input_data_size - whitebox->nonce_size;

  const uint8_t* nonce_start = payload_start + payload_size;
  const size_t nonce_size = whitebox->nonce_size;

  // Remember, the plaintext will be smaller than the payload, but we are going
  // to start with the payload size as it is our best estimate and will scale it
  // down when we have the final size.
  std::vector<uint8_t> plaintext(payload_size);
  size_t plaintext_size;

  const int result = EVP_AEAD_CTX_open(
      &whitebox->context, plaintext.data(), &plaintext_size, plaintext.size(),
      nonce_start, nonce_size, payload_start, payload_size, nullptr, 0);

  if (result != kAeadSuccess) {
    DVLOG(1) << "Data verification error: failed to verify input data.";
    return WB_RESULT_DATA_VERIFICATION_ERROR;
  }

  // Know that EVP_AEAD_CTX_open() has opened the payload, we know the actual
  // plaintext size and can now shrink the vector to the correct size.
  plaintext.resize(plaintext_size);

  if (!widevine::MemCopy(plaintext.data(), plaintext.size(), output_data,
                         *output_data_size)) {
    DVLOG(1) << "Buffer too small: output needs " << plaintext.size() << ".";
    *output_data_size = plaintext.size();
    return WB_RESULT_BUFFER_TOO_SMALL;
  }

  *output_data_size = plaintext.size();
  return WB_RESULT_OK;
}