// 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_CDM_USAGE_TABLE_H_
#define WVCDM_CORE_CDM_USAGE_TABLE_H_

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

#include "clock.h"
#include "crypto_session.h"
#include "device_files.h"
#include "file_store.h"
#include "metrics_collections.h"
#include "wv_cdm_types.h"
#include "wv_class_utils.h"

#if defined(UNIT_TEST)
#  include <gtest/gtest_prod.h>
#endif

namespace wvcdm {
// Offline licenses/secure stops may be securely tracked using usage
// tables (OEMCrypto v9-12) or usage table headers+usage entries
// (OEMCrypto v13+). This class assists with the latter, synchronizing
// access to usage table header and associated data-structures and controlling
// when they are read in or written out to non-secure persistent storage.
//
// Each OEMCrypto (for each security level) will maintain its own usage table
// header. Each license will have an associated usage entry that is also
// stored in persistent memory and is noted in the usage table header.
// Usage entry information will be verified when licenses are loaded.
//
// OEMCrypto for each security level have their own usage table
// headers. They are loaded on initialization and written out periodically.
// The lifecycle of this class is tied to when OEMCrypto is
// initialized/terminated.
//
// Sessions and licenses are however handled by CdmSession and so most
// calls to manipulate the usage table header related to usage entries
// are by CdmSession.
//
// Upgrades from a fixed size usage table (supported by previous
// versions of the OEMCrypto API v9-12) are handled by this class.
// |entry| and |entry_index|s need to be saved in the license
// and usage info records by the caller.
class CdmUsageTable {
 public:
  CdmUsageTable();
  WVCDM_DISALLOW_COPY_AND_MOVE(CdmUsageTable);
  virtual ~CdmUsageTable() {}

  // |crypto_session| is used to create or load a usage master table
  // and not cached beyound this call.
  // First attempts to restore the usage table from the device files.
  // If restoring fails, then a new usage table is created.
  // Note: No OEMCrypto session for the same security level should be
  // opened before calling.
  // Threading: Method requires care of caller for exclusive access
  // (ie. test setup or CryptoSession).
  bool Init(CdmSecurityLevel security_level, CryptoSession* crypto_session);

  // Adds a new entry to the OEMCrypto usage table header, and associates
  // the entry with the provided |crypto_session|.  The index of the new
  // usage entry will be returned by |entry_index|.
  //
  // Type of entry depends on the value of |persistent_license|:
  //    false - usage info / secure stop record
  //    true  - offline license
  //
  // Threading: Method takes exclusive use of |lock_|.
  virtual CdmResponseType AddEntry(CryptoSession* crypto_session,
                                   bool persistent_license,
                                   const CdmKeySetId& key_set_id,
                                   const std::string& usage_info_filename,
                                   const CdmKeyResponse& license_message,
                                   UsageEntryIndex* entry_index);
  // Threading: Method takes exclusive use of |lock_|.
  virtual CdmResponseType LoadEntry(CryptoSession* crypto_session,
                                    const UsageEntry& entry,
                                    UsageEntryIndex entry_index);
  // Threading: Method takes exclusive use of |lock_|.
  virtual CdmResponseType UpdateEntry(UsageEntryIndex entry_index,
                                      CryptoSession* crypto_session,
                                      UsageEntry* entry);

  // The licenses or usage info records specified by |entry_index|
  // should not be in use by any open CryptoSession objects when calls
  // to InvalidateEntry and MoveEntry are made.
  // If |defrag_table| is true, the table will be defragmented after
  // the entry has been invalidated.
  //
  // Threading: Method takes exclusive use of |lock_|.
  virtual CdmResponseType InvalidateEntry(UsageEntryIndex entry_index,
                                          bool defrag_table,
                                          DeviceFiles* device_files,
                                          metrics::CryptoMetrics* metrics);

  // Test only method. This method emulates the behavior of InvalidateEntry
  // without actually invoking OEMCrypto (through CryptoSession)
  // or storage (through DeviceFiles). It modifies internal data structures
  // when InvalidateEntry is mocked. This allows one to test methods that are
  // dependent on InvalidateEntry without having to set expectations
  // for the objects that InvalidateEntry depends on.
  //
  // Threading: Method requires care of caller for exclusive access.
  void InvalidateEntryForTest(UsageEntryIndex entry_index);

  // == Table information methods ==
  // Threading: None of the following are thread safe.  Intended for
  // testing or internal use.

  size_t size() { return entry_info_list_.size(); }

  size_t potential_table_capacity() const { return potential_table_capacity_; }

  bool HasUnlimitedTableCapacity() const {
    return potential_table_capacity_ == 0;
  }

  // Returns the number of entries currently tracked by the CDM that
  // are related to usage info (streaming licenses).
  size_t UsageInfoCount() const;
  // Returns the number of entries currently tracked by the CDM that
  // are related to offline licenses.
  size_t OfflineEntryCount() const;

