////////////////////////////////////////////////////////////////////////////////
// Copyright 2019 Google LLC.
//
// This software is licensed under the terms defined in the Widevine Master
// License Agreement. For a copy of this agreement, please contact
// widevine-licensing@google.com.
////////////////////////////////////////////////////////////////////////////////
//
// Description:
//   Unit tests for the ECPrivateKey and ECPublicKey classes.

#include "common/ec_key.h"

#include <memory>
#include <string>

#include "testing/gmock.h"
#include "testing/gunit.h"
#include "openssl/sha.h"
#include "common/ec_test_keys.h"
#include "common/ec_util.h"
#include "common/random_util.h"

using ::testing::IsNull;

namespace widevine {

class ECKeyTest : public ::testing::Test {
 public:
  static std::vector<std::tuple<std::string, std::string> > GetTestKeyList() {
    std::vector<std::tuple<std::string, std::string> > test_key_list;
    ECTestKeys test_keys;
    test_key_list.push_back(
        std::make_tuple(test_keys.private_test_key_1_secp521r1(),
                        test_keys.public_test_key_1_secp521r1()));
    test_key_list.push_back(
        std::make_tuple(test_keys.private_test_key_1_secp384r1(),
                        test_keys.public_test_key_1_secp384r1()));
    test_key_list.push_back(
        std::make_tuple(test_keys.private_test_key_1_secp256r1(),
                        test_keys.public_test_key_1_secp256r1()));
    return test_key_list;
  }

 protected:
  ECKeyTest() { plaintext_message_ = "This is a plaintext message."; }

  void GenerateEphemeralKeyPair(ECPrivateKey::EllipticCurve curve) {
    ASSERT_NE(curve, ECPrivateKey::UNDEFINED_CURVE);
    bssl::UniquePtr<EC_KEY> ephemeral_key_pair(
        EC_KEY_new_by_curve_name(ec_util::CurveToNid(curve)));
    ASSERT_TRUE(ephemeral_key_pair != nullptr);
    ASSERT_TRUE(EC_KEY_generate_key(ephemeral_key_pair.get()));
    // Separate them out into their private and public parts.
    std::string private_key;
    std::string public_key;
    ASSERT_TRUE(
        ec_util::SerializeECPrivateKey(ephemeral_key_pair.get(), &private_key));
    ASSERT_TRUE(
        ec_util::SerializeECPublicKey(ephemeral_key_pair.get(), &public_key));
    ephemeral_private_key_ = ECPrivateKey::Create(private_key);
    ephemeral_public_key_ = ECPublicKey::Create(public_key);
  }

