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

#ifndef WVCDM_CORE_CRYPTO_SESSION_H_
#define WVCDM_CORE_CRYPTO_SESSION_H_

#include <atomic>
#include <map>
#include <memory>
#include <mutex>
#include <string>
#include <vector>

#include "OEMCryptoCENC.h"
#include "crypto_wrapped_key.h"
#include "disallow_copy_and_assign.h"
#include "key_session.h"
#include "metrics_collections.h"
#include "oemcrypto_adapter.h"
#include "rw_lock.h"
#include "timer_metric.h"
#include "wv_cdm_types.h"

namespace wvcdm {
class CryptoKey;
class CryptoSessionFactory;
class OtaKeyboxProvisioner;
class UsageTableHeader;

namespace okp {
class SystemFallbackPolicy;
}  // namespace okp

using CryptoKeyMap = std::map<std::string, CryptoKey*>;

// Crypto session utility functions used by KeySession implementations.
void GenerateMacContext(const std::string& input_context,
                        std::string* deriv_context);
void GenerateEncryptContext(const std::string& input_context,
                            std::string* deriv_context);
size_t GetOffset(std::string message, std::string field);
OEMCrypto_Substring GetSubstring(const std::string& message = "",
                                 const std::string& field = "",
                                 bool set_zero = false);
OEMCryptoCipherMode ToOEMCryptoCipherMode(CdmCipherMode cipher_mode);

class CryptoSession {
 public:
  using HdcpCapability = OEMCrypto_HDCP_Capability;
  enum UsageDurationStatus {
    kUsageDurationsInvalid = 0,
    kUsageDurationPlaybackNotBegun = 1,
    kUsageDurationsValid = 2,
  };

  struct SupportedCertificateTypes {
    bool rsa_2048_bit;
    bool rsa_3072_bit;
    bool rsa_cast;
    bool ecc_secp256r1;
    bool ecc_secp384r1;
    bool ecc_secp521r1;
  };

  // Creates an instance of CryptoSession with the given |crypto_metrics|.
  // |crypto_metrics| is owned by the caller, must NOT be null, and must
  // exist as long as the new CryptoSession exists.
  static CryptoSession* MakeCryptoSession(
      metrics::CryptoMetrics* crypto_metrics);

  virtual ~CryptoSession();

  // This method will try to terminate OEMCrypto if |session_size_| is 0.
  // A platform configured property |delay_oem_crypto_termination| will
  // determine if termination occurs immediately or after a delay.
  // If termination is delayed, a countdown mechanism is employed.
  // Call |TryTerminate| periodically until it no longer returns true.
  // To immediately terminate call |DisableDelayedTermination| before calling
  // |TryTerminate|.
  static bool TryTerminate();

  static void DisableDelayedTermination();

  virtual CdmResponseType GetProvisioningToken(std::string* client_token);
  virtual CdmClientTokenType GetPreProvisionTokenType() {
    return pre_provision_token_type_;
  }

  // The overloaded methods with |requested_level| may be called
  // without a preceding call to Open. The other method must call Open first.
  virtual CdmSecurityLevel GetSecurityLevel();
  virtual CdmSecurityLevel GetSecurityLevel(SecurityLevel requested_level);
  virtual bool GetApiVersion(uint32_t* version);
  virtual bool GetApiVersion(SecurityLevel requested_level, uint32_t* version);
  virtual bool GetApiMinorVersion(SecurityLevel requested_level,
                                  uint32_t* minor_version);

  // This method will return, for devices with a
  // * keybox: the 32 byte device ID from the keybox.
  // * OEM certificate:
  //   - that implements |OEMCrypto_GetDeviceID|: the (1 to 64 byte) device ID.
  //   - that does not implement |OEMCrypto_GetDeviceID|: the OEM public
  //     certificate.
  virtual CdmResponseType GetInternalDeviceUniqueId(std::string* device_id);