  const std::vector<CdmUsageEntryInfo>& entry_info_list() const {
    return entry_info_list_;
  }

  // Set the reference clock used for the method GetCurrentTime().
  void SetClock(wvutil::Clock* clock) {
    if (clock != nullptr)
      clock_ref_ = clock;
    else
      clock_ref_ = &clock_;
  }

  static bool DetermineLicenseToRemoveForTesting(
      const std::vector<CdmUsageEntryInfo>& entry_info_list,
      int64_t current_time, size_t unexpired_threshold,
      UsageEntryIndex* entry_to_remove) {
    return DetermineLicenseToRemove(entry_info_list, current_time,
                                    unexpired_threshold, entry_to_remove);
  }

 private:
  // == Initialization methods ==
  // These should only be called during table initialization.
  // All parameters are assumed valid.

  // Creates a new, empty usage table.  Any existing usage table files
  // will be deleted.
  // Threading: Method takes exclusive use of |lock_|
  // when required.
  bool CreateNewTable(CryptoSession* crypto_session);
  // Attempts to restore the usage table from persistent storage, and
  // loads the usage table header into OEMCrypto.
  // Note: No other OEMCrypto session should be opened before calling.
  // Threading: Method takes exclusive use of |lock_|
  // when required.
  bool RestoreTable(CryptoSession* crypto_session);

  // Performs a check that there are no open OEMCrypto sessions for
  // the current security level of the usage table.
  // Threading: No special threading requirements.
  bool OpenSessionCheck(CryptoSession* crypto_session);
  // Performs a check that the OEMCrypto table can support at least
  // one more entry if the table is at or near the reported capacity.
  // If this check fails, a new usage table SHOULD be created.
  // Threading: Method requires caller to take exclusive use of
  // |lock_|.
  bool CapacityCheck(CryptoSession* crypto_session);

  // Attempts to determine the capacity of the OEMCrypto usage table.
  // Sets the result to |potential_table_capacity_|.
  // Threading: Method requires caller to take exclusive use of
  // |lock_|.
  bool DetermineTableCapacity(CryptoSession* crypto_session);

  // == Table operation methods ==
  // Threading: All of the following methods require caller to take
  // exclusive use of |lock_|.

  // Creates a new entry for the provided crypto session.  If the
  // entry is created successfully in OEMCrypto, then a new entry
  // info is added to the table's vector of entry info.
  CdmResponseType CreateEntry(CryptoSession* crypto_session,
                              UsageEntryIndex* entry_index);

  // Attempts to relocate a newly created usage entry associated with
  // the provided |crypto_session| to the lowest unoccupied position in
  // the table.
  // |entry_index| is treated as both an input and output.
  // Returns NO_ERROR so long as no internal operation fails,
  // regardless of whether the entry was moved or not.
  CdmResponseType RelocateNewEntry(CryptoSession* crypto_session,
                                   UsageEntryIndex* entry_index);

  // Checks if the specified |entry_index| is known to be
  // unoccupied (released).
  bool IsEntryUnoccupied(UsageEntryIndex entry_index) const;

  // SetOfflineEntryInfo() and SetUsageInfoEntryInfo() populate the
  // entry meta-data with the required information based on the type
  // of entry.
  void SetOfflineEntryInfo(UsageEntryIndex entry_index,
                           const std::string& key_set_id,
                           const CdmKeyResponse& license_message);
  void SetUsageInfoEntryInfo(UsageEntryIndex entry_index,
                             const std::string& key_set_id,
                             const std::string& usage_info_file_name);

  // Shrinks the table, removing all trailing unoccupied entries.
  // |entry_info_list_| will be resized appropriately.
  // Caller must store the table after a successful call.
  CdmResponseType RefitTable(CryptoSession* crypto_session);

  virtual CdmResponseType InvalidateEntryInternal(
      UsageEntryIndex entry_index, bool defrag_table, DeviceFiles* device_files,
      metrics::CryptoMetrics* metrics);

  CdmResponseType MoveEntry(UsageEntryIndex from, const UsageEntry& from_entry,
                            UsageEntryIndex to, DeviceFiles* device_files,
                            metrics::CryptoMetrics* metrics);

  CdmResponseType GetEntry(UsageEntryIndex entry_index,
                           DeviceFiles* device_files, UsageEntry* entry);
  CdmResponseType StoreEntry(UsageEntryIndex entry_index,
                             DeviceFiles* device_files,
                             const UsageEntry& entry);

  // Stores the usage table and it's info.  This will increment
  // |store_table_counter_| if successful.
  bool StoreTable();

  CdmResponseType Shrink(metrics::CryptoMetrics* metrics,
                         uint32_t number_of_entries_to_delete);

  // Must lock table before calling.
  CdmResponseType DefragTable(DeviceFiles* device_files,
                              metrics::CryptoMetrics* metrics);

