// 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 of OEMCrypto APIs
//
#include <stdlib.h>

#include <mutex>

#include <gtest/gtest.h>

#include "OEMCryptoCENC.h"
#include "log.h"
#include "oemcrypto_ref_test_utils.h"
#include "oemcrypto_rsa_key.h"

namespace wvoec_ref {
constexpr size_t kMessageSize = 4 * 1024;  // 4 kB
constexpr size_t kCastMessageSize = 83;    // Special max size.

class OEMCryptoRsaKeyTest : public ::testing::TestWithParam<RsaFieldSize> {
 public:
  void SetUp() override {
    // RSA key generation is slow (~2 seconds) compared to the
    // operations they perform (<50 ms).  Each key type is generated
    // once and globally stored in serialized form.
    // Caching the instance may result in test failures for
    // memory-leak detection.
    const RsaFieldSize field_size = GetParam();
    std::lock_guard<std::mutex> rsa_key_lock(rsa_key_mutex_);
    // Use of a switch case is intentional to cause compiler warnings
    // if a new field size is introduced without updating the test.
    switch (field_size) {
      case kRsa2048Bit: {
        if (!rsa_2048_key_data_.empty()) {
          key_ = RsaPrivateKey::Load(rsa_2048_key_data_);
        }
        if (!key_) {
          key_ = RsaPrivateKey::New(kRsa2048Bit);
        }
        if (rsa_2048_key_data_.empty() && key_) {
          rsa_2048_key_data_ = key_->Serialize();
        }
      } break;
      case kRsa3072Bit: {
        if (!rsa_3072_key_data_.empty()) {
          key_ = RsaPrivateKey::Load(rsa_3072_key_data_);
        }
        if (!key_) {
          key_ = RsaPrivateKey::New(kRsa3072Bit);
        }
        if (rsa_3072_key_data_.empty() && key_) {
          rsa_3072_key_data_ = key_->Serialize();
        }
      } break;
      case kRsaFieldUnknown:  // Suppress compiler warnings
        LOGE("RSA test was incorrectly instantiation");
        exit(EXIT_FAILURE);
        break;
    }
  }

  void TearDown() override { key_.reset(); }

 protected:
  std::unique_ptr<RsaPrivateKey> key_;
  static std::mutex rsa_key_mutex_;
  static std::vector<uint8_t> rsa_2048_key_data_;
  static std::vector<uint8_t> rsa_3072_key_data_;
};

std::mutex OEMCryptoRsaKeyTest::rsa_key_mutex_;
std::vector<uint8_t> OEMCryptoRsaKeyTest::rsa_2048_key_data_;
std::vector<uint8_t> OEMCryptoRsaKeyTest::rsa_3072_key_data_;

// Basic verification of RSA private key generation.
TEST_P(OEMCryptoRsaKeyTest, KeyProperties) {
  ASSERT_TRUE(key_);
  const RsaFieldSize expected_field_size = GetParam();

  EXPECT_EQ(key_->field_size(), expected_field_size);
  EXPECT_NE(nullptr, key_->GetRsaKey());
}

// Checks that the private key serialization APIs are compatible
// and performing in a manner that is similar to other OEMCrypto methods
// that retrieve data.
TEST_P(OEMCryptoRsaKeyTest, SerializePrivateKey) {
  ASSERT_TRUE(key_);

  constexpr size_t kInitialBufferSize = 10;  // Definitely too small.
  size_t buffer_size = kInitialBufferSize;
  std::vector<uint8_t> buffer(buffer_size);

  EXPECT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
            key_->Serialize(buffer.data(), &buffer_size));
  EXPECT_GT(buffer_size, kInitialBufferSize);

  buffer.resize(buffer_size);
  EXPECT_EQ(OEMCrypto_SUCCESS, key_->Serialize(buffer.data(), &buffer_size));
  buffer.resize(buffer_size);

  const std::vector<uint8_t> direct_key_data = key_->Serialize();
  EXPECT_FALSE(direct_key_data.empty());
  ASSERT_EQ(buffer.size(), direct_key_data.size());
  for (size_t i = 0; i < buffer.size(); i++) {
    ASSERT_EQ(buffer[i], direct_key_data[i]) << "i = " << std::to_string(i);
  }
}

