// 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 <inttypes.h>
#include <stddef.h>

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

#include <openssl/ec.h>

#include "OEMCryptoCENCCommon.h"
#include "wv_class_utils.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:
  ~EccPublicKey();
  WVCDM_DISALLOW_COPY_AND_MOVE(EccPublicKey);

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

  // Loads a ASN.1 DER 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 = ... -- SEC 1 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 a SEC 1 serialized EC public key.
  //
  // The provided |buffer| must contain a valid SEC 1 encoded EC point
  // belonging to the provided |curve|.
  //
  // SEC 1 section 2.3.3 specifies two supported formats, compressed or
  // uncompressed.
  //
  // Case uncompressed:
  //  buffer: 0x04 || X || Y
  //
  //  Where X and Y are fixed-width byte encodings of the public keys
  //  x and y component.
  //
  // Case compressed:
  //  buffer: (0x02 or 0x03) || X
  //
  //  Where X is a fixed-width byte encoding of the public keys x
  //  component; and the y component is derived using the curve
  //  equation, and the sign of y is positive if lead byte is 0x02,
  //  or negative if lead byte is 0x03.
  //
  // Note: The EC point encoding in SEC 1 is derived from the X9.62
  // format; however, the "evenness" version of the X9.62 compressed
  // point is NOT supported and is generally not recommended.
  static std::unique_ptr<EccPublicKey> LoadKeyPoint(EccCurve curve,
                                                    const uint8_t* buffer,
                                                    size_t length);
  static std::unique_ptr<EccPublicKey> LoadKeyPoint(EccCurve curve,
                                                    const std::string& buffer);
  static std::unique_ptr<EccPublicKey> LoadKeyPoint(
      EccCurve curve, const std::vector<uint8_t>& buffer);

  // 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;
  // Serializes the public key into a SEC 1 encoded public key point.
  // To restore a key from the returned key point, the caller must
  // keep track of the specific curve of the key.
  //
  // If |compressed| is false (default), the key point will be
  // serialized in its uncompressed form; otherwise, the key will be
  // serialized in its compressed form.
  //
  // Directly returns the serialized public key.
  // Returns an empty vector on error.
  std::vector<uint8_t> SerializeAsSec1KeyPoint(bool compressed = false) 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 the unsigned integers r and
  // s as two equal length octet strings using big-endian encoding.
  //
  // The |message| is digested using the same ECDSA algorithm as
  // VerifySignature().
  //
  // Returns:
  //    OEMCrypto_SUCCESS if signature is valid
  //    OEMCrypto_ERROR_SIGNATURE_FAILURE if the |signature| is invalid
  //    Any other result indicates an unexpected error
  OEMCryptoResult VerifyRawSignature(const uint8_t* message,
                                     size_t message_length,
                                     const uint8_t* signature,
                                     size_t signature_length) const;
  OEMCryptoResult VerifyRawSignature(
      const std::vector<uint8_t>& message,
      const std::vector<uint8_t>& signature) const;

 private:
  EccPublicKey() = default;

  // 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
  // SEC 1 encoded EC key point |buffer|.
  bool InitFromSec1KeyPoint(EccCurve curve, const uint8_t* buffer,
                            size_t length);
  // Digests the |message| and verifies signature against the provided
  // ECDSA signature point |sig_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:
  ~EccPrivateKey();
  WVCDM_DISALLOW_COPY_AND_MOVE(EccPrivateKey);

  // 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 = ... -- SEC 1 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;

  // Serializes the public component of the private key into an SEC 1
  // public key point.
  // See EccPublicKey::SerializeAsSec1KeyPoint() for details.
  std::vector<uint8_t> SerializeAsPublicSec1KeyPoint(
      bool compressed = false) 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;

  // Special test method used to generate a raw ECDSA signature.
  // A raw ECDSA signature is a concatenation of a same-width-big-endian
  // encoding of the ECDSA signature point components r and s.
  std::vector<uint8_t> GenerateRawSignature(
      const std::vector<uint8_t>& message) 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;

 private:
  EccPrivateKey() = default;

  // 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_