// Copyright 2023 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//

#include "oemcrypto_provisioning_test.h"

#include <stdint.h>

#include <string>

#include "bcc_validator.h"
#include "device_info_validator.h"
#include "log.h"
#include "oec_device_features.h"
#include "platform.h"
#include "signed_csr_payload_validator.h"
#include "test_sleep.h"

namespace wvoec {

/// @addtogroup provision
/// @{

// This test is used to print the device ID to stdout.
TEST_F(OEMCryptoKeyboxTest, NormalGetDeviceId) {
  OEMCryptoResult sts;
  uint8_t dev_id[128] = {0};
  size_t dev_id_len = 128;
  sts = OEMCrypto_GetDeviceID(dev_id, &dev_id_len);
  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
  cout << "             NormalGetDeviceId: dev_id = "
       << MaybeHex(dev_id, dev_id_len) << " len = " << dev_id_len << endl;
  RecordWvProperty("device_id", wvutil::HexEncode(dev_id, dev_id_len));
}

TEST_F(OEMCryptoKeyboxTest, GetDeviceIdShortBuffer) {
  OEMCryptoResult sts;
  uint8_t dev_id[128];
  for (int i = 0; i < 128; ++i) {
    dev_id[i] = 0x55;
  }
  dev_id[127] = '\0';
  size_t dev_id_len = 0;
  sts = OEMCrypto_GetDeviceID(dev_id, &dev_id_len);
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
  // On short buffer error, function should return minimum buffer length
  ASSERT_GT(dev_id_len, 0u);
  // Should also return short buffer if passed a zero length and a null buffer.
  dev_id_len = 0;
  sts = OEMCrypto_GetDeviceID(nullptr, &dev_id_len);
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, sts);
  // On short buffer error, function should return minimum buffer length
  ASSERT_GT(dev_id_len, 0u);
}

TEST_F(OEMCryptoKeyboxTest, NormalGetKeyData) {
  OEMCryptoResult sts;
  uint8_t key_data[256];
  size_t key_data_len = sizeof(key_data);
  sts = OEMCrypto_GetKeyData(key_data, &key_data_len);

  uint32_t* data = reinterpret_cast<uint32_t*>(key_data);
  const uint32_t system_id = htonl(data[1]);
  const uint32_t version = htonl(data[0]);
  printf("             NormalGetKeyData: system_id = %u = 0x%04X, version=%u\n",
         system_id, system_id, version);
  RecordWvProperty("system_id", std::to_string(system_id));
  RecordWvProperty("key_data_version", std::to_string(version));
  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
}

TEST_F(OEMCryptoKeyboxTest, GetKeyDataNullPointer) {
  OEMCryptoResult sts;
  uint8_t key_data[256];
  sts = OEMCrypto_GetKeyData(key_data, nullptr);
  ASSERT_NE(OEMCrypto_SUCCESS, sts);
}

// This test makes sure the installed keybox is valid.  It doesn't really check
// that it is a production keybox.  That must be done by an integration test.
TEST_F(OEMCryptoKeyboxTest, ProductionKeyboxValid) {
  ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_IsKeyboxValid());
}

// This verifies that the device really does claim to have a certificate.
// It should be filtered out for devices that have a keybox.
TEST_F(OEMCryptoProv30Test, DeviceClaimsOEMCertificate) {
  ASSERT_EQ(OEMCrypto_OEMCertificate, OEMCrypto_GetProvisioningMethod());
}

TEST_F(OEMCryptoProv30Test, GetDeviceId) {
  OEMCryptoResult sts;
  std::vector<uint8_t> dev_id(128, 0);
  size_t dev_id_len = dev_id.size();
  sts = OEMCrypto_GetDeviceID(dev_id.data(), &dev_id_len);
  if (sts == OEMCrypto_ERROR_SHORT_BUFFER) {
    ASSERT_GT(dev_id_len, 0u);
    dev_id.resize(dev_id_len);
    sts = OEMCrypto_GetDeviceID(dev_id.data(), &dev_id_len);
  }
  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
  dev_id.resize(dev_id_len);
  cout << "             NormalGetDeviceId: dev_id = " << MaybeHex(dev_id)
       << " len = " << dev_id_len << endl;
  RecordWvProperty("device_id",
                   wvutil::HexEncode(dev_id.data(), dev_id.size()));
}

// The OEM certificate must be valid.
TEST_F(OEMCryptoProv30Test, CertValidAPI15) {
  ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_IsKeyboxOrOEMCertValid());
}

TEST_F(OEMCryptoProv30Test, OEMCertValid) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  bool kVerify = true;
  ASSERT_NO_FATAL_FAILURE(s.LoadOEMCert(kVerify));  // Load and verify.
}

/** This verifies that the OEM Certificate cannot be used with
 * GenerateRSASignature.
 */
TEST_F(OEMCryptoProv30Test, OEMCertForbidGenerateRSASignature1) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(s.LoadOEMCert());
  DisallowForbiddenPadding(s.session_id(), kSign_PKCS1_Block1, 80);
}

/** This verifies that the OEM Certificate cannot be used with
 * GenerateRSASignature.
 */
TEST_F(OEMCryptoProv30Test, OEMCertForbidGenerateRSASignature2) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(s.LoadOEMCert());
  DisallowForbiddenPadding(s.session_id(), kSign_RSASSA_PSS, 80);
}

// Calling OEMCrypto_GetOEMPublicCertificate should not change the session's
// private key.
TEST_F(OEMCryptoProv30Test, GetCertOnlyAPI16) {
  if (wrapped_drm_key_.size() == 0) {
    // If we don't have a wrapped key yet, create one.
    // This wrapped key will be shared by all sessions in the test.
    ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  }
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  // Install the DRM Cert's RSA key.
  ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));
  ASSERT_NO_FATAL_FAILURE(s.SetTestRsaPublicKey());
  // Request the OEM Cert. -- This should NOT load the OEM Private key.
  vector<uint8_t> public_cert;
  size_t public_cert_length = 0;
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
            OEMCrypto_GetOEMPublicCertificate(nullptr, &public_cert_length));
  ASSERT_LT(0u, public_cert_length);
  public_cert.resize(public_cert_length);
  ASSERT_EQ(OEMCrypto_SUCCESS, OEMCrypto_GetOEMPublicCertificate(
                                   public_cert.data(), &public_cert_length));
  // Derive keys from the session key -- this should use the DRM Cert's key.
  // It should NOT use the OEM Private key because that key should not have
  // been loaded.
  // Now fill a message and try to load it.
  LicenseRoundTrip license_messages(&s);
  license_messages.set_control(0);
  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());
}

