// 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
//
#ifndef WVOEC_UTIL_ECC_KEY_H_
#define WVOEC_UTIL_ECC_KEY_H_

#include <stddef.h>
#include <stdint.h>

#include <memory>
#include <string>
#include <vector>

#include <openssl/ec.h>

#include "OEMCryptoCENCCommon.h"

namespace wvoec {
namespace util {
enum EccCurve {
  kEccCurveUnknown = 0,
  kEccSecp256r1 = 256,
  kEccSecp384r1 = 384,
  kEccSecp521r1 = 521
};

// Returns the string representation of the provided curve.
// Intended for logging purposes.
std::string EccCurveToString(EccCurve curve);

class EccPrivateKey;

class EccPublicKey {
 public:
  // Creates a new public key equivalent of the provided private key.
  static std::unique_ptr<EccPublicKey> New(const EccPrivateKey& private_key);

  // Loads a serialized EC public key.
  // The provided |buffer| must contain a valid ASN.1 DER encoded
  // SubjectPublicKey.  Only supported curves by this API are those
  // enumerated by EccCurve.
  //
  //  buffer: SubjectPublicKeyInfo = {
  //    algorithm: AlgorithmIdentifier = {
  //      algorithm: OID = id-ecPublicKey,
  //      parameters: ECParameters = {
  //        namedCurve: OID = secp256r1 | secp384r1 | secp521r1
  //      }
  //    },
  //    subjectPublicKey: BIT STRING = ... -- SEC1 encoded ECPoint
  //  }
  //
  // Failure will occur if the provided |buffer| does not contain a
  // valid SubjectPublicKey, or if the specified curve is not
  // supported.
  static std::unique_ptr<EccPublicKey> Load(const uint8_t* buffer,
                                            size_t length);
  static std::unique_ptr<EccPublicKey> Load(const std::string& buffer);
  static std::unique_ptr<EccPublicKey> Load(const std::vector<uint8_t>& buffer);
  // Loads EC public key from the |curve| and |buffer|.
  // The provided |buffer| must contain an EC point serialized from raw X9.62
  // format. For uncompressed form, it is a 1-byte prefix plus two 32-byte
  // integers representing X, Y coordinates.
  static std::unique_ptr<EccPublicKey> LoadKeyPoint(EccCurve curve,
                                                    const uint8_t* buffer,
                                                    size_t length);

  // Loads a serialized ECC private key, but only converting the public key.
  static std::unique_ptr<EccPublicKey> LoadPrivateKeyInfo(const uint8_t* buffer,
                                                          size_t length);
  static std::unique_ptr<EccPublicKey> LoadPrivateKeyInfo(
      const std::string& buffer);
  static std::unique_ptr<EccPublicKey> LoadPrivateKeyInfo(
      const std::vector<uint8_t>& buffer);

  EccCurve curve() const { return curve_; }
  const EC_KEY* GetEcKey() const { return key_; }

  // Checks if the provided |private_key| is the EC private key of this
  // public key.
  bool IsMatchingPrivateKey(const EccPrivateKey& private_key) const;

  // Serializes the public key into an ASN.1 DER encoded SubjectPublicKey
  // representation.
  // On success, |buffer_size| is populated with the number of bytes
  // written to |buffer|, and OEMCrypto_SUCCESS is returned.
  // If the provided |buffer_size| is too small, ERROR_SHORT_BUFFER
  // is returned and |buffer_size| is set to the required buffer size.
  OEMCryptoResult Serialize(uint8_t* buffer, size_t* buffer_size) const;
  // Same as above, except directly returns the serialized key.
  // Returns an empty vector on error.
  std::vector<uint8_t> Serialize() const;

