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

#include "OEMCryptoCENCCommon.h"
#include "oem_cert_test.h"
#include "oemcrypto_oem_cert.h"
#include "oemcrypto_rsa_key.h"

namespace wvoec {
namespace util {
namespace {
const std::vector<uint8_t> kOEMPrivateKeyVector(kOEMPrivateKey,
                                                kOEMPrivateKey +
                                                    kOEMPrivateKeySize);

const std::vector<uint8_t> kOEMPublicCertVector(kOEMPublicCert,
                                                kOEMPublicCert +
                                                    kOEMPublicCertSize);
}  // namespace

// Creates an OemCertificate wrapper around the built-in reference
// OEM cert.
// Creating the OemCertificate should succeed so long as the data
// is well-formed.
// Validating the OEM cert should succeed (assuming built-in cert+key
// are valid).
TEST(OEMCryptoOemCertTest, CreateFromArray) {
  std::unique_ptr<OemCertificate> oem_cert = OemCertificate::Create(
      kOEMPrivateKey, kOEMPrivateKeySize, kOEMPublicCert, kOEMPublicCertSize);
  ASSERT_TRUE(oem_cert);

  EXPECT_EQ(OemCertificate::kRsa, oem_cert->key_type());

  const std::vector<uint8_t> private_key = oem_cert->GetPrivateKey();
  EXPECT_EQ(kOEMPrivateKeyVector, private_key);

  size_t public_cert_size = 10;
  std::vector<uint8_t> public_cert(public_cert_size, 0);
  EXPECT_EQ(
      OEMCrypto_ERROR_SHORT_BUFFER,
      oem_cert->GetPublicCertificate(public_cert.data(), &public_cert_size));
  public_cert.resize(public_cert_size);
  EXPECT_EQ(OEMCrypto_SUCCESS, oem_cert->GetPublicCertificate(
                                   public_cert.data(), &public_cert_size));
  EXPECT_EQ(kOEMPublicCertSize, public_cert_size);
  EXPECT_EQ(kOEMPublicCertVector, public_cert);

  EXPECT_EQ(OEMCrypto_SUCCESS, oem_cert->IsCertificateValid());
}

TEST(OEMCryptoOemCertTest, CreateFromVector) {
  std::unique_ptr<OemCertificate> oem_cert =
      OemCertificate::Create(kOEMPrivateKeyVector, kOEMPublicCertVector);
  ASSERT_TRUE(oem_cert);

  EXPECT_EQ(OemCertificate::kRsa, oem_cert->key_type());

  const std::vector<uint8_t> private_key = oem_cert->GetPrivateKey();
  EXPECT_EQ(kOEMPrivateKeyVector, private_key);

  const std::vector<uint8_t> public_cert = oem_cert->GetPublicCertificate();
  EXPECT_EQ(kOEMPublicCertVector, public_cert);

  EXPECT_EQ(OEMCrypto_SUCCESS, oem_cert->IsCertificateValid());
}

// Creation of OemCertificate wrapper should fail if the provided
// key is not well-formed.
TEST(OEMCryptoOemCertTest, CreateWithABadPrivateKey) {
  static const uint8_t kBadPrivateKeyData[] = {'n', 'o', 't', ' ', 'a', ' ',
                                               'p', 'r', 'i', 'v', 'a', 't',
                                               'e', 'k', 'e', 'y'};
  std::unique_ptr<OemCertificate> oem_cert =
      OemCertificate::Create(kBadPrivateKeyData, sizeof(kBadPrivateKeyData),
                             kOEMPublicCert, kOEMPublicCertSize);
  EXPECT_FALSE(oem_cert);
}

// Creation of OemCertificate wrapper should fail if the provided
// OEM Public Cert is not well-formed.
TEST(OEMCryptoOemCertTest, CreateWithABadPublicCert) {
  static const uint8_t kBadPublicCert[] = {'n', 'o', 't', ' ', 'a', ' ', 'o',
                                           'e', 'm', ' ', 'p', 'u', 'b', 'l',
                                           'i', 'c', ' ', 'c', 'e', 'r', 't'};
  std::unique_ptr<OemCertificate> oem_cert =
      OemCertificate::Create(kOEMPrivateKey, kOEMPrivateKeySize, kBadPublicCert,
                             sizeof(kBadPublicCert));
  EXPECT_FALSE(oem_cert);
}

// It is possible to create an OEM Certificate using a non-matching
// public-private key pair so long as the key types are the same.
// However, OEM Cert validation should catch the problem.
TEST(OEMCryptoOemCertTest, CreateWithDifferentPrivateRsaKey) {
  std::unique_ptr<RsaPrivateKey> key = RsaPrivateKey::New(kRsa2048Bit);
  ASSERT_TRUE(key);
  const std::vector<uint8_t> private_key = key->Serialize();
  ASSERT_FALSE(private_key.empty());

  // Creating the OEM Certificate should succeed.
  std::unique_ptr<OemCertificate> oem_cert =
      OemCertificate::Create(private_key, kOEMPublicCertVector);
  ASSERT_TRUE(oem_cert);

  EXPECT_EQ(OemCertificate::kRsa, oem_cert->key_type());

  // Validating key should return an error.
  OEMCryptoResult status = oem_cert->IsCertificateValid();
  // Test still allows deprecated error.
  if (status != OEMCrypto_ERROR_INVALID_RSA_KEY) {
    EXPECT_EQ(OEMCrypto_ERROR_INVALID_KEY, status);
  }
}
}  // namespace util
}  // namespace wvoec