  // This will use the LRU algorithm to decide which entry is to be
  // evicted.
  CdmResponseType ReleaseOldestEntry(metrics::CryptoMetrics* metrics);

  // Performs and LRU upgrade on all loaded CdmUsageEntryInfo from a
  // device file that had not yet been upgraded to use the LRU data.
  virtual bool LruUpgradeAllUsageEntries();

  virtual bool GetRemovalCandidate(uint32_t* entry_to_remove);

  int64_t GetCurrentTime() { return clock_ref_->GetCurrentTime(); }

  // Sets LRU related metrics based on the provided |staleness| (in
  // seconds) and |storage_type| of the entry removed.
  void RecordLruEventMetrics(metrics::CryptoMetrics* metrics,
                             uint64_t staleness,
                             CdmUsageEntryStorageType storage_type);

  // Uses an LRU-base algorithm to determine which license should be
  // removed.  This is intended to be used if the usage table is full
  // and a new entry needs to be added.
  //
  // Algorithm overview:
  //   Given the set of all usage entry infos, the entries which are
  //   of unknown storage type or are the most stale will be returned,
  //   with some exceptions for offline licenses.
  //     1) Expired licenses will always be considered.  Expiration is
  //        determined using the usage entry info's
  //        |offline_license_expiry| compared to the provided
  //        |current_time|.
  //     2) Unexpired offline licenses will only be considered for
  //        removal if the number of unexpired offline licenses exceeds
  //        |unexpired_threshold|.
  //
  //   Unknown storage types will be considered above all other entry
  //   types.
  //
  // Parameters:
  //    [in] entry_info_list: The complete list of known usage
  //            entries.
  //    [in] current_time: The current time to compare expiration times
  //            against.
  //    [in] unexpired_threshold: The maximum number of unexpired
  //            offline licenses that are present, before offline
  //            licenses would be considered for removal.
  //    [out] entry_to_remove: Usage entry index of the entry selected
  //            to be removed.  Assume to be unaffected if the
  //            function returns |false|.
  //
  // Returns:
  //   |true| if an entry has been determined to be removed.
  //   Otherwise returns |false|.
  static bool DetermineLicenseToRemove(
      const std::vector<CdmUsageEntryInfo>& entry_info_list,
      int64_t current_time, size_t unexpired_threshold,
      UsageEntryIndex* entry_to_remove);

  // This handle and file system is only to be used when accessing
  // header.  Usage entries should use the file system provided
  // by CdmSession.
  std::unique_ptr<DeviceFiles> device_files_;
  std::unique_ptr<wvutil::FileSystem> file_system_;
  CdmSecurityLevel security_level_ = kSecurityLevelUninitialized;
  RequestedSecurityLevel requested_security_level_ = kLevelDefault;

  UsageTableHeader header_;
  std::vector<CdmUsageEntryInfo> entry_info_list_;

  // Table is sync with persistent storage and can be used by the CDM
  // to interact with OEMCrypto.
  bool is_initialized_ = false;

  // Synchonizes access to the Usage Table Header and bookkeeping
  // data-structures
  mutable std::mutex lock_;

  metrics::CryptoMetrics alternate_crypto_metrics_;

  // |clock_| represents the system's "wall clock".  For the clock's purpose
  // we do not need a more secure clock.
  wvutil::Clock clock_;
  // |clock_ref_| is a pointer to the clock which is to be used for
  // obtaining the current time.  By default, this points to the internal
  // |clock_| variable, however, it can be overridden for testing purpose.
  wvutil::Clock* clock_ref_;

  // The maximum number of entries that the underlying OEMCrypto
  // implementation can support.  Some implementations might not
  // support reporting the table capacity, if so, then this value is
  // assumed to be |kMinimumUsageTableEntriesSupported|.
  size_t potential_table_capacity_ = 0u;

  // Counts the number of successful calls to |StoreTable()|.  Used
  // to reduce the number of calls to device files for certain
  // table operations.
  uint32_t store_table_counter_ = 0u;

#if defined(UNIT_TEST)
  // Test related declarations
  friend class CdmUsageTableTest;

  FRIEND_TEST(CdmUsageTableTest, Shrink_NoneOfTable);
  FRIEND_TEST(CdmUsageTableTest, Shrink_PartOfTable);
  FRIEND_TEST(CdmUsageTableTest, Shrink_AllOfTable);
  FRIEND_TEST(CdmUsageTableTest, Shrink_MoreThanTable);
#endif

  // These setters are for testing only. Takes ownership of the pointers.
  void SetDeviceFiles(DeviceFiles* device_files) {
    device_files_.reset(device_files);
  }
  void SetCryptoSession(CryptoSession* crypto_session) {
    test_crypto_session_.reset(crypto_session);
  }

  // Test related data members
  std::unique_ptr<CryptoSession> test_crypto_session_;
};  // class CdmUsageTable
}  // namespace wvcdm
#endif  // WVCDM_CORE_CDM_USAGE_TABLE_H_