// Copyright 2021 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine License
// Agreement.
//
//  Reference implementation utilities of OEMCrypto APIs
//
#include "cmac.h"

#include <openssl/evp.h>

#include "log.h"
#include "scoped_object.h"

namespace wvoec {
namespace util {
namespace {
using ScopedCmacCtx = ScopedObject<CMAC_CTX, CMAC_CTX_free>;
constexpr size_t kAes128KeySize = 16;
constexpr size_t kAes256KeySize = 32;
constexpr size_t kCmacOutputSize = 16;

// Gets the appropriate AES block cipher for the CMAC algortihm
// based on the key size.
// Ownership of the pointer returned by this function is retained by
// the OpenSSL/BoringSSL framework.
const EVP_CIPHER* KeySizeToCipher(size_t key_size) {
  switch (key_size) {
    case kAes128KeySize:
      return EVP_aes_128_cbc();
    case kAes256KeySize:
      return EVP_aes_256_cbc();
  }
  LOGE("Unexpected key size: size = %zu", key_size);
  return nullptr;
}
}  // namespace

// static
std::unique_ptr<Cmac> Cmac::Create(const uint8_t* key, size_t key_size) {
  std::unique_ptr<Cmac> cmac;
  if (key == nullptr) {
    LOGE("CMAC key is null");
    return cmac;
  }
  if (key_size != kAes128KeySize && key_size != kAes256KeySize) {
    LOGE("Invalid CMAC key size: size = %zu", key_size);
    return cmac;
  }
  cmac.reset(new Cmac());
  if (!cmac->Init(key, key_size)) {
    cmac.reset();
  }
  return cmac;
}

// static
std::unique_ptr<Cmac> Cmac::Create(const std::vector<uint8_t>& key) {
  if (key.empty()) {
    LOGE("CMAC key is empty");
    return std::unique_ptr<Cmac>();
  }
  return Create(key.data(), key.size());
}

bool Cmac::Init(const uint8_t* key, size_t key_size) {
  const EVP_CIPHER* const cipher = KeySizeToCipher(key_size);
  if (cipher == nullptr) {
    LOGE("Failed to get block cipher for CMAC");
    return false;
  }
  ScopedCmacCtx ctx(CMAC_CTX_new());
  if (!ctx) {
    LOGE("Failed allocate CMAC CTX");
    return false;
  }
  if (!CMAC_Init(ctx.get(), key, key_size, cipher, nullptr)) {
    LOGE("Failed to initialize CMAC CTX");
    return false;
  }
  ctx_ = ctx.release();
  ready_ = true;
  return true;
}

bool Cmac::Update(const uint8_t* data, size_t data_length) {
  if (data == nullptr) {
    LOGE("Data is null");
    return false;
  }
  if (data_length == 0) {
    return true;
  }
  if (!ready_) {
    LOGE("CMAC must be reset before updating");
    return false;
  }
  if (!CMAC_Update(ctx_, data, data_length)) {
    LOGE("Failed to update CMAC CTX");
    ready_ = false;
    return false;
  }
  return true;
}

bool Cmac::Update(const std::vector<uint8_t>& data) {
  return Update(data.data(), data.size());
}

bool Cmac::Update(uint8_t datum) { return Update(&datum, 1); }

bool Cmac::Finalize(std::vector<uint8_t>* mac) {
  if (mac == nullptr) {
    LOGE("Output MAC buffer is null");
    return false;
  }
  mac->clear();
  return FinalizeAppend(mac);
}

bool Cmac::FinalizeAppend(std::vector<uint8_t>* mac) {
  if (mac == nullptr) {
    LOGE("Output MAC buffer is null");
    return false;
  }
  if (!ready_) {
    LOGE("CMAC must be reset before finalizing");
    return false;
  }
  const size_t end = mac->size();
  size_t mac_size = kCmacOutputSize;
  mac->resize(end + mac_size);
  if (!CMAC_Final(ctx_, &mac->at(end), &mac_size)) {
    LOGE("Failed to finalize CMAC CTX");
    mac->resize(end);
    ready_ = false;
    return false;
  }
  ready_ = false;
  return true;
}

#ifdef OPENSSL_IS_BORINGSSL
// BoringSSL allows for resetting a CMAC context explicitly, whereas
// OpenSSL does so by reinitializing using all nulls/zeros.  This
// causes segfaults on systems using BoringSSL.
void Cmac::Reset() {
  if (!CMAC_Reset(ctx_)) {
    LOGE("Failed to reset CMAC CTX");
    ready_ = false;
  } else {
    ready_ = true;
  }
}
#else  // OpenSSL is OpenSSL
void Cmac::Reset() {
  if (!CMAC_Init(ctx_, nullptr, 0, nullptr, nullptr)) {
    LOGE("Failed to reset CMAC CTX");
    ready_ = false;
  } else {
    ready_ = true;
  }
}
#endif

Cmac::~Cmac() {
  if (ctx_ != nullptr) {
    CMAC_CTX_free(ctx_);
    ctx_ = nullptr;
  }
  ready_ = false;
}
}  // namespace util
}  // namespace wvoec