/** This verifies that the OEM Certificate cannot be used with
 * GenerateRSASignature.
 */
TEST_F(OEMCryptoProv40Test, OEMCertForbidGenerateRSASignature1) {
  // Create an OEM Cert and save it for later.
  Session s1;
  ASSERT_NO_FATAL_FAILURE(s1.open());
  ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&s1));
  ASSERT_EQ(s1.IsPublicKeySet(), true);
  s1.close();
  Session s2;
  ASSERT_NO_FATAL_FAILURE(s2.open());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            OEMCrypto_InstallOemPrivateKey(
                s2.session_id(), oem_key_type_,
                reinterpret_cast<const uint8_t*>(wrapped_oem_key_.data()),
                wrapped_oem_key_.size()));
  DisallowForbiddenPadding(s2.session_id(), kSign_PKCS1_Block1, 80);
}

/** This verifies that the OEM Certificate cannot be used with
 * GenerateRSASignature.
 */
TEST_F(OEMCryptoProv40Test, OEMCertForbidGenerateRSASignature2) {
  // Create an OEM Cert and save it for later.
  Session s1;
  ASSERT_NO_FATAL_FAILURE(s1.open());
  ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&s1));
  ASSERT_EQ(s1.IsPublicKeySet(), true);
  s1.close();
  Session s2;
  ASSERT_NO_FATAL_FAILURE(s2.open());
  ASSERT_EQ(OEMCrypto_SUCCESS,
            OEMCrypto_InstallOemPrivateKey(
                s2.session_id(), oem_key_type_,
                reinterpret_cast<const uint8_t*>(wrapped_oem_key_.data()),
                wrapped_oem_key_.size()));
  DisallowForbiddenPadding(s2.session_id(), kSign_RSASSA_PSS, 80);
}

// This verifies that the device really does claim to have BCC.
// It should be filtered out for devices that have a keybox or factory OEM
// cert.
TEST_F(OEMCryptoProv40Test, DeviceClaimsBootCertificateChain) {
  ASSERT_EQ(OEMCrypto_GetProvisioningMethod(), OEMCrypto_BootCertificateChain);
}

// Verifies that short buffer error returns when the buffer is short.
TEST_F(OEMCryptoProv40Test, GetBootCertificateChainShortBuffer) {
  std::vector<uint8_t> bcc;
  size_t bcc_size = 0;
  std::vector<uint8_t> additional_signature;
  size_t additional_signature_size = 0;
  ASSERT_EQ(OEMCrypto_GetBootCertificateChain(bcc.data(), &bcc_size,
                                              additional_signature.data(),
                                              &additional_signature_size),
            OEMCrypto_ERROR_SHORT_BUFFER);
  ASSERT_NE(bcc_size, 0uL);
}

// Verifies BCC can be successfully returned.
TEST_F(OEMCryptoProv40Test, GetBootCertificateChainSuccess) {
  std::vector<uint8_t> bcc;
  size_t bcc_size = 0;
  std::vector<uint8_t> additional_signature;
  size_t additional_signature_size = 0;
  ASSERT_EQ(OEMCrypto_GetBootCertificateChain(bcc.data(), &bcc_size,
                                              additional_signature.data(),
                                              &additional_signature_size),
            OEMCrypto_ERROR_SHORT_BUFFER);

  bcc.resize(bcc_size);
  additional_signature.resize(additional_signature_size);
  ASSERT_EQ(OEMCrypto_GetBootCertificateChain(bcc.data(), &bcc_size,
                                              additional_signature.data(),
                                              &additional_signature_size),
            OEMCrypto_SUCCESS);
  util::BccValidator validator;
  EXPECT_EQ(util::CborMessageStatus::kCborParseOk, validator.Parse(bcc));
  util::CborMessageStatus status = validator.Validate();
  EXPECT_LT(status, util::CborMessageStatus::kCborValidateError);
  if (status >= util::CborMessageStatus::kCborValidateError) {
    const std::string parsed_bcc = validator.GetFormattedMessage();
    LOGE("%s", parsed_bcc.c_str());
    LOGE("Validation results\n%s", validator.PrintValidateMessage().c_str());
  }
}

// Verifies BCC signature and its type if they are available.
TEST_F(OEMCryptoProv40Test, AdditionalBccSignature) {
  std::vector<uint8_t> bcc;
  size_t bcc_size = 0;
  std::vector<uint8_t> additional_signature;
  size_t additional_signature_size = 0;
  ASSERT_EQ(OEMCrypto_GetBootCertificateChain(bcc.data(), &bcc_size,
                                              additional_signature.data(),
                                              &additional_signature_size),
            OEMCrypto_ERROR_SHORT_BUFFER);

  bcc.resize(bcc_size);
  additional_signature.resize(additional_signature_size);
  ASSERT_EQ(OEMCrypto_GetBootCertificateChain(bcc.data(), &bcc_size,
                                              additional_signature.data(),
                                              &additional_signature_size),
            OEMCrypto_SUCCESS);
  OEMCrypto_BCCSignatureType bcc_signature_type;
  const OEMCryptoResult result =
      OEMCrypto_GetBCCSignatureType(&bcc_signature_type);
  if (result == OEMCrypto_ERROR_NOT_IMPLEMENTED) return;
  ASSERT_EQ(result, OEMCrypto_SUCCESS);
  if (!additional_signature.empty()) {
    ASSERT_NE(bcc_signature_type, OEMCrypto_BCCSigType_Unknown);
  } else {
    ASSERT_EQ(bcc_signature_type, OEMCrypto_BCCSigType_Unknown);
  }
}

// Verifies that short buffer error returns when the buffer is short.
TEST_F(OEMCryptoProv40Test, GenerateCertificateKeyPairShortBuffer) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  std::vector<uint8_t> public_key;
  size_t public_key_size = 0;
  std::vector<uint8_t> public_key_signature;
  size_t public_key_signature_size = 0;
  std::vector<uint8_t> wrapped_private_key;
  size_t wrapped_private_key_size = 0;
  OEMCrypto_PrivateKeyType key_type;

  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key.data(), &public_key_size,
          public_key_signature.data(), &public_key_signature_size,
          wrapped_private_key.data(), &wrapped_private_key_size, &key_type),
      OEMCrypto_ERROR_SHORT_BUFFER);

  ASSERT_NE(public_key_size, 0uL);
  ASSERT_NE(public_key_signature_size, 0uL);
  ASSERT_NE(wrapped_private_key_size, 0uL);
}