// Checks that a private key that is serialized can be deserialized and
// reload.  Also checks that the serialization of a key produces the
// same data to ensure consistency.
TEST_P(OEMCryptoRsaKeyTest, SerializeAndReloadPrivateKey) {
  ASSERT_TRUE(key_);

  const std::vector<uint8_t> key_data = key_->Serialize();
  std::unique_ptr<RsaPrivateKey> loaded_key = RsaPrivateKey::Load(key_data);
  ASSERT_TRUE(loaded_key);

  EXPECT_EQ(key_->field_size(), loaded_key->field_size());

  const std::vector<uint8_t> loaded_key_data = loaded_key->Serialize();
  ASSERT_EQ(key_data.size(), loaded_key_data.size());
  for (size_t i = 0; i < key_data.size(); i++) {
    ASSERT_EQ(key_data[i], loaded_key_data[i]) << "i = " << std::to_string(i);
  }
}

// Checks that a private key with explicitly indicated schemes include
// the scheme fields in the reserialized key.
TEST_P(OEMCryptoRsaKeyTest, SerializeAndReloadPrivateKeyWithAllowedSchemes) {
  ASSERT_TRUE(key_);

  const std::vector<uint8_t> raw_key_data = key_->Serialize();
  std::vector<uint8_t> key_data = {'S', 'I', 'G', 'N', 0x00, 0x00, 0x00, 0x03};
  key_data.insert(key_data.end(), raw_key_data.begin(), raw_key_data.end());

  std::unique_ptr<RsaPrivateKey> explicit_key = RsaPrivateKey::Load(key_data);
  ASSERT_TRUE(explicit_key);
  EXPECT_EQ(key_->field_size(), explicit_key->field_size());
  EXPECT_EQ(explicit_key->allowed_schemes(), 0x03);

  const std::vector<uint8_t> explicit_key_data = explicit_key->Serialize();
  ASSERT_EQ(key_data.size(), explicit_key_data.size());
  ASSERT_EQ(key_data, explicit_key_data);
}

// Checks that a public key can be created from the private key.
TEST_P(OEMCryptoRsaKeyTest, DerivePublicKey) {
  ASSERT_TRUE(key_);

  std::unique_ptr<RsaPublicKey> pub_key = key_->MakePublicKey();
  ASSERT_TRUE(pub_key);

  EXPECT_TRUE(key_->IsMatchingPublicKey(*pub_key));
}

// Checks that the public key serialization APIs are compatible
// and performing in a manner that is similar to other OEMCrypto methods
// that retrieve data.
TEST_P(OEMCryptoRsaKeyTest, SerializePublicKey) {
  ASSERT_TRUE(key_);

  std::unique_ptr<RsaPublicKey> pub_key = key_->MakePublicKey();
  ASSERT_TRUE(pub_key);

  constexpr size_t kInitialBufferSize = 10;  // Definitely too small.
  size_t buffer_size = kInitialBufferSize;
  std::vector<uint8_t> buffer(buffer_size);

  EXPECT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
            pub_key->Serialize(buffer.data(), &buffer_size));
  EXPECT_GT(buffer_size, kInitialBufferSize);

  buffer.resize(buffer_size);
  EXPECT_EQ(OEMCrypto_SUCCESS, pub_key->Serialize(buffer.data(), &buffer_size));
  buffer.resize(buffer_size);

  const std::vector<uint8_t> direct_key_data = pub_key->Serialize();
  EXPECT_FALSE(direct_key_data.empty());
  ASSERT_EQ(buffer.size(), direct_key_data.size());
  for (size_t i = 0; i < buffer.size(); i++) {
    ASSERT_EQ(buffer[i], direct_key_data[i]) << "i = " << std::to_string(i);
  }
}

// Checks that a public key that is serialized can be deserialized and
// reload.  Also checks that the serialization of a key produces the
// same data to ensure consistency.
TEST_P(OEMCryptoRsaKeyTest, SerializeAndReloadPublicKey) {
  ASSERT_TRUE(key_);

  std::unique_ptr<RsaPublicKey> pub_key = key_->MakePublicKey();
  ASSERT_TRUE(pub_key);

  const std::vector<uint8_t> key_data = pub_key->Serialize();

  std::unique_ptr<RsaPublicKey> loaded_key = RsaPublicKey::Load(key_data);
  ASSERT_TRUE(loaded_key);

  EXPECT_EQ(pub_key->field_size(), loaded_key->field_size());
  EXPECT_EQ(pub_key->allowed_schemes(), loaded_key->allowed_schemes());

  const std::vector<uint8_t> loaded_key_data = loaded_key->Serialize();
  ASSERT_EQ(key_data.size(), loaded_key_data.size());
  for (size_t i = 0; i < key_data.size(); i++) {
    ASSERT_EQ(key_data[i], loaded_key_data[i]) << "i = " << std::to_string(i);
  }
}