  // This method will return, for devices with a
  // * keybox: the 32 byte device ID from the keybox.
  // * OEM certificate:
  //   - that implements |OEMCrypto_GetDeviceID|: the (1 to 64 byte) device ID.
  //   - that does not implement |OEMCrypto_GetDeviceID|: the 32 byte hash
  //     of the OEM public certificate.
  virtual CdmResponseType GetExternalDeviceUniqueId(std::string* device_id);
  virtual bool GetSystemId(uint32_t* system_id);
  virtual CdmResponseType GetProvisioningId(std::string* provisioning_id);
  virtual uint8_t GetSecurityPatchLevel();

  virtual CdmResponseType Open() { return Open(kLevelDefault); }
  virtual CdmResponseType Open(SecurityLevel requested_security_level);
  virtual void Close();

  virtual bool IsOpen() { return open_; }
  virtual CryptoSessionId oec_session_id() { return oec_session_id_; }

  // All request/responses
  virtual const std::string& request_id() { return request_id_; }
  virtual CdmResponseType GenerateNonce(uint32_t* nonce);

  // License request/responses
  virtual CdmResponseType PrepareAndSignLicenseRequest(
      const std::string& message, std::string* core_message,
      std::string* signature);
  // V15 licenses.
  virtual CdmResponseType LoadKeys(const std::string& message,
                                   const std::string& signature,
                                   const std::string& mac_key_iv,
                                   const std::string& mac_key,
                                   const std::vector<CryptoKey>& key_array,
                                   const std::string& provider_session_token,
                                   const std::string& srm_requirement,
                                   CdmLicenseKeyType key_type);
  // V16 licenses.
  virtual CdmResponseType LoadLicense(const std::string& signed_message,
                                      const std::string& core_message,
                                      const std::string& signature,
                                      CdmLicenseKeyType key_type);

  // Renewal request/responses
  virtual CdmResponseType PrepareAndSignRenewalRequest(
      const std::string& message, std::string* core_message,
      std::string* signature);
  // V15 licenses.
  virtual CdmResponseType RefreshKeys(const std::string& message,
                                      const std::string& signature,
                                      const std::vector<CryptoKey>& key_array);
  // V16 licenses.
  virtual CdmResponseType LoadRenewal(const std::string& signed_message,
                                      const std::string& core_message,
                                      const std::string& signature);

  // Entitled content Keys.
  virtual CdmResponseType LoadEntitledContentKeys(
      const std::vector<CryptoKey>& key_array);

  // Provisioning request/responses
  virtual CdmResponseType GenerateDerivedKeys(const std::string& message);
  virtual CdmResponseType GenerateDerivedKeys(const std::string& message,
                                              const std::string& session_key);
  virtual CdmResponseType PrepareAndSignProvisioningRequest(
      const std::string& message, std::string* core_message,
      std::string* signature);
  virtual CdmResponseType LoadProvisioning(const std::string& signed_message,
                                           const std::string& core_message,
                                           const std::string& signature,
                                           std::string* wrapped_private_key);
  virtual CdmResponseType LoadCertificatePrivateKey(
      const CryptoWrappedKey& private_key);

  // Media data path
  virtual CdmResponseType Decrypt(const CdmDecryptionParametersV16& params);

  virtual bool IsAntiRollbackHwPresent();

  // The overloaded methods with |security_level| may be called without a
  // preceding call to Open. The other methods must call Open first.
  virtual CdmResponseType GetHdcpCapabilities(HdcpCapability* current,
                                              HdcpCapability* max);
  virtual CdmResponseType GetHdcpCapabilities(SecurityLevel security_level,
                                              HdcpCapability* current,
                                              HdcpCapability* max);
  virtual bool GetResourceRatingTier(uint32_t* tier);
  virtual bool GetResourceRatingTier(SecurityLevel security_level,
                                     uint32_t* tier);

  virtual bool GetSupportedCertificateTypes(SupportedCertificateTypes* support);
  virtual CdmResponseType GetRandom(size_t data_length, uint8_t* random_data);
  virtual CdmResponseType GetNumberOfOpenSessions(SecurityLevel security_level,
                                                  size_t* count);
  virtual CdmResponseType GetMaxNumberOfSessions(SecurityLevel security_level,
                                                 size_t* max);

  virtual CdmResponseType GetSrmVersion(uint16_t* srm_version);
  virtual bool IsSrmUpdateSupported();
  virtual CdmResponseType LoadSrm(const std::string& srm);