  // Verifies the |signature| matches the provided |message| by the
  // private equivalent of this public key.
  // The |signature| should be a valid ASN.1 DER encoded
  // ECDSA-Sig-Value.
  // This implementation uses ECDSA with the following digest
  // algorithms for the supported curve types.
  //  - SHA-256 / secp256r1
  //  - SHA-384 / secp384r1 (optional support)
  //  - SHA-512 / secp521r1 (optional support)
  // Returns:
  //    OEMCrypto_SUCCESS if signature is valid
  //    OEMCrypto_ERROR_SIGNATURE_FAILURE if the signature is invalid
  //    Any other result indicates an unexpected error
  OEMCryptoResult VerifySignature(const uint8_t* message, size_t message_length,
                                  const uint8_t* signature,
                                  size_t signature_length) const;
  OEMCryptoResult VerifySignature(const std::string& message,
                                  const std::string& signature) const;
  OEMCryptoResult VerifySignature(const std::vector<uint8_t>& message,
                                  const std::vector<uint8_t>& signature) const;
  // Verifies the raw |signature| matches the provided |message| by the
  // private equivalent of this public key.
  // A raw ECDSA signature consists of a pair of integers (r,s). The |signature|
  // is a concatenation of two octet strings resulting from the integer-to-octet
  // encoding of the values of r and s, in the order of (r||s).
  OEMCryptoResult VerifyRawSignature(const uint8_t* message,
                                     size_t message_length,
                                     const uint8_t* signature,
                                     size_t signature_length) const;

  ~EccPublicKey();

  EccPublicKey(const EccPublicKey&) = delete;
  EccPublicKey(EccPublicKey&&) = delete;
  const EccPublicKey& operator=(const EccPublicKey&) = delete;
  EccPublicKey& operator=(EccPublicKey&&) = delete;

 private:
  EccPublicKey() {}

  // Initializes the public key object using the provided |buffer|.
  // In case of any failure, false is return and the key should be
  // discarded.
  bool InitFromSubjectPublicKeyInfo(const uint8_t* buffer, size_t length);
  bool InitFromPrivateKeyInfo(const uint8_t* buffer, size_t length);
  // Initializes the public key object from a private.
  bool InitFromPrivateKey(const EccPrivateKey& private_key);
  // Initializes the public key object from the provided curve and key point
  // |buffer|.
  bool InitFromKeyPoint(EccCurve curve, const uint8_t* buffer, size_t length);
  // Digests the |message| and verifies signature against the provided signature
  // point.
  OEMCryptoResult DigestAndVerify(const uint8_t* message, size_t message_length,
                                  const ECDSA_SIG* sig_point) const;

  // OpenSSL/BoringSSL implementation of an ECC key.
  // As a public key, this will only have key point initialized.
  EC_KEY* key_ = nullptr;
  EccCurve curve_ = kEccCurveUnknown;
};  // class EccPublicKey

class EccPrivateKey {
 public:
  // Creates a new, pseudorandom ECC private key belonging to the
  // curve specified.
  static std::unique_ptr<EccPrivateKey> New(EccCurve curve);

  // Loads a serialized ECC private key.
  // The provided |buffer| must contain a valid ASN.1 DER encoded
  // PrivateKeyInfo containing a valid ECC key description.  Only
  // supported curves by this API are those enumerated by EccCurve.
  //
  //  PrivateKeyInfo := {
  //    version: INTEGER = v1(0) | v2(1),
  //    privateKeyAlgorithm: AlgorithmIdentifier := {
  //      algorithm: OID = id-ecPublicKey,
  //      parameters: ECParameters = {
  //        namedCurve: OID = secp256r1 | secp384r1 | secp521r1
  //      }
  //    },
  //    privateKey: OCTET STRING = ..., -- BER encoding of ECPrivateKey
  //  }
  //
  //  ECPrivateKey := {
  //    version: INTEGER = ecPrivateKeyVer1(1),
  //    privateKey: OCTET STRING = ...,  -- I2OSP of private key point
  //    -- |parameters| are obtained from PrivateKeyInfo
  //    publicKey: BIT STRING OPTIONAL = ... -- SEC1 encoded ECPoint
  //  }
  // Note: If the public key is not included, then it is computed from
  // the private key.
  //
  // References:
  //  RFC 5208 - Description of PrivateKeyInfo
  //  RFC 5480 - Curve OIDs
  //  RFC 5915 - Description of ECPrivateKey in PrivateKeyInfo
  //
  // Failure will occur if the provided |buffer| does not contain a
  // valid PrivateKeyInfo, key is not an ECC key, the specified
  // curve is not supported, or the key is not valid.
  static std::unique_ptr<EccPrivateKey> Load(const uint8_t* buffer,
                                             size_t length);
  static std::unique_ptr<EccPrivateKey> Load(const std::string& buffer);
  static std::unique_ptr<EccPrivateKey> Load(
      const std::vector<uint8_t>& buffer);

  // Creates a new ECC public key of this private key.
  // Equivalent to calling EccPublicKey::New with this private
  // key.
  std::unique_ptr<EccPublicKey> MakePublicKey() const;

