// 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 <gtest/gtest.h>

#include "OEMCryptoCENCCommon.h"
#include "oemcrypto_ecc_key.h"
#include "oemcrypto_ref_test_utils.h"

namespace wvoec_ref {
constexpr size_t kMessageSize = 4 * 1024;  // 4 kB

class OEMCryptoEccKeyTest : public ::testing::TestWithParam<EccCurve> {
 public:
  void SetUp() override { key_ = EccPrivateKey::New(GetParam()); }

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

 protected:
  std::unique_ptr<EccPrivateKey> key_;
};

// Basic verification of ECC private key generation.
TEST_P(OEMCryptoEccKeyTest, KeyProperties) {
  ASSERT_TRUE(key_);
  const EccCurve expected_curve = GetParam();

  EXPECT_EQ(key_->curve(), expected_curve);
  EXPECT_NE(nullptr, key_->GetEcKey());
}

// 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(OEMCryptoEccKeyTest, 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 = " << 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(OEMCryptoEccKeyTest, SerializeAndReloadPrivateKey) {
  ASSERT_TRUE(key_);

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

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

  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 = " << i;
  }
}

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

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

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

// 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(OEMCryptoEccKeyTest, SerializePublicKey) {
  ASSERT_TRUE(key_);

  std::unique_ptr<EccPublicKey> 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 = " << 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.
// It is anticipated that OEMCrypto will need to parse the licensing
// server's ephemerial key when deriving the session key.
TEST_P(OEMCryptoEccKeyTest, SerializeAndReloadPublicKey) {
  ASSERT_TRUE(key_);

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

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

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

  EXPECT_EQ(pub_key->curve(), loaded_key->curve());

  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 = " << i;
  }
}

// Checks that the ECC signature generating API operates similar to
// existing signature generation functions.
TEST_P(OEMCryptoEccKeyTest, 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(),
                                    signature.data(), &signature_size));
  EXPECT_GT(signature_size, kInitialBufferSize);

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

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

// Checks that ECC signatures can be verified by an ECC public key.
TEST_P(OEMCryptoEccKeyTest, 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<EccPublicKey> 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));
}

// Verifies the session key exchange protocol used by the licensing
// server.
TEST_P(OEMCryptoEccKeyTest, DeriveSessionKey) {
  ASSERT_TRUE(key_);
  // Set up Alice.
  EccPrivateKey* alice_private_key = key_.get();
  std::unique_ptr<EccPublicKey> alice_public_key =
      alice_private_key->MakePublicKey();
  ASSERT_TRUE(alice_public_key);
  // Set up Bob.
  std::unique_ptr<EccPrivateKey> bob_private_key =
      EccPrivateKey::New(alice_private_key->curve());
  ASSERT_TRUE(bob_private_key);
  std::unique_ptr<EccPublicKey> bob_public_key =
      bob_private_key->MakePublicKey();
  ASSERT_TRUE(bob_public_key);

  const size_t session_key_length = alice_private_key->SessionKeyLength();
  EXPECT_EQ(session_key_length, bob_private_key->SessionKeyLength());
  // From Alice's perspective.
  const std::vector<uint8_t> alice_session_key =
      alice_private_key->DeriveSessionKey(*bob_public_key);

  // From Bob's perspective.
  const std::vector<uint8_t> bob_session_key =
      bob_private_key->DeriveSessionKey(*alice_public_key);

  // Both should have the same session key.
  ASSERT_EQ(session_key_length, alice_session_key.size());
  ASSERT_EQ(session_key_length, bob_session_key.size());
  for (size_t i = 0; i < session_key_length; i++) {
    ASSERT_EQ(alice_session_key[i], bob_session_key[i]) << "i = " << i;
  }
}

INSTANTIATE_TEST_CASE_P(AllCurves, OEMCryptoEccKeyTest,
                        ::testing::Values(kEccSecp256r1, kEccSecp384r1,
                                          kEccSecp521r1));

}  // namespace wvoec_ref