  virtual bool GetBuildInformation(SecurityLevel security_level,
                                   std::string* info);
  virtual bool GetBuildInformation(std::string* info);

  virtual bool GetMaximumUsageTableEntries(SecurityLevel security_level,
                                           size_t* number_of_entries);

  virtual bool GetDecryptHashSupport(SecurityLevel security_level,
                                     uint32_t* hash_support);

  virtual CdmResponseType SetDecryptHash(uint32_t frame_number,
                                         const std::string& hash);

  virtual CdmResponseType GetDecryptHashError(std::string* error_string);

  virtual CdmResponseType GenericEncrypt(const std::string& in_buffer,
                                         const std::string& key_id,
                                         const std::string& iv,
                                         CdmEncryptionAlgorithm algorithm,
                                         std::string* out_buffer);
  virtual CdmResponseType GenericDecrypt(const std::string& in_buffer,
                                         const std::string& key_id,
                                         const std::string& iv,
                                         CdmEncryptionAlgorithm algorithm,
                                         std::string* out_buffer);
  virtual CdmResponseType GenericSign(const std::string& message,
                                      const std::string& key_id,
                                      CdmSigningAlgorithm algorithm,
                                      std::string* signature);
  virtual CdmResponseType GenericVerify(const std::string& message,
                                        const std::string& key_id,
                                        CdmSigningAlgorithm algorithm,
                                        const std::string& signature);

  // Usage table API related methods.
  //   Used to manipulate the CDM managed usage table header & entries,
  //   delegating calls to OEMCrypto.

  // Determines whether the OEMCrypto library supports usage info.
  // As of V16, the only valid type of support is usage table header +
  // usage entries.
  // The first method will use a cached value if present.
  virtual bool HasUsageInfoSupport(bool* has_support);
  virtual bool HasUsageInfoSupport(SecurityLevel security_level,
                                   bool* has_support);

  // Usage report.
  virtual CdmResponseType DeactivateUsageInformation(
      const std::string& provider_session_token);
  virtual CdmResponseType GenerateUsageReport(
      const std::string& provider_session_token, std::string* usage_report,
      UsageDurationStatus* usage_duration_status,
      int64_t* seconds_since_started, int64_t* seconds_since_last_played);

  // Usage table header.
  virtual UsageTableHeader* GetUsageTableHeader() {
    return usage_table_header_;
  }
  // The following crypto methods do not require an open session to
  // complete the operations.
  virtual CdmResponseType CreateUsageTableHeader(
      SecurityLevel requested_security_level,
      CdmUsageTableHeader* usage_table_header);
  virtual CdmResponseType LoadUsageTableHeader(
      SecurityLevel requested_security_level,
      const CdmUsageTableHeader& usage_table_header);
  virtual CdmResponseType ShrinkUsageTableHeader(
      SecurityLevel requested_security_level, uint32_t new_entry_count,
      CdmUsageTableHeader* usage_table_header);

  // Usage entry.
  virtual CdmResponseType CreateUsageEntry(uint32_t* entry_number);
  virtual CdmResponseType LoadUsageEntry(uint32_t entry_number,
                                         const CdmUsageEntry& usage_entry);
  virtual CdmResponseType UpdateUsageEntry(
      CdmUsageTableHeader* usage_table_header, CdmUsageEntry* usage_entry);

  // Adjust usage entries in usage table header.
  virtual CdmResponseType MoveUsageEntry(uint32_t new_entry_number);

  virtual bool GetAnalogOutputCapabilities(bool* can_support_output,
                                           bool* can_disable_output,
                                           bool* can_support_cgms_a);
  virtual metrics::CryptoMetrics* GetCryptoMetrics() { return metrics_; }

  virtual CdmResponseType GetProvisioningMethod(
      SecurityLevel requested_security_level, CdmClientTokenType* token_type);

  // OTA Provisioning

  static bool needs_keybox_provisioning() { return needs_keybox_provisioning_; }

  // This tells the OEMCrypto adapter to ignore the next |count| keyboxes and
  // report that it needs provisioning instead.
  static CdmResponseType SetDebugIgnoreKeyboxCount(uint32_t count);