// Verifies a pair of key can be successfully returned.
TEST_F(OEMCryptoProv40Test, GenerateCertificateKeyPairSuccess) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  std::vector<uint8_t> public_key;
  size_t public_key_size = 0;
  std::vector<uint8_t> public_key_signature;
  size_t public_key_signature_size = 0;
  std::vector<uint8_t> wrapped_private_key;
  size_t wrapped_private_key_size = 0;
  OEMCrypto_PrivateKeyType key_type;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key.data(), &public_key_size,
          public_key_signature.data(), &public_key_signature_size,
          wrapped_private_key.data(), &wrapped_private_key_size, &key_type),
      OEMCrypto_ERROR_SHORT_BUFFER);
  public_key.resize(public_key_size);
  public_key_signature.resize(public_key_signature_size);
  wrapped_private_key.resize(wrapped_private_key_size);

  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key.data(), &public_key_size,
          public_key_signature.data(), &public_key_signature_size,
          wrapped_private_key.data(), &wrapped_private_key_size, &key_type),
      OEMCrypto_SUCCESS);
  public_key.resize(public_key_size);
  public_key_signature.resize(public_key_signature_size);
  wrapped_private_key.resize(wrapped_private_key_size);
  // Parse the public key generated to make sure it is correctly formatted.
  ASSERT_NO_FATAL_FAILURE(s.SetPublicKeyFromSubjectPublicKey(
      key_type, public_key.data(), public_key_size));
}

// Verifies the generated key pairs are different on each call.
TEST_F(OEMCryptoProv40Test, GenerateCertificateKeyPairsAreDifferent) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  // Large buffer to make sure it is large enough.
  size_t public_key_size1 = 10000;
  std::vector<uint8_t> public_key1(public_key_size1);
  size_t public_key_signature_size1 = 10000;
  std::vector<uint8_t> public_key_signature1(public_key_signature_size1);
  size_t wrapped_private_key_size1 = 10000;
  std::vector<uint8_t> wrapped_private_key1(wrapped_private_key_size1);
  OEMCrypto_PrivateKeyType key_type1;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key1.data(), &public_key_size1,
          public_key_signature1.data(), &public_key_signature_size1,
          wrapped_private_key1.data(), &wrapped_private_key_size1, &key_type1),
      OEMCrypto_SUCCESS);
  EXPECT_NE(public_key_size1, 0UL);
  EXPECT_NE(public_key_signature_size1, 0UL);
  EXPECT_NE(wrapped_private_key_size1, 0UL);
  public_key1.resize(public_key_size1);
  public_key_signature1.resize(public_key_signature_size1);
  wrapped_private_key1.resize(wrapped_private_key_size1);

  size_t public_key_size2 = 10000;
  std::vector<uint8_t> public_key2(public_key_size2);
  size_t public_key_signature_size2 = 10000;
  std::vector<uint8_t> public_key_signature2(public_key_signature_size2);
  size_t wrapped_private_key_size2 = 10000;
  std::vector<uint8_t> wrapped_private_key2(wrapped_private_key_size2);
  OEMCrypto_PrivateKeyType key_type2;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key2.data(), &public_key_size2,
          public_key_signature2.data(), &public_key_signature_size2,
          wrapped_private_key2.data(), &wrapped_private_key_size2, &key_type2),
      OEMCrypto_SUCCESS);
  EXPECT_NE(public_key_size2, 0UL);
  EXPECT_NE(public_key_signature_size2, 0UL);
  EXPECT_NE(wrapped_private_key_size2, 0UL);
  public_key2.resize(public_key_size2);
  public_key_signature2.resize(public_key_signature_size2);
  wrapped_private_key2.resize(wrapped_private_key_size2);

  EXPECT_NE(public_key1, public_key2);
  EXPECT_NE(public_key_signature1, public_key_signature2);
  EXPECT_NE(wrapped_private_key1, wrapped_private_key2);
}

TEST_F(OEMCryptoProv40Test, GetDeviceInformationAPI18) {
  if (wvoec::global_features.api_version < 18) {
    GTEST_SKIP() << "Test for versions 18 and up only.";
  }
  std::vector<uint8_t> device_info;
  size_t device_info_length = 0;
  OEMCryptoResult sts =
      OEMCrypto_GetDeviceInformation(device_info.data(), &device_info_length);
  ASSERT_EQ(sts, OEMCrypto_ERROR_SHORT_BUFFER);
  ASSERT_NE(device_info_length, 0uL);
  device_info.resize(device_info_length);
  ASSERT_EQ(
      OEMCrypto_GetDeviceInformation(device_info.data(), &device_info_length),
      OEMCrypto_SUCCESS);
  EXPECT_NE(device_info_length, 0uL);
  device_info.resize(device_info_length);
  constexpr int kDeviceVersion = 3;
  util::DeviceInfoValidator validator(kDeviceVersion);
  EXPECT_EQ(util::CborMessageStatus::kCborParseOk,
            validator.Parse(device_info));
  validator.Validate();
  EXPECT_EQ(util::CborMessageStatus::kCborValidateOk, validator.Validate());
}

TEST_F(OEMCryptoProv40Test, GetDeviceSignedCsrPayloadAPI18) {
  if (wvoec::global_features.api_version < 18) {
    GTEST_SKIP() << "Test for versions 18 and up only.";
  }
  const std::vector<uint8_t> challenge(64, 0xaa);
  const std::vector<uint8_t> device_info = cppbor::Map()
                                               .add("manufacturer", "google")
                                               .add("fused", 0)
                                               .add("other", "ignored")
                                               .canonicalize()
                                               .encode();
  std::vector<uint8_t> signed_csr_payload;
  size_t signed_csr_payload_length = 0;
  OEMCryptoResult sts = OEMCrypto_GetDeviceSignedCsrPayload(
      challenge.data(), challenge.size(), device_info.data(),
      device_info.size(), signed_csr_payload.data(),
      &signed_csr_payload_length);
  ASSERT_EQ(sts, OEMCrypto_ERROR_SHORT_BUFFER);
  ASSERT_NE(signed_csr_payload_length, 0uL);
  signed_csr_payload.resize(signed_csr_payload_length);
  ASSERT_EQ(OEMCrypto_GetDeviceSignedCsrPayload(
                challenge.data(), challenge.size(), device_info.data(),
                device_info.size(), signed_csr_payload.data(),
                &signed_csr_payload_length),
            OEMCrypto_SUCCESS);
  EXPECT_NE(signed_csr_payload_length, 0uL);
  util::SignedCsrPayloadValidator validator;
  EXPECT_EQ(util::CborMessageStatus::kCborParseOk,
            validator.Parse(signed_csr_payload));
  EXPECT_EQ(util::CborMessageStatus::kCborValidateOk, validator.Validate());
}

