// Copyright 2018 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine Master
// License Agreement.

#ifndef CRYPTO_SESSION_H
#define CRYPTO_SESSION_H

#include <limits>
#include <memory>
#include <unordered_map>
#include <vector>

#include "OEMCryptoCENC.h"
#include "cas_status.h"
#include "cas_types.h"
#include "crypto_wrapped_key.h"
#include "oemcrypto_interface.h"
#include "rw_lock.h"

namespace wvcas {

struct KeySlot {
  std::vector<uint8_t> key_id;
  std::vector<uint8_t> entitlement_key_id;
  std::vector<uint8_t> wrapped_key;
  std::vector<uint8_t> wrapped_key_iv;
  std::vector<uint8_t> content_iv;
  CryptoMode cipher_mode;
};

struct SubSample {
  uint32_t num_bytes_of_clear;
  uint32_t num_bytes_of_encrypted;
};

typedef OEMCrypto_HDCP_Capability HdcpCapability;

class CryptoLock {
 public:
  CryptoLock() {};
  // These methods should be used to take the various CryptoSession mutexes in
  // preference to taking the mutexes directly.
  //
  // A lock should be taken on the Static Field Mutex before accessing any of
  // CryptoSession's non-atomic static fields. It can be taken as a reader or as
  // a writer, depending on how you will be accessing the static fields.
  //
  // Before calling into OEMCrypto, code must take locks on the OEMCrypto Mutex
  // and/or the OEMCrypto Session Mutex. Which of them should be taken and how
  // depends on the OEMCrypto function being called; consult the OEMCrypto
  // specification's threading guarantees before making any calls. The OEMCrypto
  // specification defines several classes of functions for the purposes of
  // parallelism. The methods below lock the OEMCrypto Mutex and OEMCrypto
  // Session Mutex in the correct order and manner to fulfill the guarantees in
  // the specification.
  //
  // For this function class...    | ...use this locking method
  // ------------------------------+---------------------------
  // Initialization & Termination  |         WithOecWriteLock()
  // Property                      |          WithOecReadLock()
  // Session Initialization        |         WithOecWriteLock()
  // Usage Table                   |         WithOecWriteLock()
  // Session                       |       WithOecSessionLock()
  //
  // Note that accessing |key_session_| often accesses the OEMCrypto session, so
  // WithOecSessionLock() should be used before accessing |key_session_| as
  // well.
  //
  // If a function needs to take a lock on both the Static Field Mutex and some
  // of the OEMCrypto mutexes simultaneously, it must *always* lock the Static
  // Field Mutex before the OEMCrypto mutexes.
  //
  // In general, all locks should only be held for the minimum time necessary
  // (e.g. a lock on the OEMCrypto mutexes should only be held for the duration
  // of a single call into OEMCrypto) unless there is a compelling argument
  // otherwise, such as making two calls into OEMCrypto immediately after each
  // other.
  template <class Func>
  static auto WithStaticFieldWriteLock(const char* tag, Func body)
      -> decltype(body());
  template <class Func>
  static auto WithStaticFieldReadLock(const char* tag, Func body)
      -> decltype(body());
  template <class Func>
  static auto WithOecWriteLock(const char* tag, Func body) -> decltype(body());
  template <class Func>
  static auto WithOecReadLock(const char* tag, Func body) -> decltype(body());
  template <class Func>
  auto WithOecSessionLock(const char* tag, Func body) -> decltype(body());

 private:
  // The locking methods above should be used in preference to taking these
  // mutexes directly. If code takes these manually and needs to take more
  // than one, it must *always* take them in the order they are defined here.
  static wvutil::shared_mutex static_field_mutex_;
  static wvutil::shared_mutex oem_crypto_mutex_;
  std::mutex oem_crypto_session_mutex_;
};

class CryptoInterface {
  CryptoInterface();

 public:
  virtual ~CryptoInterface();

