////////////////////////////////////////////////////////////////////////////////
// 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 ec_util.

#include "common/ec_util.h"

#include <stddef.h>

#include <memory>

#include "glog/logging.h"
#include "testing/gmock.h"
#include "testing/gunit.h"
#include "common/ec_test_keys.h"
#include "common/openssl_util.h"

namespace widevine {
namespace ec_util {

class ECUtilTest : public ::testing::Test {
 protected:
  ECTestKeys test_keys_;

  // Given serialization and deserialization methods for an EC_KEY and a test
  // key, this checks that that those methods can deserialize to a valid EC_KEY
  // and serialize to the same test key.
  template <bool (*SerializeECKey)(const EC_KEY*, std::string*),
            bool (*DeserializeECKey)(const std::string&, EC_KEY** private_key)>
  void SerializeDeserializeECKeyCheck(const std::string& test_key) {
    EC_KEY* key = nullptr;
    std::string serialized_key;
    EXPECT_TRUE(DeserializeECKey(test_key, &key));
    ASSERT_TRUE(key != nullptr);
    EXPECT_TRUE(SerializeECKey(key, &serialized_key));
    EXPECT_EQ(serialized_key, test_key);
    EXPECT_EQ(EC_KEY_check_key(key), 1);
    EC_KEY_free(key);
  }

