// Copyright 2020 Google LLC. All Rights Reserved.

#include <stddef.h>
#include <stdint.h>

#include <memory>
#include <vector>

#include "api/license_whitebox.h"
#include "api/license_whitebox_benchmark.h"
#include "api/result.h"
#include "api/test_data.h"
#include "api/test_license_builder.h"
#include "benchmarking/data_source.h"
#include "benchmarking/measurements.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace widevine {
namespace {
// The mask size can be any size (depends on the implementation), so 256 should
// be more than enough for any implementation.
constexpr size_t kMaskSize = 256;
}  // namespace

// Test Parameter:
//   - The number of blocks given to each decrypt call. Each block will be 16
//     bytes.
//   - The number of iterations to run.
class LicenseWhiteboxDecryptBenchmark
    : public LicenseWhiteboxBenchmark,
      public testing::WithParamInterface<std::tuple<size_t, size_t>> {
 protected:
  void SetUp() override {
    LicenseWhiteboxBenchmark::SetUp();

    // Extract all parameters.
    size_t blocks_per_call;
    std::tie(blocks_per_call, iterations_) = GetParam();

    // We are using AES with no padding, the input and output will be the same
    // size.
    const size_t bytes_per_call = blocks_per_call * 16;
    ciphertext_ = Data().Get(bytes_per_call);
    masked_text_.resize(bytes_per_call);
    plaintext_.resize(bytes_per_call);

    const auto init_data = GetLicenseInitData();
    ASSERT_EQ(WB_License_Create(init_data.data(), init_data.size(), &whitebox_),
              WB_RESULT_OK);

    const auto license = CreateLicense();
    ASSERT_EQ(
        WB_License_ProcessLicenseResponse(
            whitebox_, license.core_message.data(), license.core_message.size(),
            license.message.data(), license.message.size(),
            license.signature.data(), license.signature.size(),
            license.session_key.data(), license.session_key.size(),
            license.request.data(), license.request.size()),
        WB_RESULT_OK);
  }

  void TearDown() override { WB_License_Delete(whitebox_); }

  WB_License_Whitebox* whitebox_;

  std::vector<uint8_t> ciphertext_;
  std::vector<uint8_t> masked_text_;
  std::vector<uint8_t> plaintext_;

  size_t iterations_;
};

TEST_P(LicenseWhiteboxDecryptBenchmark, DecryptCBCThroughput) {
  Timer timer;
  timer.Reset();

  for (size_t i = 0; i < iterations_; i++) {
    size_t plaintext_size = plaintext_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CBC,
                                 ContentKeyId().data(), ContentKeyId().size(),
                                 ciphertext_.data(), ciphertext_.size(),
                                 ContentIV().data(), ContentIV().size(),
                                 plaintext_.data(), &plaintext_size));
  }

  Throughput throughput(timer.Get(), iterations_ * ciphertext_.size());
  PrettyPrint("License Decrypt CBC Throughput", throughput, ciphertext_.size());
}

TEST_P(LicenseWhiteboxDecryptBenchmark, DecryptCTRThroughput) {
  Timer timer;
  timer.Reset();

  for (size_t i = 0; i < iterations_; i++) {
    size_t plaintext_size = plaintext_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_License_Decrypt(whitebox_, WB_CIPHER_MODE_CTR,
                                 ContentKeyId().data(), ContentKeyId().size(),
                                 ciphertext_.data(), ciphertext_.size(),
                                 ContentIV().data(), ContentIV().size(),
                                 plaintext_.data(), &plaintext_size));
  }

  Throughput throughput(timer.Get(), iterations_ * ciphertext_.size());
  PrettyPrint("License Decrypt CTR Throughput", throughput, ciphertext_.size());
}

TEST_P(LicenseWhiteboxDecryptBenchmark, MaskedDecryptCBCThroughput) {
  std::vector<uint8_t> mask(kMaskSize);
  size_t mask_size = mask.size();
  ASSERT_EQ(WB_RESULT_OK,
            WB_License_GetSecretString(
                whitebox_, WB_CIPHER_MODE_CBC, ContentKeyId().data(),
                ContentKeyId().size(), mask.data(), &mask_size));
  mask.resize(mask_size);

  Timer timer;
  timer.Reset();

  for (size_t i = 0; i < iterations_; i++) {
    // Include the unmask in the timing as it will be part of the flow from
    // encrypted to consumption.
    size_t plaintext_size = plaintext_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_License_MaskedDecrypt(
                  whitebox_, WB_CIPHER_MODE_CBC, ContentKeyId().data(),
                  ContentKeyId().size(), ciphertext_.data(), ciphertext_.size(),
                  ContentIV().data(), ContentIV().size(), masked_text_.data(),
                  &plaintext_size));

    WB_License_Unmask(masked_text_.data(), 0, plaintext_size, mask.data(),
                      mask.size(), plaintext_.data());
  }

  Throughput throughput(timer.Get(), iterations_ * ciphertext_.size());
  PrettyPrint("License Masked Decrypt CBC Throughput", throughput,
              ciphertext_.size());
}

TEST_P(LicenseWhiteboxDecryptBenchmark, MaskedDecryptCTRThroughput) {
  std::vector<uint8_t> mask(kMaskSize);
  size_t mask_size = mask.size();
  ASSERT_EQ(WB_RESULT_OK,
            WB_License_GetSecretString(
                whitebox_, WB_CIPHER_MODE_CTR, ContentKeyId().data(),
                ContentKeyId().size(), mask.data(), &mask_size));
  mask.resize(mask_size);

  Timer timer;
  timer.Reset();

  for (size_t i = 0; i < iterations_; i++) {
    // Include the unmask in the timing as it will be part of the flow from
    // encrypted to consumption.
    size_t plaintext_size = plaintext_.size();
    ASSERT_EQ(WB_RESULT_OK,
              WB_License_MaskedDecrypt(
                  whitebox_, WB_CIPHER_MODE_CTR, ContentKeyId().data(),
                  ContentKeyId().size(), ciphertext_.data(), ciphertext_.size(),
                  ContentIV().data(), ContentIV().size(), masked_text_.data(),
                  &plaintext_size));

    WB_License_Unmask(masked_text_.data(), 0, plaintext_size, mask.data(),
                      mask.size(), plaintext_.data());
  }

  Throughput throughput(timer.Get(), iterations_ * ciphertext_.size());
  PrettyPrint("License Masked Decrypt CTR Throughput", throughput,
              ciphertext_.size());
}

// The first value is the number of blocks. Each block is 16 bytes.
INSTANTIATE_TEST_SUITE_P(
    DifferentBytesPerCall,
    LicenseWhiteboxDecryptBenchmark,
    ::testing::Values(std::make_tuple(1, 100),           //  16  B
                      std::make_tuple(16, 100),          // 256  B
                      std::make_tuple(64, 100),          //   1 KB
                      std::make_tuple(1024, 100),        //  16 KB
                      std::make_tuple(16 * 1024, 25),    // 256 KB
                      std::make_tuple(64 * 1024, 10)));  //   1 MB
}  // namespace widevine