  virtual OEMCryptoResult OEMCrypto_OpenSession(OEMCrypto_SESSION* session);
  virtual OEMCryptoResult OEMCrypto_CloseSession(OEMCrypto_SESSION session);
  virtual OEMCrypto_ProvisioningMethod OEMCrypto_GetProvisioningMethod();
  virtual OEMCryptoResult OEMCrypto_GetKeyData(uint8_t* keyData,
                                               size_t* keyDataLength);
  virtual uint32_t OEMCrypto_SupportedCertificates();
  virtual OEMCryptoResult OEMCrypto_GenerateNonce(OEMCrypto_SESSION session,
                                                  uint32_t* nonce);
  virtual OEMCryptoResult OEMCrypto_GenerateDerivedKeys(
      OEMCrypto_SESSION session, const uint8_t* mac_key_context,
      uint32_t mac_key_context_length, const uint8_t* enc_key_context,
      uint32_t enc_key_context_length);
  virtual OEMCryptoResult OEMCrypto_PrepAndSignLicenseRequest(
      OEMCrypto_SESSION session, uint8_t* message, size_t message_length,
      size_t* core_message_size, uint8_t* signature, size_t* signature_length);
  virtual OEMCryptoResult OEMCrypto_PrepAndSignRenewalRequest(
      OEMCrypto_SESSION session, uint8_t* message, size_t message_length,
      size_t* core_message_size, uint8_t* signature, size_t* signature_length);
  virtual OEMCryptoResult OEMCrypto_PrepAndSignProvisioningRequest(
      OEMCrypto_SESSION session, uint8_t* message, size_t message_length,
      size_t* core_message_size, uint8_t* signature, size_t* signature_length);
  virtual OEMCryptoResult OEMCrypto_LoadProvisioning(
      OEMCrypto_SESSION session, const uint8_t* message, size_t message_length,
      size_t core_message_length, const uint8_t* signature,
      size_t signature_length, uint8_t* wrapped_private_key,
      size_t* wrapped_private_key_length);
  virtual OEMCryptoResult OEMCrypto_GetOEMPublicCertificate(
      OEMCrypto_SESSION session, uint8_t* public_cert,
      size_t* public_cert_length);
  virtual OEMCryptoResult OEMCrypto_LoadDRMPrivateKey(
      OEMCrypto_SESSION session, OEMCrypto_PrivateKeyType key_type,
      const uint8_t* wrapped_rsa_key, size_t wrapped_rsa_key_length);
  virtual OEMCryptoResult OEMCrypto_GenerateRSASignature(
      OEMCrypto_SESSION session, const uint8_t* message, size_t message_length,
      uint8_t* signature, size_t* signature_length,
      RSA_Padding_Scheme padding_scheme);
  virtual OEMCryptoResult OEMCrypto_DeriveKeysFromSessionKey(
      OEMCrypto_SESSION session, const uint8_t* enc_session_key,
      size_t enc_session_key_length, const uint8_t* mac_key_context,
      size_t mac_key_context_length, const uint8_t* enc_key_context,
      size_t enc_key_context_length);
  virtual OEMCryptoResult OEMCrypto_LoadLicense(OEMCrypto_SESSION session,
                                                const uint8_t* message,
                                                size_t message_length,
                                                size_t core_message_length,
                                                const uint8_t* signature,
                                                size_t signature_length);
  virtual OEMCryptoResult OEMCrypto_LoadRenewal(OEMCrypto_SESSION session,
                                                const uint8_t* message,
                                                size_t message_length,
                                                size_t core_message_length,
                                                const uint8_t* signature,
                                                size_t signature_length);
  virtual OEMCryptoResult OEMCrypto_LoadCasECMKeys(
      OEMCrypto_SESSION session, const uint8_t* message, size_t message_length,
      const OEMCrypto_EntitledContentKeyObject* even_key,
      const OEMCrypto_EntitledContentKeyObject* odd_key);
  virtual OEMCryptoResult OEMCrypto_GetHDCPCapability(
      OEMCrypto_HDCP_Capability* current, OEMCrypto_HDCP_Capability* max);
  virtual OEMCryptoResult OEMCrypto_GetDeviceID(uint8_t* deviceID,
                                                size_t* idLength);
  virtual const char* OEMCrypto_SecurityLevel();
  virtual OEMCryptoResult OEMCrypto_CreateEntitledKeySession(
      OEMCrypto_SESSION session, OEMCrypto_SESSION* entitled_key_session_id);
  virtual OEMCryptoResult OEMCrypto_RemoveEntitledKeySession(
      OEMCrypto_SESSION entitled_key_session_id);
  virtual OEMCryptoResult OEMCrypto_ReassociateEntitledKeySession(
      OEMCrypto_SESSION key_sid, OEMCrypto_SESSION oec_sid);
  virtual uint32_t OEMCrypto_APIVersion();
  virtual OEMCryptoResult OEMCrypto_GetOEMKeyToken(
      OEMCrypto_SESSION key_session, uint8_t* key_token,
      size_t* key_token_length);
  virtual OEMCryptoResult OEMCrypto_GetSignatureHashAlgorithm(
      OEMCrypto_SESSION session, OEMCrypto_SignatureHashAlgorithm* algorithm);
  virtual OEMCryptoResult OEMCrypto_GetBootCertificateChain(
      uint8_t* bcc, size_t* bcc_length, uint8_t* additional_signature,
      size_t* additional_signature_length);
  virtual OEMCryptoResult OEMCrypto_GenerateCertificateKeyPair(
      OEMCrypto_SESSION session, uint8_t* public_key, size_t* public_key_length,
      uint8_t* public_key_signature, size_t* public_key_signature_length,
      uint8_t* wrapped_private_key, size_t* wrapped_private_key_length,
      OEMCrypto_PrivateKeyType* key_type);
  virtual OEMCryptoResult OEMCrypto_InstallOemPrivateKey(
      OEMCrypto_SESSION session, OEMCrypto_PrivateKeyType key_type,
      const uint8_t* wrapped_private_key, size_t wrapped_private_key_length);
  virtual uint8_t OEMCrypto_Security_Patch_Level();
  virtual OEMCryptoResult OEMCrypto_BuildInformation(char* buffer,
                                                     size_t* buffer_length);