TEST_F(OEMCryptoProv40Test, GetDeviceSignedCsrPayloadInvalid) {
  std::vector<uint8_t> signed_csr_payload;
  size_t signed_csr_payload_length = 0;
  const std::vector<uint8_t> challenge(64, 0xaa);
  const std::vector<uint8_t> device_info = cppbor::Map()
                                               .add("manufacturer", "google")
                                               .add("fused", 0)
                                               .add("other", "ignored")
                                               .canonicalize()
                                               .encode();
  std::vector<uint8_t> challenge_empty;
  OEMCryptoResult sts = OEMCrypto_GetDeviceSignedCsrPayload(
      challenge_empty.data(), challenge_empty.size(), device_info.data(),
      device_info.size(), signed_csr_payload.data(),
      &signed_csr_payload_length);
  if (sts == OEMCrypto_ERROR_NOT_IMPLEMENTED) return;
  ASSERT_EQ(sts, OEMCrypto_ERROR_INVALID_CONTEXT);

  // Oversized challenge
  const std::vector<uint8_t> challenge_long(65, 0xaa);
  sts = OEMCrypto_GetDeviceSignedCsrPayload(
      challenge_long.data(), challenge_long.size(), device_info.data(),
      device_info.size(), signed_csr_payload.data(),
      &signed_csr_payload_length);
  ASSERT_EQ(sts, OEMCrypto_ERROR_INVALID_CONTEXT);

  std::vector<uint8_t> device_empty;
  sts = OEMCrypto_GetDeviceSignedCsrPayload(
      challenge.data(), challenge.size(), device_empty.data(),
      device_empty.size(), signed_csr_payload.data(),
      &signed_csr_payload_length);
  ASSERT_EQ(sts, OEMCrypto_ERROR_INVALID_CONTEXT);
}

// Verifies that an OEM private key can be installed.
TEST_F(OEMCryptoProv40Test, InstallOemPrivateKeySuccess) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  // First generate a key pair.
  // Large buffer to make sure it is large enough.
  size_t public_key_size = 10000;
  std::vector<uint8_t> public_key(public_key_size);
  size_t public_key_signature_size = 10000;
  std::vector<uint8_t> public_key_signature(public_key_signature_size);
  size_t wrapped_private_key_size = 10000;
  std::vector<uint8_t> wrapped_private_key(wrapped_private_key_size);
  OEMCrypto_PrivateKeyType key_type;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key.data(), &public_key_size,
          public_key_signature.data(), &public_key_signature_size,
          wrapped_private_key.data(), &wrapped_private_key_size, &key_type),
      OEMCrypto_SUCCESS);
  public_key.resize(public_key_size);
  public_key_signature.resize(public_key_signature_size);
  wrapped_private_key.resize(wrapped_private_key_size);

  // Install the generated private key.
  ASSERT_EQ(OEMCrypto_InstallOemPrivateKey(s.session_id(), key_type,
                                           wrapped_private_key.data(),
                                           wrapped_private_key_size),
            OEMCrypto_SUCCESS);
}

// If data is empty or random, the API should return non-success status.
TEST_F(OEMCryptoProv40Test, InstallOemPrivateKeyInvalidDataFail) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());

  // Empty key fails.
  std::vector<uint8_t> wrapped_private_key;
  OEMCrypto_PrivateKeyType key_type = OEMCrypto_RSA_Private_Key;
  ASSERT_NE(OEMCrypto_InstallOemPrivateKey(s.session_id(), key_type,
                                           wrapped_private_key.data(),
                                           wrapped_private_key.size()),
            OEMCrypto_SUCCESS);

  // Random key data fails.
  wrapped_private_key = {1, 2, 3};
  ASSERT_NE(OEMCrypto_InstallOemPrivateKey(s.session_id(), key_type,
                                           wrapped_private_key.data(),
                                           wrapped_private_key.size()),
            OEMCrypto_SUCCESS);
}

// Verifies that an OEM private key can be installed, and used by
// GenerateCertificateKeyPair call.
TEST_F(OEMCryptoProv40Test, InstallOemPrivateKeyCanBeUsed) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  // First generate a key pair.
  size_t public_key_size1 = 10000;
  std::vector<uint8_t> public_key1(public_key_size1);
  size_t public_key_signature_size1 = 10000;
  std::vector<uint8_t> public_key_signature1(public_key_signature_size1);
  size_t wrapped_private_key_size1 = 10000;
  std::vector<uint8_t> wrapped_private_key1(wrapped_private_key_size1);
  OEMCrypto_PrivateKeyType key_type1;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key1.data(), &public_key_size1,
          public_key_signature1.data(), &public_key_signature_size1,
          wrapped_private_key1.data(), &wrapped_private_key_size1, &key_type1),
      OEMCrypto_SUCCESS);
  EXPECT_NE(public_key_size1, 0UL);
  EXPECT_NE(public_key_signature_size1, 0UL);
  EXPECT_NE(wrapped_private_key_size1, 0UL);
  public_key1.resize(public_key_size1);
  public_key_signature1.resize(public_key_signature_size1);
  wrapped_private_key1.resize(wrapped_private_key_size1);

  // Install the generated private key.
  ASSERT_EQ(OEMCrypto_InstallOemPrivateKey(s.session_id(), key_type1,
                                           wrapped_private_key1.data(),
                                           wrapped_private_key_size1),
            OEMCrypto_SUCCESS);

  // Now calling GenerateCertificateKeyPair should use wrapped_private_key to
  // sign the newly generated public key.
  size_t public_key_size2 = 10000;
  std::vector<uint8_t> public_key2(public_key_size2);
  size_t public_key_signature_size2 = 10000;
  std::vector<uint8_t> public_key_signature2(public_key_signature_size2);
  size_t wrapped_private_key_size2 = 10000;
  std::vector<uint8_t> wrapped_private_key2(wrapped_private_key_size2);
  OEMCrypto_PrivateKeyType key_type2;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key2.data(), &public_key_size2,
          public_key_signature2.data(), &public_key_signature_size2,
          wrapped_private_key2.data(), &wrapped_private_key_size2, &key_type2),
      OEMCrypto_SUCCESS);
  EXPECT_NE(public_key_size2, 0UL);
  EXPECT_NE(public_key_signature_size2, 0UL);
  EXPECT_NE(wrapped_private_key_size2, 0UL);
  public_key2.resize(public_key_size2);
  public_key_signature2.resize(public_key_signature_size2);
  wrapped_private_key2.resize(wrapped_private_key_size2);

  // Verify public_key_signature2 with public_key1.
  if (key_type1 == OEMCrypto_PrivateKeyType::OEMCrypto_RSA_Private_Key) {
    ASSERT_NO_FATAL_FAILURE(s.SetRsaPublicKeyFromSubjectPublicKey(
        public_key1.data(), public_key1.size()));
    ASSERT_NO_FATAL_FAILURE(
        s.VerifyRsaSignature(public_key2, public_key_signature2.data(),
                             public_key_signature2.size(), kSign_RSASSA_PSS));
  } else if (key_type1 == OEMCrypto_PrivateKeyType::OEMCrypto_ECC_Private_Key) {
    ASSERT_NO_FATAL_FAILURE(s.SetEccPublicKeyFromSubjectPublicKey(
        public_key1.data(), public_key1.size()));
    ASSERT_NO_FATAL_FAILURE(s.VerifyEccSignature(public_key2,
                                                 public_key_signature2.data(),
                                                 public_key_signature2.size()));
  }
}