// Checks that the RSA signature generating API operates similar to
// existing signature generation functions.
TEST_P(OEMCryptoRsaKeyTest, GenerateSignature) {
  ASSERT_TRUE(key_);

  const std::vector<uint8_t> message = RandomData(kMessageSize);
  ASSERT_FALSE(message.empty()) << "CdmRandom failed";

  constexpr size_t kInitialBufferSize = 10;  // Definitely too small.
  size_t signature_size = kInitialBufferSize;
  std::vector<uint8_t> signature(signature_size);
  EXPECT_EQ(
      OEMCrypto_ERROR_SHORT_BUFFER,
      key_->GenerateSignature(message.data(), message.size(), kRsaPssDefault,
                              signature.data(), &signature_size));
  EXPECT_GT(signature_size, kInitialBufferSize);

  signature.resize(signature_size);
  EXPECT_EQ(
      OEMCrypto_SUCCESS,
      key_->GenerateSignature(message.data(), message.size(), kRsaPssDefault,
                              signature.data(), &signature_size));
  signature.resize(signature_size);

  EXPECT_LE(signature_size, key_->SignatureSize());
}

// Checks that RSA signatures can be verified by an RSA public key.
TEST_P(OEMCryptoRsaKeyTest, VerifySignature) {
  ASSERT_TRUE(key_);
  const std::vector<uint8_t> message = RandomData(kMessageSize);
  ASSERT_FALSE(message.empty()) << "CdmRandom failed";

  const std::vector<uint8_t> signature = key_->GenerateSignature(message);

  std::unique_ptr<RsaPublicKey> pub_key = key_->MakePublicKey();
  ASSERT_TRUE(pub_key);

  EXPECT_EQ(OEMCrypto_SUCCESS, pub_key->VerifySignature(message, signature));

  // Check with different message.
  const std::vector<uint8_t> message_two = RandomData(kMessageSize);
  EXPECT_EQ(OEMCrypto_ERROR_SIGNATURE_FAILURE,
            pub_key->VerifySignature(message_two, signature));

  // Check with bad signature.
  const std::vector<uint8_t> bad_signature = RandomData(signature.size());
  EXPECT_EQ(OEMCrypto_ERROR_SIGNATURE_FAILURE,
            pub_key->VerifySignature(message, bad_signature));
}

// Checks that the special CAST receiver signature scheme works
// to the degree that it is possible to test.
TEST_P(OEMCryptoRsaKeyTest, GenerateAndVerifyRsaSignature) {
  ASSERT_TRUE(key_);
  // Key must be enabled for PKCS1 Block 1 padding scheme.
  // To do so, the key is serialized and the padding scheme is
  // added to the key data.
  const std::vector<uint8_t> key_data = key_->Serialize();
  ASSERT_FALSE(key_data.empty());
  std::vector<uint8_t> pkcs_enabled_key_data = {
      'S', 'I', 'G', 'N', 0x00, 0x00, 0x00, kSign_PKCS1_Block1};
  pkcs_enabled_key_data.insert(pkcs_enabled_key_data.end(), key_data.begin(),
                               key_data.end());
  std::unique_ptr<RsaPrivateKey> pkcs_enabled_key =
      RsaPrivateKey::Load(pkcs_enabled_key_data);
  ASSERT_TRUE(pkcs_enabled_key);

  // The actual cast message is a domain specific hash of the message,
  // however, random data works for testing purposes.
  const std::vector<uint8_t> message = RandomData(kCastMessageSize);

  // Generate signature.
  constexpr size_t kInitialBufferSize = 10;  // Definitely too small.
  size_t signature_size = kInitialBufferSize;
  std::vector<uint8_t> signature(signature_size);
  EXPECT_EQ(OEMCrypto_ERROR_SHORT_BUFFER,
            pkcs_enabled_key->GenerateSignature(message.data(), message.size(),
                                                kRsaPkcs1Cast, signature.data(),
                                                &signature_size));
  EXPECT_GT(signature_size, kInitialBufferSize);
  signature.resize(signature_size);
  EXPECT_EQ(OEMCrypto_SUCCESS,
            pkcs_enabled_key->GenerateSignature(message.data(), message.size(),
                                                kRsaPkcs1Cast, signature.data(),
                                                &signature_size));
  signature.resize(signature_size);

  EXPECT_LE(signature_size, pkcs_enabled_key->SignatureSize());

  // Verify signature.
  std::unique_ptr<RsaPublicKey> pub_key = pkcs_enabled_key->MakePublicKey();
  ASSERT_TRUE(pub_key);

  EXPECT_EQ(OEMCrypto_SUCCESS,
            pub_key->VerifySignature(message, signature, kRsaPkcs1Cast));
}

