// 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_LICENSE_TEST_
#define CDM_OEMCRYPTO_LICENSE_TEST_

#include <gtest/gtest.h>

#include "OEMCryptoCENC.h"
#include "log.h"
#include "oec_session_util.h"
#include "oemcrypto_basic_test.h"
#include "oemcrypto_resource_test.h"
#include "oemcrypto_session_tests_helper.h"

using ::testing::WithParamInterface;

namespace wvoec {

// Used for testing oemcrypto APIs with huge buffers.
typedef const std::function<OEMCryptoResult(size_t)> oemcrypto_function;
void TestHugeLengthDoesNotCrashAPI(oemcrypto_function f,
                                   size_t start_buffer_length,
                                   size_t end_buffer_length, bool check_status);
void TestHugeLengthDoesNotCrashAPI(oemcrypto_function f, bool check_status);

void TestMaxKeys(SessionUtil* util, size_t num_keys_per_session);

class OEMCryptoSessionTests : public OEMCryptoClientTest {
 public:
  vector<uint8_t> encrypted_usage_header_;

 protected:
  OEMCryptoSessionTests() {}

  void SetUp() override {
    OEMCryptoClientTest::SetUp();
    EnsureTestROT();
    if (global_features.usage_table) {
      CreateUsageTableHeader();
    }
  }

  void CreateUsageTableHeader(bool expect_success = true) {
    size_t header_buffer_length = 0;
    OEMCryptoResult sts =
        OEMCrypto_CreateUsageTableHeader(nullptr, &header_buffer_length);
    if (expect_success) {
      ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
    } else {
      ASSERT_NE(OEMCrypto_SUCCESS, sts);
      if (sts != OEMCrypto_ERROR_SHORT_BUFFER) return;
    }
    encrypted_usage_header_.resize(header_buffer_length);
    sts = OEMCrypto_CreateUsageTableHeader(encrypted_usage_header_.data(),
                                           &header_buffer_length);
    if (expect_success) {
      ASSERT_EQ(OEMCrypto_SUCCESS, sts);
      encrypted_usage_header_.resize(header_buffer_length);
    } else {
      ASSERT_NE(OEMCrypto_SUCCESS, sts);
    }
  }

  void TestPrepareLicenseRequestForHugeBufferLengths(
      const std::function<void(size_t, LicenseRoundTrip*)> f,
      bool check_status) {
    auto oemcrypto_function = [&](size_t message_length) {
      Session s;
      s.open();
      InstallTestDrmKey(&s);
      LicenseRoundTrip license_messages(&s);
      f(message_length, &license_messages);
      OEMCryptoResult result =
          license_messages.SignAndCreateRequestWithCustomBufferLengths();
      s.close();
      return result;
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, check_status);
  }

  OEMCryptoResult LoadLicense(Session& s, LicenseRoundTrip& license_messages) {
    InstallTestDrmKey(&s);
    license_messages.SignAndVerifyRequest();
    license_messages.CreateDefaultResponse();
    license_messages.EncryptAndSignResponse();
    return license_messages.LoadResponse();
  }
};

class OEMCryptoSessionTestKeyboxTest : public OEMCryptoSessionTests {};

// This class is for testing a single license with the default API version
// of 16.
class OEMCryptoLicenseTestAPI16 : public OEMCryptoSessionTests {
 public:
  OEMCryptoLicenseTestAPI16()
      : license_api_version_(kCurrentAPI), license_messages_(&session_) {}

  void SetUp() override {
    OEMCryptoSessionTests::SetUp();
    ASSERT_NO_FATAL_FAILURE(session_.open());
    ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(&session_));
  }

  void TearDown() override {
    ASSERT_NO_FATAL_FAILURE(session_.close());
    OEMCryptoSessionTests::TearDown();
  }

 protected:
  Session session_;
  uint32_t license_api_version_;
  LicenseRoundTrip license_messages_;
};