/** Verify that the private key from an OEM Cert cannot be loaded as a DRM
 * cert.
 */
TEST_F(OEMCryptoProv40Test, OEMPrivateKeyCannotBeDRMKey) {
  // Create an OEM Cert and save it for later.
  Session s1;
  ASSERT_NO_FATAL_FAILURE(s1.open());
  ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&s1));
  ASSERT_EQ(s1.IsPublicKeySet(), true);
  s1.close();
  const std::vector<uint8_t> wrapped_oem_key1 = wrapped_oem_key_;
  // Now create a new OEM cert, load the second key, and try to load key1
  // as the DRM key.
  Session s2;
  ASSERT_NO_FATAL_FAILURE(s2.open());
  ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&s2));
  s2.close();
  // Load the current key as the OEM key in session 3.
  Session s3;
  ASSERT_NO_FATAL_FAILURE(s3.open());
  // Now try to load key 1 as a DRM key. That should fail.
  ASSERT_EQ(OEMCrypto_ERROR_INVALID_KEY,
            OEMCrypto_LoadDRMPrivateKey(s3.session_id(), oem_key_type_,
                                        wrapped_oem_key1.data(),
                                        wrapped_oem_key1.size()));
}

/** The private key for a DRM Cert cannot be loaded as an OEM Certificate. */
TEST_F(OEMCryptoProv40Test, DRMPrivateKeyCannotBeOEMKey) {
  // Create a DRM cert and save it for later.
  Session s1;
  // Make sure the drm private key exists.
  ASSERT_NO_FATAL_FAILURE(s1.open());
  ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(&s1));
  ASSERT_NE(wrapped_drm_key_.size(), 0u);
  // Now try to load the drm private key as an OEM key.
  Session s2;
  ASSERT_NO_FATAL_FAILURE(s2.open());
  ASSERT_EQ(OEMCrypto_ERROR_INVALID_KEY,
            OEMCrypto_InstallOemPrivateKey(
                s2.session_id(), drm_key_type_,
                reinterpret_cast<const uint8_t*>(wrapped_drm_key_.data()),
                wrapped_drm_key_.size()));
}

TEST_F(OEMCryptoProv40Test, GetDeviceId) {
  OEMCryptoResult sts;
  std::vector<uint8_t> dev_id;
  size_t dev_id_len = dev_id.size();
  sts = OEMCrypto_GetDeviceID(dev_id.data(), &dev_id_len);
  if (sts == OEMCrypto_ERROR_SHORT_BUFFER) {
    ASSERT_GT(dev_id_len, 0u);
    dev_id.resize(dev_id_len);
    sts = OEMCrypto_GetDeviceID(dev_id.data(), &dev_id_len);
  }
  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
  dev_id.resize(dev_id_len);
  cout << "             NormalGetDeviceId: dev_id = " << MaybeHex(dev_id)
       << " len = " << dev_id_len << endl;
  RecordWvProperty("device_id",
                   wvutil::HexEncode(dev_id.data(), dev_id.size()));
  // Device id should be stable. Query again.
  std::vector<uint8_t> dev_id2(dev_id_len);
  sts = OEMCrypto_GetDeviceID(dev_id2.data(), &dev_id_len);
  ASSERT_EQ(OEMCrypto_SUCCESS, sts);
  ASSERT_EQ(dev_id2, dev_id);
}

// Verifies provisioning stage 1 OEM cert provisioning round trip works
TEST_F(OEMCryptoProv40Test, ProvisionOemCert) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(CreateProv4OEMKey(&s));
  ASSERT_EQ(s.IsPublicKeySet(), true);
}

// Verifies both provisioning stages OEM and DRM cert provisioning round trip
// works
TEST_F(OEMCryptoProv40Test, ProvisionDrmCert) {
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(&s));
  ASSERT_EQ(s.IsPublicKeySet(), true);
}

TEST_P(OEMCryptoProv40CastTest, ProvisionCastWorks) {
  // Generate an OEM key first, to load into next session
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  size_t public_key_size = 10000;
  std::vector<uint8_t> public_key(public_key_size);
  size_t public_key_signature_size = 10000;
  std::vector<uint8_t> public_key_signature(public_key_signature_size);
  size_t wrapped_private_key_size = 10000;
  std::vector<uint8_t> wrapped_private_key(wrapped_private_key_size);
  OEMCrypto_PrivateKeyType key_type;
  ASSERT_EQ(
      OEMCrypto_GenerateCertificateKeyPair(
          s.session_id(), public_key.data(), &public_key_size,
          public_key_signature.data(), &public_key_signature_size,
          wrapped_private_key.data(), &wrapped_private_key_size, &key_type),
      OEMCrypto_SUCCESS);
  public_key.resize(public_key_size);
  public_key_signature.resize(public_key_signature_size);
  wrapped_private_key.resize(wrapped_private_key_size);
  ASSERT_NO_FATAL_FAILURE(s.close());

  // Install OEM key and get cast RSA
  Session s1;
  ASSERT_NO_FATAL_FAILURE(s1.open());
  ASSERT_EQ(OEMCrypto_InstallOemPrivateKey(s1.session_id(), key_type,
                                           wrapped_private_key.data(),
                                           wrapped_private_key_size),
            OEMCrypto_SUCCESS);

  ASSERT_NO_FATAL_FAILURE(CreateProv4CastKey(&s1, GetParam()));
}