  EccCurve curve() const { return curve_; }
  const EC_KEY* GetEcKey() const { return key_; }

  // Checks if the provided |public_key| is the EC public key of this
  // private key.
  bool IsMatchingPublicKey(const EccPublicKey& public_key) const;

  // Serializes the private key into an ASN.1 DER encoded PrivateKeyInfo
  // representation.
  // On success, |buffer_size| is populated with the number of bytes
  // written to |buffer|, and SUCCESS is returned.
  // If the provided |buffer_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |buffer_size| is
  // set to the required buffer size.
  OEMCryptoResult Serialize(uint8_t* buffer, size_t* buffer_size) const;
  // Same as above, except directly returns the serialized key.
  // Returns an empty vector on error.
  std::vector<uint8_t> Serialize() const;

  // Serializes the public component of the private key into an ASN.1
  // DER encoded SubjectPublicKey representation.
  // On success, |buffer_size| is populated with the number of bytes
  // written to |buffer|, and SUCCESS is returned.
  // If the provided |buffer_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |buffer_size| is
  // set to the required buffer size.
  OEMCryptoResult SerializeAsPublicKey(uint8_t* buffer,
                                       size_t* buffer_size) const;
  // Same as above, except directly returns the serialized key.
  // Returns an empty vector on error.
  std::vector<uint8_t> SerializeAsPublicKey() const;

  // Signs the provided |message| and serializes the signature
  // point to |signature| as a ASN.1 DER encoded ECDSA-Sig-Value.
  // This implementation uses ECDSA with the following digest
  // algorithms for the supported curve types.
  //  - SHA-256 / secp256r1
  //  - SHA-384 / secp384r1 (optional support)
  //  - SHA-512 / secp521r1 (optional support)
  // On success, |signature_length| is populated with the number of
  // bytes written to |signature|, and SUCCESS is returned.
  // If the provided |signature_length| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |signature_length|
  // is set to the required signature size.
  OEMCryptoResult GenerateSignature(const uint8_t* message,
                                    size_t message_length, uint8_t* signature,
                                    size_t* signature_length) const;
  // Same as above, except directly returns the serialized signature.
  // Returns an empty vector on error.
  std::vector<uint8_t> GenerateSignature(
      const std::vector<uint8_t>& message) const;
  std::vector<uint8_t> GenerateSignature(const std::string& message) const;
  // Returns an upper bound for the signature size.  May be larger than
  // the actual signature generated by GenerateSignature().
  size_t SignatureSize() const;

  // Derives the OEMCrypto session key used for deriving other keys.
  // The provided public key must be of the same curve.
  // On success, |session_key_size| is populated with the number of
  // bytes written to |session_key|, and OEMCrypto_SUCCESS is returned.
  // If the provided |session_key_size| is too small,
  // OEMCrypto_ERROR_SHORT_BUFFER is returned and |session_key_size|
  // is set to the required buffer size.
  OEMCryptoResult DeriveSessionKey(const EccPublicKey& public_key,
                                   uint8_t* session_key,
                                   size_t* session_key_size) const;
  // Same as above, except directly returns the derived key.
  std::vector<uint8_t> DeriveSessionKey(const EccPublicKey& public_key) const;
  // Returns the byte length of the symmetric key that would be derived
  // by DeriveSymmetricKey().
  size_t SessionKeyLength() const;

  ~EccPrivateKey();

  EccPrivateKey(const EccPrivateKey&) = delete;
  EccPrivateKey(EccPrivateKey&&) = delete;
  const EccPrivateKey& operator=(const EccPrivateKey&) = delete;
  EccPrivateKey& operator=(EccPrivateKey&&) = delete;

 private:
  EccPrivateKey() {}

  // Initializes the public key object using the provided |buffer|.
  // In case of any failure, false is return and the key should be
  // discarded.
  bool InitFromPrivateKeyInfo(const uint8_t* buffer, size_t length);
  // Generates a new key based on the provided curve.
  bool InitFromCurve(EccCurve curve);

  // OpenSSL/BoringSSL implementation of an ECC key.
  // The public point of the key will always be present.
  EC_KEY* key_ = nullptr;
  EccCurve curve_ = kEccCurveUnknown;
};  // class EccPrivateKey
}  // namespace util
}  // namespace wvoec
#endif  // WVOEC_UTIL_ECC_KEY_H_