  // Given serialization and deserialization methods for an EC_KEY and a test
  // key, this method checks that deserialization fails for a bad input std::string
  // and both serialization and deserialization fail on nullptrs.
  template <bool (*SerializeECKey)(const EC_KEY*, std::string*),
            bool (*DeserializeECKey)(const std::string&, EC_KEY** private_key)>
  void SerializeDeserializeECKeyValidateInput(const std::string& test_key) {
    EC_KEY* key = nullptr;
    // Invalid key
    EXPECT_FALSE(DeserializeECKey("this is a bad key", &key));
    EXPECT_EQ(key, nullptr);
    // Invalid pointers
    EXPECT_TRUE(DeserializeECKey(test_key, &key));
    EXPECT_FALSE(SerializeECKey(key, nullptr));
    EC_KEY_free(key);
    EXPECT_FALSE(DeserializeECKey(test_key, nullptr));
  }
};

TEST_F(ECUtilTest, SerializeDeserializePrivateKey) {
  SerializeDeserializeECKeyCheck<SerializeECPrivateKey,
                                 DeserializeECPrivateKey>(
      test_keys_.private_test_key_1_secp521r1());
  SerializeDeserializeECKeyCheck<SerializeECPrivateKey,
                                 DeserializeECPrivateKey>(
      test_keys_.private_test_key_1_secp384r1());
  SerializeDeserializeECKeyCheck<SerializeECPrivateKey,
                                 DeserializeECPrivateKey>(
      test_keys_.private_test_key_1_secp256r1());
  SerializeDeserializeECKeyValidateInput<SerializeECPrivateKey,
                                         DeserializeECPrivateKey>(
      test_keys_.private_test_key_1_secp521r1());
}

TEST_F(ECUtilTest, SerializeDeserializePublicKey) {
  SerializeDeserializeECKeyCheck<SerializeECPublicKey, DeserializeECPublicKey>(
      test_keys_.public_test_key_1_secp521r1());
  SerializeDeserializeECKeyCheck<SerializeECPublicKey, DeserializeECPublicKey>(
      test_keys_.public_test_key_1_secp384r1());
  SerializeDeserializeECKeyCheck<SerializeECPublicKey, DeserializeECPublicKey>(
      test_keys_.public_test_key_1_secp256r1());
  SerializeDeserializeECKeyValidateInput<SerializeECPublicKey,
                                         DeserializeECPublicKey>(
      test_keys_.public_test_key_1_secp521r1());
}

TEST_F(ECUtilTest, PublicKeyExtraction) {
  EC_KEY* private_key = nullptr;
  std::string serialized_public_key;
  // Key 1
  EXPECT_TRUE(DeserializeECPrivateKey(test_keys_.private_test_key_1_secp521r1(),
                                      &private_key));
  ASSERT_TRUE(private_key != nullptr);
  EXPECT_TRUE(SerializeECPublicKey(private_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_1_secp521r1());
  EC_KEY_free(private_key);
  private_key = nullptr;
  // Key 2
  EXPECT_TRUE(DeserializeECPrivateKey(test_keys_.private_test_key_1_secp384r1(),
                                      &private_key));
  ASSERT_TRUE(private_key != nullptr);
  EXPECT_TRUE(SerializeECPublicKey(private_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_1_secp384r1());
  EC_KEY_free(private_key);
  private_key = nullptr;
  // Key 3
  EXPECT_TRUE(DeserializeECPrivateKey(test_keys_.private_test_key_1_secp256r1(),
                                      &private_key));
  ASSERT_TRUE(private_key != nullptr);
  EXPECT_TRUE(SerializeECPublicKey(private_key, &serialized_public_key));
  EXPECT_EQ(serialized_public_key, test_keys_.public_test_key_1_secp256r1());
  EC_KEY_free(private_key);
}

TEST_F(ECUtilTest, GetPublicKeyPointSizeAll) {
  ASSERT_GT(ec_util::GetPublicKeyPointSize(ECPrivateKey::SECP256R1), 0);
  ASSERT_GT(ec_util::GetPublicKeyPointSize(ECPrivateKey::SECP384R1), 0);
  ASSERT_GT(ec_util::GetPublicKeyPointSize(ECPrivateKey::SECP521R1), 0);

  ASSERT_EQ(0, ec_util::GetPublicKeyPointSize(ECPrivateKey::UNDEFINED_CURVE));
}

TEST_F(ECUtilTest, GetPublicKeyFromKeyPointDeath) {
  EXPECT_DEATH(GetPublicKeyFromKeyPoint(ECPrivateKey::SECP256R1, "", nullptr),
               "");
}

class ECKeyPointTest : public ::testing::Test,
                       public ::testing::WithParamInterface<std::string> {
 public:
  static std::vector<std::string> GetTestKeyList() {
    ECTestKeys test_keys;
    return {test_keys.public_test_key_1_secp256r1(),
            test_keys.public_test_key_1_secp384r1(),
            test_keys.public_test_key_1_secp521r1()};
  }

  void SetUp() override {
    serialized_public_key_ = GetParam();

    EC_KEY* public_key = nullptr;
    ASSERT_TRUE(DeserializeECPublicKey(serialized_public_key_, &public_key));
    ASSERT_TRUE(public_key != nullptr);
    ec_public_key_.reset(public_key);

    curve_ = NidToCurve(
        EC_GROUP_get_curve_name(EC_KEY_get0_group(ec_public_key_.get())));
  }

 protected:
  std::string serialized_public_key_;
  ScopedECKEY ec_public_key_;
  ECPrivateKey::EllipticCurve curve_;
};

INSTANTIATE_TEST_SUITE_P(ECUtilKeyPointTest, ECKeyPointTest,
                         ::testing::ValuesIn(ECKeyPointTest::GetTestKeyList()));

TEST_P(ECKeyPointTest, GetPublicKeyPointDeath) {
  EXPECT_DEATH(GetPublicKeyPoint(ec_public_key_.get(), nullptr), "");
}

TEST_P(ECKeyPointTest, GetPublicKeyFromKeyPointInvalidCurve) {
  std::string encoded_key;
  ASSERT_TRUE(GetPublicKeyPoint(ec_public_key_.get(), &encoded_key));
  EC_KEY* new_public_key = nullptr;
  EXPECT_FALSE(GetPublicKeyFromKeyPoint(ECPrivateKey::UNDEFINED_CURVE,
                                        encoded_key, &new_public_key));
}

TEST_P(ECKeyPointTest, GetPublicKeyFromKeyPointInvalidSize) {
  EC_KEY* new_public_key = nullptr;
  EXPECT_FALSE(GetPublicKeyFromKeyPoint(curve_, "", &new_public_key));
}

TEST_P(ECKeyPointTest, PublicKeyPointFormatTest) {
  std::string encoded_key;
  ASSERT_TRUE(GetPublicKeyPoint(ec_public_key_.get(), &encoded_key));
  ASSERT_FALSE(encoded_key.empty());
  EXPECT_EQ(4, encoded_key[0]);
  EXPECT_EQ(GetPublicKeyPointSize(curve_), encoded_key.size());

  // Create a new EC key from the key point.
  EC_KEY* ec_key = nullptr;
  ASSERT_TRUE(GetPublicKeyFromKeyPoint(curve_, encoded_key, &ec_key));
  ASSERT_TRUE(ec_key != nullptr);
  ScopedECKEY new_public_key(ec_key);

  // EC_POINT_cmp returns 0 if the keys match.
  ASSERT_EQ(
      0, EC_POINT_cmp(EC_KEY_get0_group(ec_public_key_.get()),
                      EC_KEY_get0_public_key(ec_public_key_.get()),
                      EC_KEY_get0_public_key(new_public_key.get()), nullptr));
}

}  // namespace ec_util
}  // namespace widevine