  // Returns a system-wide singleton instance of SystemFallbackPolicy
  // to be used for communicating OTA keybox provisioning state between
  // apps.  Returns a null pointer if OTA provisioning is not supported,
  // or not required.
  static okp::SystemFallbackPolicy* GetOkpFallbackPolicy();

  // Generates an OTA provisioning request.
  // This should only be called by an instance of OtaKeyboxProvisioner.
  virtual CdmResponseType PrepareOtaProvisioningRequest(bool use_test_key,
                                                        std::string* request);

  // Loads an OTA provisioning response.
  // This should only be called by an instance of OtaKeyboxProvisioner.
  virtual CdmResponseType LoadOtaProvisioning(bool use_test_key,
                                              const std::string& response);

 protected:
  // Creates an instance of CryptoSession with the given |crypto_metrics|.
  // |crypto_metrics| is owned by the caller, must NOT be null, and must
  // exist as long as the new CryptoSession exists.
  explicit CryptoSession(metrics::CryptoMetrics* crypto_metrics);

  int session_count() const { return session_count_; }
  bool initialized() const { return initialized_; }
  void set_initialized(bool initialized) { initialized_ = initialized; }
  // Cache api version and fallback policy. Call this once at initialization.
  void CacheVersion();
  void OverrideNeedKeyboxForTesting(bool needs_keybox_provisioning) {
    needs_keybox_provisioning_ = needs_keybox_provisioning;
  }

 private:
  friend class CryptoSessionForTest;
  friend class CryptoSessionFactory;
#if defined(UNIT_TEST)
  friend class CertificateProvisioningTest;
  friend class WvCdmTestBase;
  friend class CdmOtaKeyboxTest;
#endif

  // The global factory method can be set to generate special crypto sessions
  // just for testing.  These sessions will avoid nonce floods and will ask
  // OEMCrypto to use a test keybox.
  // Ownership of the object is transfered to CryptoSession.
  static void SetCryptoSessionFactory(CryptoSessionFactory* factory) {
    std::unique_lock<std::mutex> auto_lock(factory_mutex_);
    factory_.reset(factory);
  }

  void Init();

  // Will set up the UsageTableHeader for this session.  This may require
  // creating a new UsageTableHeader if the global instance has not
  // been initialized.
  // Note:  This function will lock the global static field lock in write mode.
  bool SetUpUsageTableHeader(SecurityLevel requested_security_level);

  CdmResponseType GetTokenFromKeybox(std::string* token);
  CdmResponseType GetTokenFromOemCert(std::string* token);
  static bool ExtractSystemIdFromOemCert(const std::string& oem_cert,
                                         uint32_t* system_id);
  CdmResponseType GetSystemIdInternal(uint32_t* system_id);
  CdmResponseType GenerateRsaSignature(const std::string& message,
                                       std::string* signature);
  size_t GetMaxSubsampleRegionSize();

  bool SetDestinationBufferType();

  CdmResponseType SelectKey(const std::string& key_id,
                            CdmCipherMode cipher_mode);

  // Retrieves the OEMCrypto usage info support for the specified
  // |requested_security_level|.
  // Caller should acquire the OEMCrypto read lock before calling.
  bool HasUsageInfoSupportInternal(SecurityLevel requested_security_level,
                                   bool* has_support);

  // These methods fall back into each other in the order given, depending on
  // how much data they were given and how much data OEMCrypto can accept in one
  // call.
  OEMCryptoResult DecryptMultipleSamples(
      const std::vector<OEMCrypto_SampleDescription>& samples,
      CdmCipherMode cipher_mode,
      const OEMCrypto_CENCEncryptPatternDesc& pattern);
  OEMCryptoResult DecryptSample(
      const OEMCrypto_SampleDescription& sample, CdmCipherMode cipher_mode,
      const OEMCrypto_CENCEncryptPatternDesc& pattern);
  OEMCryptoResult LegacyDecrypt(
      const OEMCrypto_SampleDescription& sample, CdmCipherMode cipher_mode,
      const OEMCrypto_CENCEncryptPatternDesc& pattern);
  OEMCryptoResult LegacyCopyBufferInChunks(
      const OEMCrypto_SampleDescription& sample, size_t max_chunk_size);
  OEMCryptoResult LegacyDecryptInChunks(
      const OEMCrypto_SampleDescription& sample, CdmCipherMode cipher_mode,
      const OEMCrypto_CENCEncryptPatternDesc& pattern, size_t max_chunk_size);