  // This is the factory method used to enable the oemcrypto interface.
  static OEMCryptoResult create(std::unique_ptr<CryptoInterface>* init) {
    // This is *the* oemcrypto interface used in the normal running. There is
    // only one and there is not need to destroy it.
    static OEMCryptoInterface oemcrypto_interface;
    return create_internal(&oemcrypto_interface, init);
  }

  // This initializer factory method is templated to allow tests to pass in
  // a mocked OEMCryptoInterface. The caller retains ownership of
  // |oemcrypto_interface|.
  template <typename CI>
  static OEMCryptoResult create(std::unique_ptr<CryptoInterface>* init,
                                CI* oemcrypto_interface) {
    return create_internal(oemcrypto_interface, init);
  }

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

 private:
  static OEMCryptoResult create_internal(
      OEMCryptoInterface* oemcrypto_interface,
      std::unique_ptr<CryptoInterface>* init);

  static bool initialized_;
  static int session_count_;
  static std::unique_ptr<CryptoLock> lock_;
  OEMCryptoInterface* oemcrypto_interface_;
};

class SupportedCertificates {
 public:
  explicit SupportedCertificates(uint32_t supported) : supported_(supported) {}
  bool rsa_2048bit() { return OEMCrypto_Supports_RSA_2048bit & supported_; }
  bool rsa_3072bit() { return OEMCrypto_Supports_RSA_3072bit & supported_; }
  bool rsa_CASTbit() { return OEMCrypto_Supports_RSA_CAST & supported_; }