INSTANTIATE_TEST_SUITE_P(Prov4CastProvisioningBasic, OEMCryptoProv40CastTest,
                         testing::Values(true, false));

// Verify that you cannot use GenerateRSASignature with a normal DRM Cert.
// that function needs a cast cert.
TEST_F(OEMCryptoLoadsCertificate, ForbidRSASignatureForDRMKey1) {
  DisallowForbiddenPadding(session_.session_id(), kSign_RSASSA_PSS, 80);
}

TEST_F(OEMCryptoLoadsCertificate, ForbidRSASignatureForDRMKey2) {
  DisallowForbiddenPadding(session_.session_id(), kSign_PKCS1_Block1, 80);
}

TEST_F(OEMCryptoLoadsCertificate, PrepAndSignLicenseRequestCounterAPI18) {
  if (wvoec::global_features.api_version < 18) {
    GTEST_SKIP() << "Test for versions 18 and up only.";
  }
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));
  s.GenerateNonce();

  size_t core_message_length = 100;
  std::vector<uint8_t> message(128, 0);
  std::vector<uint8_t> signature(256, 0);
  size_t signature_length = signature.size();

  OEMCryptoResult result = OEMCrypto_PrepAndSignLicenseRequest(
      s.session_id(), message.data(), message.size(), &core_message_length,
      signature.data(), &signature_length);

  ASSERT_EQ(OEMCrypto_SUCCESS, result);
}

// This test verifies that we can create a wrapped RSA key, and then reload it.
TEST_F(OEMCryptoLoadsCertificate, LoadRSASessionKey) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));
}

TEST_F(OEMCryptoLoadsCertificate, SignProvisioningRequest) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.PrepareSession(keybox_));
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
}

// This tests a large message size. The size is larger than we required in v15.
TEST_F(OEMCryptoLoadsCertificate, SignLargeProvisioningRequestAPI16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  const size_t max_size = GetResourceValue(kLargeMessageSize);
  provisioning_messages.set_message_size(max_size);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.PrepareSession(keybox_));
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
}

// This creates a wrapped RSA key, and then does the sanity check that the
// unencrypted key is not found in the wrapped key.  The wrapped key should be
// encrypted.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvision) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.PrepareSession(keybox_));
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_EQ(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  // We should not be able to find the rsa key in the wrapped key. It should
  // be encrypted.
  EXPECT_EQ(nullptr, find(provisioning_messages.wrapped_rsa_key(),
                          provisioning_messages.encoded_rsa_key()));
}

