// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
// Test data for OEMCrypto unit tests.
//
#ifndef CDM_OEMCRYPTO_DECRYPT_TEST_
#define CDM_OEMCRYPTO_DECRYPT_TEST_

#include <gtest/gtest.h>

#include "OEMCryptoCENC.h"
#include "log.h"
#include "oec_decrypt_fallback_chain.h"
#include "oemcrypto_basic_test.h"
#include "oemcrypto_license_test.h"

namespace wvoec {

// Used to test the different HDCP versions.  This test is parameterized by the
// required HDCP version in the key control block.
class OEMCryptoSessionTestDecryptWithHDCP : public OEMCryptoSessionTests,
                                            public WithParamInterface<int> {
 protected:
  void DecryptWithHDCP(OEMCrypto_HDCP_Capability version) {
    OEMCryptoResult sts;
    OEMCrypto_HDCP_Capability current, maximum;
    sts = OEMCrypto_GetHDCPCapability(&current, &maximum);
    ASSERT_EQ(OEMCrypto_SUCCESS, sts);
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(&s));
    LicenseRoundTrip license_messages(&s);
    license_messages.set_control((version << wvoec::kControlHDCPVersionShift) |
                                 wvoec::kControlObserveHDCP |
                                 wvoec::kControlHDCPRequired);
    ASSERT_NO_FATAL_FAILURE(license_messages.SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(license_messages.CreateDefaultResponse());
    ASSERT_NO_FATAL_FAILURE(license_messages.EncryptAndSignResponse());
    ASSERT_EQ(OEMCrypto_SUCCESS, license_messages.LoadResponse());

    if (((version <= HDCP_V2_3 || current >= HDCP_V1_0) && version > current) ||
        (current == HDCP_V1 && version >= HDCP_V1_0)) {
      if (global_features.api_version >= 16) {
        // Can provide either OEMCrypto_WARNING_MIXED_OUTPUT_PROTECTION or
        // OEMCrypto_ERROR_INSUFFICIENT_HDCP. TestDecryptCTR allows either to be
        // reported if OEMCrypto_WARNING_MIXED_OUTPUT_PROTECTION is expected.
        ASSERT_NO_FATAL_FAILURE(
            s.TestDecryptCTR(true, OEMCrypto_WARNING_MIXED_OUTPUT_PROTECTION))
            << "Failed when current HDCP = " << HDCPCapabilityAsString(current)
            << ", maximum HDCP = " << HDCPCapabilityAsString(maximum)
            << ", license HDCP = " << HDCPCapabilityAsString(version);
      } else {
        ASSERT_NO_FATAL_FAILURE(
            s.TestDecryptCTR(true, OEMCrypto_ERROR_INSUFFICIENT_HDCP))
            << "Failed when current HDCP = " << HDCPCapabilityAsString(current)
            << ", maximum HDCP = " << HDCPCapabilityAsString(maximum)
            << ", license HDCP = " << HDCPCapabilityAsString(version);
      }
    } else {
      ASSERT_NO_FATAL_FAILURE(s.TestDecryptCTR(true, OEMCrypto_SUCCESS))
          << "Failed when current HDCP = " << HDCPCapabilityAsString(current)
          << ", maximum HDCP = " << HDCPCapabilityAsString(maximum)
          << ", license HDCP = " << HDCPCapabilityAsString(version);
    }
  }
};

struct SubsampleSize {
  size_t clear_size;
  size_t encrypted_size;
  SubsampleSize(size_t clear, size_t encrypted)
      : clear_size(clear), encrypted_size(encrypted) {}
};

// Struct for holding the data for one test sample in the decrypt tests.
struct TestSample {
  // Encrypted data -- this is input to OEMCrypto, and output from EncryptData.
  std::vector<uint8_t> encrypted_buffer;
  std::vector<uint8_t> clear_buffer;  // OEMCrypto store clear output here.
  std::vector<uint8_t> truth_buffer;  // Truth data for clear text.
  OEMCrypto_SampleDescription description;
  std::vector<OEMCrypto_SubSampleDescription> subsamples;
  int secure_buffer_fid;
};

// A class of tests that test decryption for a variety of patterns and modes.
// This test is parameterized by three parameters:
// 1. The pattern used for pattern decryption.
// 2. The cipher mode for decryption: either CTR or CBC.
// 3. A boolean that determines if decrypt in place should be done.  When the
//    output buffer is clear, it should be possible for the input and output
//    buffers to be the same.
class OEMCryptoSessionTestsDecryptTests
    : public OEMCryptoLicenseTestAPI16,
      public WithParamInterface<tuple<OEMCrypto_CENCEncryptPatternDesc,
                                      OEMCryptoCipherMode, OutputType>> {
 protected:
  void SetUp() override {
    OEMCryptoLicenseTestAPI16::SetUp();
    if (wvoec::global_features.derive_key_method ==
        wvoec::DeviceFeatures::NO_METHOD) {
      GTEST_SKIP() << "Test for devices that can derive session keys only.";
    }
    pattern_ = ::testing::get<0>(GetParam());
    cipher_mode_ = ::testing::get<1>(GetParam());
    decrypt_inplace_ = ::testing::get<2>(GetParam()).decrypt_inplace;
    output_buffer_type_ = ::testing::get<2>(GetParam()).type;
    verify_crc_ = global_features.supports_crc;
    // Pick a random key.
    EXPECT_EQ(GetRandBytes(key_, sizeof(key_)), 1);
    // Pick a random starting iv.  Some tests override this before using it.
    EXPECT_EQ(GetRandBytes(initial_iv_, sizeof(initial_iv_)), 1);
  }

  void TearDown() override {
    FreeSecureBuffers();
    OEMCryptoLicenseTestAPI16::TearDown();
  }

  void SetSubsampleSizes(std::vector<SubsampleSize> subsample_sizes) {
    // This is just sugar for having one sample with the given subsamples in it.
    SetSampleSizes({std::move(subsample_sizes)});
  }

  void SetSampleSizes(std::vector<std::vector<SubsampleSize>> sample_sizes) {
    ASSERT_GT(sample_sizes.size(), 0u);
    samples_.clear();
    samples_.reserve(sample_sizes.size());

    // Convert all the size arrays to TestSample structs
    for (const std::vector<SubsampleSize>& subsample_sizes : sample_sizes) {
      // This could be one line if we had C++17
      samples_.emplace_back();
      TestSample& sample = samples_.back();

      ASSERT_GT(subsample_sizes.size(), 0u);
      sample.subsamples.reserve(subsample_sizes.size());

      // Convert all the sizes to subsample descriptions and tally the total
      // sample size
      size_t sample_size = 0;
      size_t current_block_offset = 0;
      for (const SubsampleSize& size : subsample_sizes) {
        sample.subsamples.push_back(OEMCrypto_SubSampleDescription{
            size.clear_size, size.encrypted_size,
            0,  // Subsample Flags, to be filled in after the loop
            current_block_offset});

        // Update the rolling variables
        sample_size += size.clear_size + size.encrypted_size;
        if (cipher_mode_ == OEMCrypto_CipherMode_CENC) {
          current_block_offset =
              (current_block_offset + size.encrypted_size) % AES_BLOCK_SIZE;
        }
      }

      // Set the subsample flags now that all the subsamples are processed
      sample.subsamples.front().subsample_flags |= OEMCrypto_FirstSubsample;
      sample.subsamples.back().subsample_flags |= OEMCrypto_LastSubsample;

      // Set related information on the sample description
      sample.description.subsamples = sample.subsamples.data();
      sample.description.subsamples_length = sample.subsamples.size();
      sample.description.buffers.input_data_length = sample_size;
    }
  }

  // Set up the input buffer and either a clear or secure output buffer for each
  // test sample. This should be called after SetSubsampleSizes().
  void MakeBuffers() {
    for (TestSample& sample : samples_) {
      const size_t total_size = sample.description.buffers.input_data_length;
      ASSERT_GT(total_size, 0u);
      sample.encrypted_buffer.clear();
      sample.truth_buffer.clear();
      sample.clear_buffer.clear();
      sample.encrypted_buffer.resize(total_size);
      sample.truth_buffer.resize(total_size);
      for (size_t i = 0; i < total_size; i++) sample.truth_buffer[i] = i % 256;

      OEMCrypto_DestBufferDesc& output_descriptor =
          sample.description.buffers.output_descriptor;
      output_descriptor.type = output_buffer_type_;
      switch (output_descriptor.type) {
        case OEMCrypto_BufferType_Clear:
          if (decrypt_inplace_) {
            // Add some padding to verify there is no overrun.
            sample.encrypted_buffer.resize(total_size + kBufferOverrunPadding,
                                           0xaa);
            output_descriptor.buffer.clear.clear_buffer =
                sample.encrypted_buffer.data();
          } else {
            // Add some padding to verify there is no overrun.
            sample.clear_buffer.resize(total_size + kBufferOverrunPadding,
                                       0xaa);
            output_descriptor.buffer.clear.clear_buffer =
                sample.clear_buffer.data();
          }
          output_descriptor.buffer.clear.clear_buffer_length = total_size;
          break;

        case OEMCrypto_BufferType_Secure:
          output_descriptor.buffer.secure.secure_buffer_length = total_size;
          ASSERT_EQ(OEMCrypto_AllocateSecureBuffer(
                        session_.session_id(), total_size, &output_descriptor,
                        &sample.secure_buffer_fid),
                    OEMCrypto_SUCCESS);
          ASSERT_NE(output_descriptor.buffer.secure.secure_buffer, nullptr);
          // It is OK if OEMCrypto changes the maximum size, but there must
          // still be enough room for our data.
          ASSERT_GE(output_descriptor.buffer.secure.secure_buffer_length,
                    total_size);
          output_descriptor.buffer.secure.offset = 0;
          break;

        case OEMCrypto_BufferType_Direct:
          output_descriptor.buffer.direct.is_video = false;
          break;
      }  // switch (output_descriptor.type)
    }    // sample loop
  }

  void FreeSecureBuffers() {
    for (TestSample& sample : samples_) {
      OEMCrypto_DestBufferDesc& output_descriptor =
          sample.description.buffers.output_descriptor;
      if (output_descriptor.type == OEMCrypto_BufferType_Secure) {
        ASSERT_EQ(OEMCrypto_FreeSecureBuffer(session_.session_id(),
                                             &output_descriptor,
                                             sample.secure_buffer_fid),
                  OEMCrypto_SUCCESS);
      }
    }
  }

  void EncryptData() {
    AES_KEY aes_key;
    AES_set_encrypt_key(key_, AES_BLOCK_SIZE * 8, &aes_key);

    for (TestSample& sample : samples_) {
      uint8_t iv[KEY_IV_SIZE];  // Current IV
      memcpy(iv, initial_iv_, KEY_IV_SIZE);
      memcpy(sample.description.iv, initial_iv_, KEY_IV_SIZE);

      size_t buffer_index = 0;  // byte index into in and out.
      size_t block_offset = 0;  // byte index into current block.
      for (const OEMCrypto_SubSampleDescription& subsample :
           sample.subsamples) {
        // Copy clear content.
        if (subsample.num_bytes_clear > 0) {
          memcpy(&sample.encrypted_buffer[buffer_index],
                 &sample.truth_buffer[buffer_index], subsample.num_bytes_clear);
          buffer_index += subsample.num_bytes_clear;
        }

        // The IV resets at the start of each subsample in the 'cbcs' schema.
        if (cipher_mode_ == OEMCrypto_CipherMode_CBCS) {
          memcpy(iv, initial_iv_, KEY_IV_SIZE);
        }

        size_t pattern_offset = 0;
        const size_t subsample_end =
            buffer_index + subsample.num_bytes_encrypted;
        while (buffer_index < subsample_end) {
          const size_t size =
              min(subsample_end - buffer_index, AES_BLOCK_SIZE - block_offset);
          const size_t pattern_length = pattern_.encrypt + pattern_.skip;
          const bool skip_block =
              (pattern_offset >= pattern_.encrypt) && (pattern_length > 0);
          if (pattern_length > 0) {
            pattern_offset = (pattern_offset + 1) % pattern_length;
          }
          // CBC mode should just copy a partial block at the end.  If there
          // is a partial block at the beginning, an error is returned, so we
          // can put whatever we want in the output buffer.
          if (skip_block || ((cipher_mode_ == OEMCrypto_CipherMode_CBCS) &&
                             (size < AES_BLOCK_SIZE))) {
            memcpy(&sample.encrypted_buffer[buffer_index],
                   &sample.truth_buffer[buffer_index], size);
            block_offset = 0;  // Next block should be complete.
          } else {
            if (cipher_mode_ == OEMCrypto_CipherMode_CENC) {
              uint8_t aes_output[AES_BLOCK_SIZE];
              AES_encrypt(iv, aes_output, &aes_key);
              for (size_t n = 0; n < size; n++) {
                sample.encrypted_buffer[buffer_index + n] =
                    aes_output[n + block_offset] ^
                    sample.truth_buffer[buffer_index + n];
              }
              if (size + block_offset < AES_BLOCK_SIZE) {
                // Partial block.  Don't increment iv.  Compute next block
                // offset.
                block_offset = block_offset + size;
              } else {
                EXPECT_EQ(block_offset + size,
                          static_cast<size_t>(AES_BLOCK_SIZE));
                // Full block.  Increment iv, and set offset to 0 for next
                // block.
                ctr128_inc64(1, iv);
                block_offset = 0;
              }
            } else {
              uint8_t aes_input[AES_BLOCK_SIZE];
              for (size_t n = 0; n < size; n++) {
                aes_input[n] = sample.truth_buffer[buffer_index + n] ^ iv[n];
              }
              AES_encrypt(aes_input, &sample.encrypted_buffer[buffer_index],
                          &aes_key);
              memcpy(iv, &sample.encrypted_buffer[buffer_index],
                     AES_BLOCK_SIZE);
              // CBC mode should always start on block boundary.
              block_offset = 0;
            }
          }
          buffer_index += size;
        }  // encryption loop
      }    // per-subsample loop
    }      // per-sample loop
  }

  void LoadLicense() {
    uint32_t control = wvoec::kControlNonceEnabled;
    if (verify_crc_) control |= kControlAllowHashVerification;
    if (output_buffer_type_ == OEMCrypto_BufferType_Secure)
      control |= kControlObserveDataPath | kControlDataPathSecure;
    license_messages_.set_control(control);
    ASSERT_NO_FATAL_FAILURE(license_messages_.SignAndVerifyRequest());

    ASSERT_NO_FATAL_FAILURE(license_messages_.CreateDefaultResponse());
    license_messages_.core_response()
        .timer_limits.initial_renewal_duration_seconds = kDuration;
    memcpy(license_messages_.response_data().keys[0].key_data, key_,
           sizeof(key_));
    license_messages_.response_data().keys[0].cipher_mode = cipher_mode_;
    ASSERT_NO_FATAL_FAILURE(license_messages_.EncryptAndSignResponse());
    ASSERT_EQ(OEMCrypto_SUCCESS, license_messages_.LoadResponse());
    ASSERT_EQ(GetKeyHandleIntoVector(session_.session_id(),
                                     session_.license().keys[0].key_id,
                                     session_.license().keys[0].key_id_length,
                                     cipher_mode_, key_handle_),
              OEMCrypto_SUCCESS);
  }

  void TestDecryptCENC() { ASSERT_EQ(DecryptCENC(), OEMCrypto_SUCCESS); }

  void ValidateDecryptedData() {
    for (TestSample& sample : samples_) {
      if (sample.description.buffers.output_descriptor.type ==
          OEMCrypto_BufferType_Clear) {
        const size_t total_size = sample.description.buffers.input_data_length;
        // To verify there is no buffer overrun after decrypting, look at the
        // padded bytes just after the data buffer that was written.  It
        // should not have changed from the original 0xaa that we set in
        // MakeBuffer function.
        if (decrypt_inplace_) {
          EXPECT_EQ(std::count(sample.encrypted_buffer.begin() + total_size,
                               sample.encrypted_buffer.end(), 0xaa),
                    static_cast<int32_t>(kBufferOverrunPadding))
              << "Buffer overrun.";
          sample.encrypted_buffer.resize(total_size);  // Remove padding.
          // We expect encrypted buffer to have been changed by OEMCrypto.
          EXPECT_EQ(sample.encrypted_buffer, sample.truth_buffer);
        } else {
          EXPECT_EQ(std::count(sample.clear_buffer.begin() + total_size,
                               sample.clear_buffer.end(), 0xaa),
                    static_cast<int32_t>(kBufferOverrunPadding))
              << "Buffer overrun.";
          sample.clear_buffer.resize(total_size);  // Remove padding.
          EXPECT_EQ(sample.clear_buffer, sample.truth_buffer);
        }
      }
    }
    if (verify_crc_) {
      uint32_t frame;
      ASSERT_EQ(OEMCrypto_GetHashErrorCode(session_.session_id(), &frame),
                OEMCrypto_SUCCESS);
    }
  }

  OEMCryptoResult DecryptCENC() {
    // OEMCrypto only supports providing a decrypt hash for one sample.
    if (samples_.size() > 1) verify_crc_ = false;

    // If supported, check the decrypt hashes.
    if (verify_crc_) {
      const TestSample& sample = samples_[0];

      uint32_t hash =
          util::wvcrc32(sample.truth_buffer.data(), sample.truth_buffer.size());
      OEMCrypto_SetDecryptHash(session_.session_id(), 1,
                               reinterpret_cast<const uint8_t*>(&hash),
                               sizeof(hash));
    }

    // Build an array of just the sample descriptions.
    std::vector<OEMCrypto_SampleDescription> sample_descriptions;
    sample_descriptions.reserve(samples_.size());
    for (TestSample& sample : samples_) {
      // This must be deferred until this point in case the test modifies the
      // buffer before testing decrypt.
      sample.description.buffers.input_data = sample.encrypted_buffer.data();
      // Append to the description array.
      sample_descriptions.push_back(sample.description);
    }

    // Perform decryption using the test data that was previously set up.
    OEMCryptoResult result = DecryptFallbackChain::Decrypt(
        key_handle_.data(), key_handle_.size(), sample_descriptions.data(),
        sample_descriptions.size(), cipher_mode_, &pattern_);
    if (result != OEMCrypto_SUCCESS) return result;
    ValidateDecryptedData();
    return result;
  }

  // Parameters of test case
  OEMCrypto_CENCEncryptPatternDesc pattern_;
  OEMCryptoCipherMode cipher_mode_;
  bool decrypt_inplace_;  // If true, input and output buffers are the same.
  OEMCryptoBufferType output_buffer_type_;

  bool verify_crc_;
  uint8_t key_[AES_BLOCK_SIZE];      // Encryption Key.
  uint8_t initial_iv_[KEY_IV_SIZE];  // Starting IV for every sample.
  std::vector<TestSample> samples_;
  std::vector<uint8_t> key_handle_;
};

}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_DECRYPT_TEST_