  // 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 with the exception
  // of the Usage Table Mutex.  The Static Field Mutex can be taken as
  // a reader or as a writer, depending on how you will be accessing
  // the static fields.  The Usage Table Mutex should be taken when
  // reading and writing to the static usage table fields (creating,
  // destroying or taking a pointer of the handles).  The purpose of
  // having a separate mutex for usage table is due to the recursive
  // nature of initializing the global usage table.
  //
  // 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 Header & Entries  |         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());

  static bool IsInitialized();

  // 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 shared_mutex static_field_mutex_;
  static shared_mutex oem_crypto_mutex_;
  std::mutex oem_crypto_session_mutex_;
  // Usage table mutex used only when performing write operations on
  // the static usage table pointers.
  static std::recursive_mutex usage_table_mutex_;

  static bool initialized_;
  static int session_count_;
  static int termination_counter_;
  static bool needs_keybox_provisioning_;

  enum CachedBooleanProperty {
    // Property has not yet been checked/cached.
    kBooleanUnset,
    // Property has been cached as false.
    kBooleanFalse,
    // Property has been cached as true.
    kBooleanTrue
  };

  metrics::CryptoMetrics* metrics_;
  metrics::Timer life_span_;
  uint32_t system_id_;

  bool open_;
  CdmClientTokenType pre_provision_token_type_;
  std::string oem_token_;  // Cached OEMCrypto Public Key
  bool update_usage_table_after_close_session_;
  CryptoSessionId oec_session_id_;
  std::unique_ptr<KeySession> key_session_;

  OEMCryptoBufferType destination_buffer_type_;
  bool is_destination_buffer_type_valid_;
  SecurityLevel requested_security_level_;

  // Open session-cached result of OEMCrypto_SupportsUsageTable().
  CachedBooleanProperty has_usage_info_support_ = kBooleanUnset;
  UsageTableHeader* usage_table_header_ = nullptr;
  // These fields are protected by |usage_table_mutex_| and not
  // |static_field_mutex_|.
  static std::unique_ptr<UsageTableHeader> usage_table_header_l1_;
  static std::unique_ptr<UsageTableHeader> usage_table_header_l3_;

  std::string request_id_;
  static std::atomic<uint64_t> request_id_index_source_;

  uint32_t api_version_;
  size_t max_subsample_region_size_;

  // Stores the most recent error code returned from a call to
  // OEMCrypto_DecryptCENC.  This is used to reduce the total number of
  // error logs for decrypt calls, as there could be a large number of
  // same error code in sequence of each other.  A value of
  // OEMCrypto_SUCCESS indicates that no error have yet occurred.
  OEMCryptoResult last_decrypt_error_ = OEMCrypto_SUCCESS;

  // In order to avoid creating a deadlock if instantiation needs to take any
  // of the CryptoSession static mutexes, |factory_| is protected by its own
  // mutex that is only used in the two funtions that interact with it.
  static std::mutex factory_mutex_;
  static std::unique_ptr<CryptoSessionFactory> factory_;

  // A singleton instance of SystemFallbackPolicy.  Only one will
  // be created for the system if OTA keybox provisioning is both
  // required and supported by L1.
  static std::unique_ptr<okp::SystemFallbackPolicy> okp_fallback_policy_l1_;

  CORE_DISALLOW_COPY_AND_ASSIGN(CryptoSession);
};  // class CryptoSession

class CryptoSessionFactory {
 public:
  virtual ~CryptoSessionFactory() {}
  virtual CryptoSession* MakeCryptoSession(
      metrics::CryptoMetrics* crypto_metrics);

 protected:
  friend class CryptoSession;
  CryptoSessionFactory() {}

 private:
  CORE_DISALLOW_COPY_AND_ASSIGN(CryptoSessionFactory);
};

}  // namespace wvcdm

#endif  // WVCDM_CORE_CRYPTO_SESSION_H_