// This class is used to test a license that is from a server with the specified
// version parameter. Up to two versions old.
class OEMCryptoLicenseTest : public OEMCryptoLicenseTestAPI16,
                             public WithParamInterface<uint32_t> {
 protected:
  void SetUp() override {
    // The only difference between this class and its parent is that we use a
    // different license api:
    license_api_version_ = GetParam();
    license_messages_.set_api_version(license_api_version_);
    OEMCryptoLicenseTestAPI16::SetUp();
  }

  void LoadLicense() {
    license_messages_.SignAndVerifyRequest();
    license_messages_.CreateDefaultResponse();
    license_messages_.EncryptAndSignResponse();
    license_messages_.LoadResponse();
  }

  void TestDecryptCENCForHugeBufferLengths(
      const std::function<void(size_t, OEMCrypto_SampleDescription*)> f,
      bool check_status) {
    LoadLicense();
    auto oemcrypto_function = [&](size_t message_length) {
      vector<uint8_t> key_handle;
      GetKeyHandleIntoVector(session_.session_id(),
                             session_.license().keys[0].key_id,
                             session_.license().keys[0].key_id_length,
                             OEMCrypto_CipherMode_CENC, key_handle);

      size_t input_buffer_size = 1;
      vector<uint8_t> in_buffer(input_buffer_size + message_length);
      vector<uint8_t> out_buffer(in_buffer.size());

      OEMCrypto_SampleDescription sample_description;
      OEMCrypto_SubSampleDescription subsample_description;
      GenerateSimpleSampleDescription(
          in_buffer, out_buffer, &sample_description, &subsample_description);

      OEMCrypto_SubSampleDescription* sub_samples =
          const_cast<OEMCrypto_SubSampleDescription*>(
              sample_description.subsamples);
      // Actual tests modifies either of these fields to match clear + encrypted
      // = in_buffer.size().
      sub_samples[0].num_bytes_clear = 0;
      sub_samples[0].num_bytes_encrypted = input_buffer_size;

      // Create the pattern description (always 0,0 for CTR)
      OEMCrypto_CENCEncryptPatternDesc pattern = {0, 0};
      int secure_fd = 0;
      f(message_length, &sample_description);
      if (sample_description.buffers.output_descriptor.type ==
          OEMCrypto_BufferType_Secure) {
        OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
            session_.session_id(), in_buffer.size(),
            &sample_description.buffers.output_descriptor, &secure_fd);
        if (sts != OEMCrypto_SUCCESS) {
          LOGI("Secure buffers are not supported.");
          return sts;
        }
      }
      // Try to decrypt the data
      OEMCryptoResult result =
          OEMCrypto_DecryptCENC(key_handle.data(), key_handle.size(),
                                &sample_description, 1, &pattern);
      if (sample_description.buffers.output_descriptor.type ==
          OEMCrypto_BufferType_Secure) {
        OEMCrypto_FreeSecureBuffer(
            session_.session_id(),
            &sample_description.buffers.output_descriptor, secure_fd);
      }
      return result;
    };
    TestHugeLengthDoesNotCrashAPI(oemcrypto_function, check_status);
  }

  void TestDecryptCENCForOutOfRangeOffsetsAndLengths(
      const std::function<void(OEMCrypto_SampleDescription*)> f,
      bool update_secure_buffer) {
    LoadLicense();
    vector<uint8_t> key_handle;
    GetKeyHandleIntoVector(session_.session_id(),
                           session_.license().keys[0].key_id,
                           session_.license().keys[0].key_id_length,
                           OEMCrypto_CipherMode_CENC, key_handle);

    vector<uint8_t> in_buffer(256);
    vector<uint8_t> out_buffer(in_buffer.size());

    OEMCrypto_SampleDescription sample_description;
    OEMCrypto_SubSampleDescription subsample_description;
    GenerateSimpleSampleDescription(in_buffer, out_buffer, &sample_description,
                                    &subsample_description);

    // Create the pattern description (always 0,0 for CTR)
    OEMCrypto_CENCEncryptPatternDesc pattern = {0, 0};
    int secure_fd = 0;
    if (update_secure_buffer) {
      OEMCryptoResult sts = OEMCrypto_AllocateSecureBuffer(
          session_.session_id(), in_buffer.size(),
          &sample_description.buffers.output_descriptor, &secure_fd);
      if (sts != OEMCrypto_SUCCESS) {
        LOGI("Secure buffers are not supported.");
        return;
      }
    }
    f(&sample_description);
    // Try to decrypt the data
    OEMCryptoResult result = OEMCrypto_DecryptCENC(
        key_handle.data(), key_handle.size(), &sample_description, 1, &pattern);
    if (update_secure_buffer) {
      OEMCrypto_FreeSecureBuffer(session_.session_id(),
                                 &sample_description.buffers.output_descriptor,
                                 secure_fd);
    }
    ASSERT_NE(OEMCrypto_SUCCESS, result);
  }
};

}  // namespace wvoec

#endif  // CDM_OEMCRYPTO_LICENSE_TEST_