// Verifies the session key exchange protocol used by the licensing
// server.
TEST_P(OEMCryptoRsaKeyTest, ShareSessionKey) {
  constexpr size_t kSessionKeySize = 16;
  ASSERT_TRUE(key_);
  std::unique_ptr<RsaPublicKey> public_key = key_->MakePublicKey();
  ASSERT_TRUE(public_key);

  // Generate session key.
  const std::vector<uint8_t> session_key = RandomData(kSessionKeySize);
  ASSERT_FALSE(session_key.empty());

  // Server's perspective.
  constexpr size_t kInitialBufferSize = 10;  // Definitely too small.
  size_t enc_session_key_size = kInitialBufferSize;
  std::vector<uint8_t> enc_session_key(enc_session_key_size);
  OEMCryptoResult result = public_key->EncryptSessionKey(
      session_key.data(), session_key.size(), enc_session_key.data(),
      &enc_session_key_size);
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, result);
  enc_session_key.resize(enc_session_key_size);
  result = public_key->EncryptSessionKey(session_key.data(), session_key.size(),
                                         enc_session_key.data(),
                                         &enc_session_key_size);
  ASSERT_EQ(OEMCrypto_SUCCESS, result);
  enc_session_key.resize(enc_session_key_size);

  // Client's perspective.
  size_t received_session_key_size = kInitialBufferSize;
  std::vector<uint8_t> received_session_key(received_session_key_size);
  result = key_->DecryptSessionKey(
      enc_session_key.data(), enc_session_key.size(),
      received_session_key.data(), &received_session_key_size);
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, result);
  received_session_key.resize(received_session_key_size);
  result = key_->DecryptSessionKey(
      enc_session_key.data(), enc_session_key.size(),
      received_session_key.data(), &received_session_key_size);
  ASSERT_EQ(OEMCrypto_SUCCESS, result);
  received_session_key.resize(received_session_key_size);

  // Compare keys.
  ASSERT_EQ(session_key.size(), received_session_key.size());
  ASSERT_EQ(session_key, received_session_key);
}

// Verifies the encryption key exchange protocol used by the licensing
// server.
TEST_P(OEMCryptoRsaKeyTest, ShareEncryptionKey) {
  constexpr size_t kEncryptionKeySize = 16;
  ASSERT_TRUE(key_);
  std::unique_ptr<RsaPublicKey> public_key = key_->MakePublicKey();
  ASSERT_TRUE(public_key);

  // Generate session key.
  const std::vector<uint8_t> encryption_key = RandomData(kEncryptionKeySize);
  ASSERT_FALSE(encryption_key.empty());

  // Server's perspective.
  constexpr size_t kInitialBufferSize = 10;  // Definitely too small.
  size_t enc_encryption_key_size = kInitialBufferSize;
  std::vector<uint8_t> enc_encryption_key(enc_encryption_key_size);
  OEMCryptoResult result = public_key->EncryptEncryptionKey(
      encryption_key.data(), encryption_key.size(), enc_encryption_key.data(),
      &enc_encryption_key_size);
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, result);
  enc_encryption_key.resize(enc_encryption_key_size);
  result = public_key->EncryptEncryptionKey(
      encryption_key.data(), encryption_key.size(), enc_encryption_key.data(),
      &enc_encryption_key_size);
  ASSERT_EQ(OEMCrypto_SUCCESS, result);
  enc_encryption_key.resize(enc_encryption_key_size);

  // Client's perspective.
  size_t received_encryption_key_size = kInitialBufferSize;
  std::vector<uint8_t> received_encryption_key(received_encryption_key_size);
  result = key_->DecryptEncryptionKey(
      enc_encryption_key.data(), enc_encryption_key.size(),
      received_encryption_key.data(), &received_encryption_key_size);
  ASSERT_EQ(OEMCrypto_ERROR_SHORT_BUFFER, result);
  received_encryption_key.resize(received_encryption_key_size);
  result = key_->DecryptEncryptionKey(
      enc_encryption_key.data(), enc_encryption_key.size(),
      received_encryption_key.data(), &received_encryption_key_size);
  ASSERT_EQ(OEMCrypto_SUCCESS, result);
  received_encryption_key.resize(received_encryption_key_size);

  // Compare keys.
  ASSERT_EQ(encryption_key.size(), received_encryption_key.size());
  ASSERT_EQ(encryption_key, received_encryption_key);
}

INSTANTIATE_TEST_CASE_P(AllFieldSizes, OEMCryptoRsaKeyTest,
                        ::testing::Values(kRsa2048Bit, kRsa3072Bit));

}  // namespace wvoec_ref