 private:
  uint32_t supported_;
};

// CryptoSession implements the core methods need to interface with OEMCrypto.
class CryptoSession {
 public:
  explicit CryptoSession();
  virtual ~CryptoSession();
  virtual CasStatus initialize();
  virtual CasStatus reset();
  virtual CasStatus close();
  virtual CasProvisioningMethod provisioning_method();
  virtual CasStatus GetKeyData(uint8_t* keyData, size_t* keyDataLength);
  virtual SupportedCertificates supported_certificates();
  virtual CasStatus GenerateNonce(uint32_t* nonce);
  virtual CasStatus GenerateDerivedKeys(const uint8_t* mac_key_context,
                                        uint32_t mac_key_context_length,
                                        const uint8_t* enc_key_context,
                                        uint32_t enc_key_context_length);
  virtual CasStatus PrepareAndSignLicenseRequest(
      const std::string& message, std::string* core_message,
      std::string* signature, bool& should_specify_algorithm,
      OEMCrypto_SignatureHashAlgorithm& algorithm);
  virtual CasStatus PrepareAndSignRenewalRequest(const std::string& message,
                                                 std::string* core_message,
                                                 std::string* signature);
  virtual CasStatus PrepareAndSignProvisioningRequest(
      const std::string& message, std::string* core_message,
      std::string* signature, bool& should_specify_algorithm,
      OEMCrypto_SignatureHashAlgorithm& algorithm);
  virtual CasStatus LoadProvisioning(const std::string& signed_message,
                                     const std::string& core_message,
                                     const std::string& signature,
                                     std::string* wrapped_private_key);
  virtual CasStatus GetOEMPublicCertificate(uint8_t* public_cert,
                                            size_t* public_cert_length);
  virtual CasStatus GenerateRSASignature(const uint8_t* message,
                                         size_t message_length,
                                         uint8_t* signature,
                                         size_t* signature_length,
                                         RSA_Padding_Scheme padding_scheme);
  virtual CasStatus DeriveKeysFromSessionKey(const uint8_t* enc_session_key,
                                             size_t enc_session_key_length,
                                             const uint8_t* mac_key_context,
                                             size_t mac_key_context_length,
                                             const uint8_t* enc_key_context,
                                             size_t enc_key_context_length);
  virtual CasStatus LoadLicense(const std::string& signed_message,
                                const std::string& core_message,
                                const std::string& signature);
  virtual CasStatus LoadRenewal(const std::string& signed_message,
                                const std::string& core_message,
                                const std::string& signature);
  // LoadCasECMKeys loads the ecm keys into the crypto library making them
  // available for use.
  // |odd_key| - if not null, contains control word data.
  // |even_key| - if not null, contains control word data.
  virtual CasStatus LoadCasECMKeys(OEMCrypto_SESSION session,
                                   const KeySlot* even_key,
                                   const KeySlot* odd_key);
  virtual bool GetHdcpCapabilities(HdcpCapability* current,
                                   HdcpCapability* max);
  virtual CasStatus GetDeviceID(std::string* buffer);
  virtual const char* SecurityLevel();
  virtual CasStatus CreateEntitledKeySession(
      OEMCrypto_SESSION* entitled_key_session_id);
  virtual CasStatus RemoveEntitledKeySession(
      OEMCrypto_SESSION entitled_key_session_id);
  virtual CasStatus ReassociateEntitledKeySession(
      OEMCrypto_SESSION entitled_key_session_id);
  virtual CasStatus APIVersion(uint32_t* api_version);
  virtual CasStatus GetOEMKeyToken(OEMCrypto_SESSION entitled_key_session_id,
                                   std::vector<uint8_t>& token);
  virtual CasStatus GetBootCertificateChain(std::string* bcc,
                                            std::string* additional_signature);
  virtual CasStatus GenerateCertificateKeyPair(
      std::string* public_key, std::string* public_key_signature,
      std::string* wrapped_private_key, CryptoWrappedKey::Type* key_type);
  virtual CasStatus LoadOemCertificatePrivateKey(
      const CryptoWrappedKey& private_key);
  virtual CasStatus LoadCertificatePrivateKey(
      const CryptoWrappedKey& private_key);
  virtual uint8_t GetSecurityPatchLevel();
  virtual bool GetBuildInformation(std::string* info);

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

 private:
  virtual OEMCryptoResult getCryptoInterface(
      std::unique_ptr<CryptoInterface>* interface);

  // TODO(jfore, widevine-eng): Merge CryptoInterface into CryptoSession and
  // drop this shared pointer.
  std::unique_ptr<CryptoInterface> crypto_interface_;
  OEMCrypto_SESSION session_;
};

}  // namespace wvcas

#endif  // CRYPTO_SESSION_H