// Verify that RewrapDeviceRSAKey checks pointers are within the provisioning
// message.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadRange1_API16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }

  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  // Encrypt and sign once, so that we can use the size of the response.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  provisioning_messages.core_response().enc_private_key.offset =
      provisioning_messages.encrypted_response_buffer().size() + 1;
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_NE(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Verify that RewrapDeviceRSAKey checks pointers are within the provisioning
// message.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadRange2_API16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  // Encrypt and sign once, so that we can use the size of the response.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  provisioning_messages.core_response().enc_private_key_iv.offset =
      provisioning_messages.encrypted_response_buffer().size() + 1;
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_NE(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Verify that RewrapDeviceRSAKey checks pointers are within the provisioning
// message.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadRange3_API16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  // Encrypt and sign once, so that we can use the size of the response.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  // If the offset is before the end, but the offset+length is bigger, then
  // the message should be rejected.
  provisioning_messages.core_response().enc_private_key.offset =
      provisioning_messages.encrypted_response_buffer().size() - 5;
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_NE(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Verify that RewrapDeviceRSAKey checks pointers are within the provisioning
// message.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadRange4_API16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  // Encrypt and sign once, so that we can use the size of the response.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  // If the offset is before the end, but the offset+length is bigger, then
  // the message should be rejected.
  provisioning_messages.core_response().enc_private_key_iv.offset =
      provisioning_messages.encrypted_response_buffer().size() - 5;
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_NE(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Verify that RewrapDeviceRSAKey checks pointers are within the provisioning
// message.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadRange5Prov30_API16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  if (global_features.provisioning_method != OEMCrypto_OEMCertificate) {
    GTEST_SKIP() << "Test for Prov 3.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  // Encrypt and sign once, so that we can use the size of the response.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  // If the offset is before the end, but the offset+length is bigger, then
  // the message should be rejected.
  provisioning_messages.core_response().encrypted_message_key.offset =
      provisioning_messages.encrypted_response_buffer().size() + 1;
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_NE(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Test that RewrapDeviceRSAKey verifies the message signature.
// TODO(b/144186970): This test should also run on Prov 3.0 devices.
TEST_F(OEMCryptoLoadsCertificate,
       CertificateProvisionBadSignatureKeyboxTestAPI16) {
  if (global_features.provisioning_method != OEMCrypto_Keybox) {
    GTEST_SKIP() << "Test for Prov 2.0 devices only.";
  }
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  if (global_features.provisioning_method != OEMCrypto_Keybox) {
    GTEST_SKIP() << "Test for Prov 2.0 devices only.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  provisioning_messages.response_signature()[4] ^= 42;  // bad signature.
  ASSERT_EQ(OEMCrypto_ERROR_SIGNATURE_FAILURE,
            provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Test that RewrapDeviceRSAKey verifies the nonce is current.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadNonce_API16) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  provisioning_messages.core_request().nonce ^= 42;  // bad nonce.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_EQ(OEMCrypto_ERROR_INVALID_NONCE,
            provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Test that RewrapDeviceRSAKey verifies the RSA key is valid.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionBadRSAKey) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates and do not support
  // key rewrapping.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  provisioning_messages.response_data().rsa_key[4] ^= 42;  // bad key.
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_NE(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Test that RewrapDeviceRSAKey verifies the RSA key is valid.
// TODO(b/144186970): This test should also run on Prov 3.0 devices.
TEST_F(OEMCryptoLoadsCertificate,
       CertificateProvisionBadRSAKeyKeyboxTestAPI16) {
  if (global_features.provisioning_method != OEMCrypto_Keybox) {
    GTEST_SKIP() << "Test for Prov 2.0 devices only.";
  }
  if (global_features.provisioning_method != OEMCrypto_Keybox) {
    GTEST_SKIP() << "Test for Prov 2.0 devices only.";
  }
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  size_t rsa_offset =
      provisioning_messages.core_response().enc_private_key.offset;
  // Offsets are relative to the message body, after the core message.
  rsa_offset += provisioning_messages.serialized_core_message().size();
  rsa_offset += 4;  // Change the middle of the key.
  provisioning_messages.encrypted_response_buffer()[rsa_offset] ^= 42;
  ASSERT_EQ(OEMCrypto_ERROR_SIGNATURE_FAILURE,
            provisioning_messages.LoadResponse());
  provisioning_messages.VerifyLoadFailed();
}

// Test that RewrapDeviceRSAKey accepts the maximum message size.
TEST_F(OEMCryptoLoadsCertificate, CertificateProvisionLargeBuffer) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  Session s;
  ProvisioningRoundTrip provisioning_messages(&s, encoded_rsa_key_);
  const size_t max_size = GetResourceValue(kLargeMessageSize);
  provisioning_messages.set_message_size(max_size);
  provisioning_messages.PrepareSession(keybox_);
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
  ASSERT_EQ(OEMCrypto_SUCCESS, provisioning_messages.LoadResponse());
  // We should not be able to find the rsa key in the wrapped key. It should
  // be encrypted.
  EXPECT_EQ(nullptr, find(provisioning_messages.wrapped_rsa_key(),
                          provisioning_messages.encoded_rsa_key()));
}

// Test that a wrapped RSA key can be loaded.
TEST_F(OEMCryptoLoadsCertificate, LoadWrappedRSAKey) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  Session s;
  ASSERT_NO_FATAL_FAILURE(s.open());
  ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));
}

class OEMCryptoLoadsCertVariousKeys : public OEMCryptoLoadsCertificate {
 public:
  void SetUp() override {
    OEMCryptoLoadsCertificate::SetUp();
    // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
    // provisioning 4. Disabled here temporarily.
    if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
      GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
    }
  }

  void TestKey(const uint8_t* key, size_t key_length) {
    encoded_rsa_key_.assign(key, key + key_length);
    ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
    Session s;
    ASSERT_NO_FATAL_FAILURE(s.open());
    ASSERT_NO_FATAL_FAILURE(s.SetRsaPublicKeyFromPrivateKeyInfo(
        encoded_rsa_key_.data(), encoded_rsa_key_.size()));
    ASSERT_NO_FATAL_FAILURE(s.LoadWrappedRsaDrmKey(wrapped_drm_key_));

    LicenseRoundTrip license_messages(&s);
    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());
    ASSERT_NO_FATAL_FAILURE(s.TestDecryptCTR());
  }
};

// Test a 3072 bit RSA key certificate.
TEST_F(OEMCryptoLoadsCertVariousKeys, TestLargeRSAKey3072) {
  if (!global_features.supports_rsa_3072) {
    GTEST_SKIP() << "OEMCrypto does not support RSA 3072";
  }
  TestKey(kTestRSAPKCS8PrivateKeyInfo3_3072,
          sizeof(kTestRSAPKCS8PrivateKeyInfo3_3072));
}

// Test an RSA key certificate which has a private key generated using the
// Carmichael totient.
TEST_F(OEMCryptoLoadsCertVariousKeys, TestCarmichaelRSAKey) {
  TestKey(kTestKeyRSACarmichael_2048, sizeof(kTestKeyRSACarmichael_2048));
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestCarmichaelNonZeroNormalDer) {
  TestKey(kCarmichaelNonZeroNormalDer, kCarmichaelNonZeroNormalDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestCarmichaelNonZeroShortDer) {
  TestKey(kCarmichaelNonZeroShortDer, kCarmichaelNonZeroShortDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestCarmichaelZeroNormalDer) {
  TestKey(kCarmichaelZeroNormalDer, kCarmichaelZeroNormalDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestCarmichaelZeroShortDer) {
  TestKey(kCarmichaelZeroShortDer, kCarmichaelZeroShortDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestDualNonZeroNormalDer) {
  TestKey(kDualNonZeroNormalDer, kDualNonZeroNormalDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestDualNonZeroShortDer) {
  TestKey(kDualNonZeroShortDer, kDualNonZeroShortDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestDualZeroNormalDer) {
  TestKey(kDualZeroNormalDer, kDualZeroNormalDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestDualZeroShortDer) {
  TestKey(kDualZeroShortDer, kDualZeroShortDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestEulerNonZeroNormalDer) {
  TestKey(kEulerNonZeroNormalDer, kEulerNonZeroNormalDerLen);
}

TEST_F(OEMCryptoLoadsCertVariousKeys, TestEulerZeroNormalDer) {
  TestKey(kEulerZeroNormalDer, kEulerZeroNormalDerLen);
}

// This tests that two sessions can use different RSA keys simultaneously.
TEST_F(OEMCryptoLoadsCertificate, TestMultipleRSAKeys) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  Session s1;  // Session s1 loads the default rsa key, but doesn't use it
               // until after s2 uses its key.
  ASSERT_NO_FATAL_FAILURE(s1.open());
  ASSERT_NO_FATAL_FAILURE(s1.SetRsaPublicKeyFromPrivateKeyInfo(
      encoded_rsa_key_.data(), encoded_rsa_key_.size()));
  ASSERT_NO_FATAL_FAILURE(s1.LoadWrappedRsaDrmKey(wrapped_drm_key_));

  Session s2;  // Session s2 uses a different rsa key.
  encoded_rsa_key_.assign(kTestRSAPKCS8PrivateKeyInfo4_2048,
                          kTestRSAPKCS8PrivateKeyInfo4_2048 +
                              sizeof(kTestRSAPKCS8PrivateKeyInfo4_2048));
  ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
  ASSERT_NO_FATAL_FAILURE(s2.open());
  ASSERT_NO_FATAL_FAILURE(s2.SetRsaPublicKeyFromPrivateKeyInfo(
      encoded_rsa_key_.data(), encoded_rsa_key_.size()));
  ASSERT_NO_FATAL_FAILURE(s2.LoadWrappedRsaDrmKey(wrapped_drm_key_));
  LicenseRoundTrip license_messages2(&s2);
  ASSERT_NO_FATAL_FAILURE(license_messages2.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(license_messages2.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(license_messages2.EncryptAndSignResponse());
  ASSERT_EQ(OEMCrypto_SUCCESS, license_messages2.LoadResponse());
  ASSERT_NO_FATAL_FAILURE(s2.TestDecryptCTR());
  s2.close();

  // After s2 has loaded its rsa key, we continue using s1's key.
  LicenseRoundTrip license_messages1(&s1);
  ASSERT_NO_FATAL_FAILURE(license_messages1.SignAndVerifyRequest());
  ASSERT_NO_FATAL_FAILURE(license_messages1.CreateDefaultResponse());
  ASSERT_NO_FATAL_FAILURE(license_messages1.EncryptAndSignResponse());
  ASSERT_EQ(OEMCrypto_SUCCESS, license_messages1.LoadResponse());
  ASSERT_NO_FATAL_FAILURE(s1.TestDecryptCTR());
}

// This tests the maximum number of DRM private keys that OEMCrypto can load
TEST_F(OEMCryptoLoadsCertificate, TestMaxDRMKeys) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  // DRM Reprovisioning CDMs have embedded certificates.
  if (global_features.provisioning_method == OEMCrypto_DrmReprovisioning) {
    GTEST_SKIP() << "Test for non DRM Reprovisioning devices.";
  }
  const size_t max_total_keys = GetResourceValue(kMaxTotalDRMPrivateKeys);
  std::vector<std::unique_ptr<Session>> sessions;
  std::vector<std::unique_ptr<LicenseRoundTrip>> licenses;

  // It should be able to load up to kMaxTotalDRMPrivateKeys keys
  for (size_t i = 0; i < max_total_keys; i++) {
    sessions.push_back(std::unique_ptr<Session>(new Session()));
    licenses.push_back(std::unique_ptr<LicenseRoundTrip>(
        new LicenseRoundTrip(sessions[i].get())));
    const size_t key_index = i % kTestRSAPKCS8PrivateKeys_2048.size();
    encoded_rsa_key_.assign(kTestRSAPKCS8PrivateKeys_2048[key_index].begin(),
                            kTestRSAPKCS8PrivateKeys_2048[key_index].end());
    ASSERT_NO_FATAL_FAILURE(CreateWrappedDRMKey());
    ASSERT_NO_FATAL_FAILURE(sessions[i]->open());
    ASSERT_NO_FATAL_FAILURE(InstallTestDrmKey(sessions[i].get()));
  }

  // Attempts to load one more key than the kMaxTotalDRMPrivateKeys
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    Session s;
    const size_t buffer_size = 5000;  // Make sure it is large enough.
    std::vector<uint8_t> public_key(buffer_size);
    size_t public_key_size = buffer_size;
    std::vector<uint8_t> public_key_signature(buffer_size);
    size_t public_key_signature_size = buffer_size;
    std::vector<uint8_t> wrapped_private_key(buffer_size);
    size_t wrapped_private_key_size = buffer_size;
    OEMCrypto_PrivateKeyType key_type;
    OEMCryptoResult result = OEMCrypto_GenerateCertificateKeyPair(
        s.session_id(), public_key.data(), &public_key_size,
        public_key_signature.data(), &public_key_signature_size,
        wrapped_private_key.data(), &wrapped_private_key_size, &key_type);
    // Key creation is allowed to fail due to resource restriction
    if (result != OEMCrypto_SUCCESS) {
      ASSERT_TRUE(result == OEMCrypto_ERROR_INSUFFICIENT_RESOURCES ||
                  result == OEMCrypto_ERROR_TOO_MANY_KEYS);
    }
  } else {
    Session s;
    encoded_rsa_key_.assign(kTestRSAPKCS8PrivateKeyInfo2_2048,
                            kTestRSAPKCS8PrivateKeyInfo2_2048 +
                                sizeof(kTestRSAPKCS8PrivateKeyInfo2_2048));
    Session ps;
    ProvisioningRoundTrip provisioning_messages(&ps, encoded_rsa_key_);
    provisioning_messages.PrepareSession(keybox_);
    ASSERT_NO_FATAL_FAILURE(provisioning_messages.SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(provisioning_messages.CreateDefaultResponse());
    ASSERT_NO_FATAL_FAILURE(provisioning_messages.EncryptAndSignResponse());
    OEMCryptoResult result = provisioning_messages.LoadResponse();
    // Key loading is allowed to fail due to resource restriction
    if (result != OEMCrypto_SUCCESS) {
      ASSERT_TRUE(result == OEMCrypto_ERROR_INSUFFICIENT_RESOURCES ||
                  result == OEMCrypto_ERROR_TOO_MANY_KEYS);
    }
  }
  // Verifies that the DRM keys which are already loaded should still function
  for (size_t i = 0; i < licenses.size(); i++) {
    ASSERT_NO_FATAL_FAILURE(licenses[i]->SignAndVerifyRequest());
    ASSERT_NO_FATAL_FAILURE(licenses[i]->CreateDefaultResponse());
    ASSERT_NO_FATAL_FAILURE(licenses[i]->EncryptAndSignResponse());
    ASSERT_EQ(OEMCrypto_SUCCESS, licenses[i]->LoadResponse());
    ASSERT_NO_FATAL_FAILURE(sessions[i]->TestDecryptCTR());
  }
}

// Devices that load certificates, should at least support RSA 2048 keys.
TEST_F(OEMCryptoLoadsCertificate, SupportsCertificatesAPI13) {
  // TODO(b/197141970): Need to revisit OEMCryptoLoadsCert* tests for
  // provisioning 4. Disabled here temporarily.
  if (global_features.provisioning_method == OEMCrypto_BootCertificateChain) {
    GTEST_SKIP() << "Test for non Prov 4.0 devices only.";
  }
  ASSERT_NE(0u,
            OEMCrypto_Supports_RSA_2048bit & OEMCrypto_SupportedCertificates())
      << "Supported certificates is only " << OEMCrypto_SupportedCertificates();
}

/// @}
}  // namespace wvoec