  std::unique_ptr<ECPrivateKey> ephemeral_private_key_;
  std::unique_ptr<ECPublicKey> ephemeral_public_key_;
  std::string plaintext_message_;
};

TEST_F(ECKeyTest, KeyCurve) {
  ECTestKeys test_keys;
  EXPECT_EQ(
      ECPrivateKey::Create(test_keys.private_test_key_1_secp521r1())->Curve(),
      ECPrivateKey::SECP521R1);
  EXPECT_EQ(
      ECPrivateKey::Create(test_keys.private_test_key_1_secp384r1())->Curve(),
      ECPrivateKey::SECP384R1);
  EXPECT_EQ(
      ECPrivateKey::Create(test_keys.private_test_key_1_secp256r1())->Curve(),
      ECPrivateKey::SECP256R1);

  EXPECT_EQ(
      ECPublicKey::Create(test_keys.public_test_key_1_secp521r1())->Curve(),
      ECPrivateKey::SECP521R1);
  EXPECT_EQ(
      ECPublicKey::Create(test_keys.public_test_key_1_secp384r1())->Curve(),
      ECPrivateKey::SECP384R1);
  EXPECT_EQ(
      ECPublicKey::Create(test_keys.public_test_key_1_secp256r1())->Curve(),
      ECPrivateKey::SECP256R1);
}

TEST_F(ECKeyTest, KeyPointEncodingNullDeath) {
  EXPECT_DEATH(ephemeral_public_key_->GetPointEncodedKey(nullptr), "");
}

class ECKeyTestKeyPairs : public ECKeyTest,
                          public ::testing::WithParamInterface<
                              std::tuple<std::string, std::string> > {
 protected:
  ECKeyTestKeyPairs() {
    test_private_key_ = std::get<0>(GetParam());
    test_public_key_ = std::get<1>(GetParam());
    private_key_ = ECPrivateKey::Create(test_private_key_);
    public_key_ = ECPublicKey::Create(test_public_key_);
  }

  std::string test_private_key_;
  std::string test_public_key_;
  std::unique_ptr<ECPrivateKey> private_key_;
  std::unique_ptr<ECPublicKey> public_key_;
};

TEST_P(ECKeyTestKeyPairs, CreateWrongKey) {
  EXPECT_EQ(ECPrivateKey::Create(test_public_key_), nullptr);
  EXPECT_EQ(ECPublicKey::Create(test_private_key_), nullptr);
}

TEST_P(ECKeyTestKeyPairs, InvalidKeyConstructorParameters) {
  EXPECT_DEATH(ECPrivateKey(nullptr), "");
  EXPECT_DEATH(ECPublicKey(nullptr), "");

  EC_KEY* key;
  ASSERT_TRUE(ec_util::DeserializeECPrivateKey(test_private_key_, &key));
  // Brace initialization to avoid most vexing parse.
  EXPECT_DEATH(ECPublicKey{key}, "");
  EC_KEY_free(key);

  ASSERT_TRUE(ec_util::DeserializeECPublicKey(test_public_key_, &key));
  EXPECT_DEATH(ECPrivateKey{key}, "");
  EC_KEY_free(key);
}

TEST_P(ECKeyTestKeyPairs, KeyMatch) {
  EXPECT_TRUE(private_key_->MatchesPrivateKey(*private_key_));
  EXPECT_TRUE(private_key_->MatchesPublicKey(*public_key_));
  EXPECT_TRUE(public_key_->MatchesPrivateKey(*private_key_));
  EXPECT_TRUE(public_key_->MatchesPublicKey(*public_key_));
}

TEST_P(ECKeyTestKeyPairs, DeriveSharedSessionKey) {
  std::string derived_shared_session_key_1;
  // Generate a temporary key pair as part of ECDH.
  GenerateEphemeralKeyPair(private_key_->Curve());
  EXPECT_TRUE(private_key_->DeriveSharedSessionKey(
      *ephemeral_public_key_, &derived_shared_session_key_1));
  EXPECT_EQ(derived_shared_session_key_1.size(), SHA256_DIGEST_LENGTH);
  std::string derived_shared_session_key_2;
  EXPECT_TRUE(ephemeral_private_key_->DeriveSharedSessionKey(
      *public_key_, &derived_shared_session_key_2));
  EXPECT_EQ(derived_shared_session_key_1, derived_shared_session_key_2);
}

TEST_P(ECKeyTestKeyPairs, InvalidDeriveSharedSessionKeyParameters) {
  GenerateEphemeralKeyPair(private_key_->Curve());
  EXPECT_FALSE(
      private_key_->DeriveSharedSessionKey(*ephemeral_public_key_, nullptr));
  EXPECT_FALSE(
      ephemeral_private_key_->DeriveSharedSessionKey(*public_key_, nullptr));
}

TEST_P(ECKeyTestKeyPairs, SignVerify) {
  std::string signature;
  EXPECT_TRUE(private_key_->GenerateSignature(plaintext_message_, &signature));
  EXPECT_TRUE(public_key_->VerifySignature(plaintext_message_, signature));
}

TEST_P(ECKeyTestKeyPairs, InvalidSignVerifyParameters) {
  std::string signature;
  EXPECT_FALSE(private_key_->GenerateSignature("", &signature));
  EXPECT_FALSE(public_key_->VerifySignature("", "signature"));
  EXPECT_FALSE(private_key_->GenerateSignature(plaintext_message_, nullptr));
  EXPECT_FALSE(public_key_->VerifySignature(plaintext_message_, ""));
}

TEST_P(ECKeyTestKeyPairs, KeyMismatch) {
  GenerateEphemeralKeyPair(private_key_->Curve());
  EXPECT_FALSE(private_key_->MatchesPrivateKey(*ephemeral_private_key_));
  EXPECT_FALSE(private_key_->MatchesPublicKey(*ephemeral_public_key_));
  EXPECT_FALSE(public_key_->MatchesPrivateKey(*ephemeral_private_key_));
  EXPECT_FALSE(public_key_->MatchesPublicKey(*ephemeral_public_key_));
}

TEST_P(ECKeyTestKeyPairs, SignVerifyKeyMismatch) {
  std::string signature;
  GenerateEphemeralKeyPair(private_key_->Curve());
  EXPECT_TRUE(private_key_->GenerateSignature(plaintext_message_, &signature));
  EXPECT_FALSE(
      ephemeral_public_key_->VerifySignature(plaintext_message_, signature));
}

TEST_P(ECKeyTestKeyPairs, KeyPointEncodingCreateBadKeyPoint) {
  ASSERT_THAT(ECPublicKey::CreateFromKeyPoint(public_key_->Curve(), ""),
              IsNull());
}

TEST_P(ECKeyTestKeyPairs, KeyPointEncodingCreateBadCurve) {
  std::string serialized_key_point;
  ASSERT_TRUE(public_key_->GetPointEncodedKey(&serialized_key_point));
  ASSERT_THAT(ECPublicKey::CreateFromKeyPoint(ECPrivateKey::UNDEFINED_CURVE,
                                              serialized_key_point),
              IsNull());
}

TEST_P(ECKeyTestKeyPairs, KeyPointEncodingSuccess) {
  std::string serialized_key_point;
  ASSERT_TRUE(public_key_->GetPointEncodedKey(&serialized_key_point));
  ASSERT_EQ(ec_util::GetPublicKeyPointSize(public_key_->Curve()),
            serialized_key_point.size());
  // Check that the leading byte indicates uncompressed.
  ASSERT_EQ(4, serialized_key_point[0]);

  std::unique_ptr<ECPublicKey> new_public_key = ECPublicKey::CreateFromKeyPoint(
      public_key_->Curve(), serialized_key_point);
  ASSERT_THAT(new_public_key, testing::NotNull());
  EXPECT_TRUE(new_public_key->MatchesPublicKey(*public_key_));
  EXPECT_TRUE(new_public_key->MatchesPrivateKey(*private_key_));
};

INSTANTIATE_TEST_SUITE_P(ECKeyTestKeyPairs, ECKeyTestKeyPairs,
                         ::testing::ValuesIn(ECKeyTest::GetTestKeyList()));

class ECKeyTestCurveMismatch
    : public ECKeyTest,
      public ::testing::WithParamInterface<
          std::tuple<std::tuple<std::string, std::string>,
                     std::tuple<std::string, std::string> > > {
 protected:
  ECKeyTestCurveMismatch() {
    private_key_ = ECPrivateKey::Create(std::get<0>(std::get<0>(GetParam())));
    public_key_ = ECPublicKey::Create(std::get<1>(std::get<0>(GetParam())));
    wrong_curve_private_key_ =
        ECPrivateKey::Create(std::get<0>(std::get<1>(GetParam())));
    wrong_curve_public_key_ =
        ECPublicKey::Create(std::get<1>(std::get<1>(GetParam())));
  }

  void SetUp() override {
    // Only run combinations where the two curves differ.
    if (std::get<0>(GetParam()) == std::get<1>(GetParam())) GTEST_SKIP();
  }

  std::unique_ptr<ECPrivateKey> private_key_;
  std::unique_ptr<ECPublicKey> public_key_;
  std::unique_ptr<ECPrivateKey> wrong_curve_private_key_;
  std::unique_ptr<ECPublicKey> wrong_curve_public_key_;
};

TEST_P(ECKeyTestCurveMismatch, CurveMismatch) {
  EXPECT_FALSE(private_key_->MatchesPrivateKey(*wrong_curve_private_key_));
  EXPECT_FALSE(private_key_->MatchesPublicKey(*wrong_curve_public_key_));
  EXPECT_FALSE(public_key_->MatchesPrivateKey(*wrong_curve_private_key_));
  EXPECT_FALSE(public_key_->MatchesPublicKey(*wrong_curve_public_key_));
}

TEST_P(ECKeyTestCurveMismatch, DeriveSharedSessionKeyCurveMismatch) {
  std::string derived_shared_session_key;
  EXPECT_FALSE(private_key_->DeriveSharedSessionKey(
      *wrong_curve_public_key_, &derived_shared_session_key));
}

TEST_P(ECKeyTestCurveMismatch, SignVerifyCurveMismatch) {
  std::string signature;
  EXPECT_TRUE(private_key_->GenerateSignature(plaintext_message_, &signature));
  EXPECT_FALSE(
      wrong_curve_public_key_->VerifySignature(plaintext_message_, signature));
}

INSTANTIATE_TEST_SUITE_P(
    ECKeyTestCurveMismatch, ECKeyTestCurveMismatch,
    ::testing::Combine(::testing::ValuesIn(ECKeyTest::GetTestKeyList()),
                       ::testing::ValuesIn(ECKeyTest::GetTestKeyList())));

}  // namespace widevine