#ifndef CDM_OEC_SESSION_UTIL_H_
#define CDM_OEC_SESSION_UTIL_H_

// Copyright 2018 Google LLC. All Rights Reserved. This file and proprietary
// source code may only be used and distributed under the Widevine
// License Agreement.
//
// OEMCrypto unit tests
//
#include <gtest/gtest.h>
#include <time.h>
#include <unordered_map>

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

#include "core_message_deserialize.h"
#include "core_message_features.h"
#include "core_message_serialize.h"
#include "odk.h"
#include "oec_device_features.h"
#include "oec_key_deriver.h"
#include "oemcrypto_ecc_key.h"
#include "oemcrypto_fuzz_structs.h"
#include "oemcrypto_rsa_key.h"
#include "oemcrypto_types.h"
#include "pst_report.h"

// GTest requires PrintTo to be in the same namespace as the thing it prints,
// which is std::vector in this case.
namespace std {

void PrintTo(const vector<uint8_t>& value, ostream* os);

}  // namespace std

namespace wvoec {

using namespace std;

// OEMCrypto Fuzzing: Set max signture length to 1mb.
const size_t MB = 1024 * 1024;

// Make sure this is larger than kMaxKeysPerSession.
constexpr size_t kMaxNumKeys = 30;

#if defined(TEST_SPEED_MULTIPLIER)  // Can slow test time limits when
                                    // debugging is slowing everything.
constexpr int kSpeedMultiplier = TEST_SPEED_MULTIPLIER;
#else
constexpr int kSpeedMultiplier = 1;
#endif
constexpr int kShortSleep = 1 * kSpeedMultiplier;
constexpr int kLongSleep = 2 * kSpeedMultiplier;
constexpr uint32_t kDuration = 2 * kSpeedMultiplier;
constexpr uint32_t kLongDuration = 5 * kSpeedMultiplier;
constexpr int32_t kTimeTolerance = 3 * kSpeedMultiplier;
constexpr int64_t kUsageTableTimeTolerance = 10 * kSpeedMultiplier;

// Note: The API does not specify a maximum key id length.  We specify a
// maximum just for these tests, so that we have a fixed message size.
constexpr size_t kTestKeyIdMaxLength = 16;

// Most content will use a key id that is 16 bytes long.
constexpr int kDefaultKeyIdLength = 16;
constexpr size_t kMaxPSTLength = 255;  // In specification.
constexpr size_t kMaxCoreMessage = 200 * kMaxNumKeys + 200;  // Rough estimate.

typedef struct {
  uint8_t key_id[kTestKeyIdMaxLength];
  size_t key_id_length;
  uint8_t key_data[MAC_KEY_SIZE];
  size_t key_data_length;
  uint8_t key_iv[KEY_IV_SIZE];
  uint8_t control_iv[KEY_IV_SIZE];
  KeyControlBlock control;
  // Note: cipher_mode may not be part of a real signed message. For these
  // tests, it is convenient to keep it in this structure anyway.
  OEMCryptoCipherMode cipher_mode;
} MessageKeyData;

// This structure will be signed to simulate a message from the server.
struct MessageData {
  MessageKeyData keys[kMaxNumKeys];
  uint8_t mac_key_iv[KEY_IV_SIZE];
  uint8_t padding[KEY_IV_SIZE];
  uint8_t mac_keys[2 * MAC_KEY_SIZE];
  uint8_t pst[kMaxPSTLength];
  SRM_Restriction_Data srm_restriction_data;
};

struct Test_PST_Report {
  Test_PST_Report(const std::string& pst_in,
                  OEMCrypto_Usage_Entry_Status status_in);

  OEMCrypto_Usage_Entry_Status status;
  int64_t seconds_since_license_received;
  int64_t seconds_since_first_decrypt;
  int64_t seconds_since_last_decrypt;
  std::string pst;
  int64_t time_created;
};

struct EntitledContentKeyData {
  uint8_t entitlement_key_id[kTestKeyIdMaxLength];
  size_t entitlement_key_id_length;
  uint8_t content_key_id[kTestKeyIdMaxLength];
  size_t content_key_id_length;
  uint8_t content_key_data_iv[KEY_IV_SIZE];
  uint8_t content_key_data[KEY_SIZE];
  uint8_t encrypted_content_key_data[KEY_SIZE];
  uint8_t content_iv[KEY_IV_SIZE];
  size_t content_iv_length;
  size_t key_index;  // Index into the license's key array. Only for testing.
};

// returns 1 on success, -1 if not supported, or 0 if other failure.
int GetRandBytes(unsigned char* buf, size_t num);

void GenerateSimpleSampleDescription(const std::vector<uint8_t>& in,
                                     std::vector<uint8_t>& out,
                                     OEMCrypto_SampleDescription* sample,
                                     OEMCrypto_SubSampleDescription* subsample);

// Encrypt a block of data using CTR mode.
void EncryptCTR(const vector<uint8_t>& in_buffer, const uint8_t* key,
                const uint8_t* starting_iv, vector<uint8_t>* out_buffer);

// Increment counter for AES-CTR.  The CENC spec specifies we increment only
// the low 64 bits of the IV counter, and leave the high 64 bits alone.  This
// is different from the OpenSSL implementation, so we implement the CTR loop
// ourselves.
void ctr128_inc64(int64_t increaseBy, uint8_t* iv);

// Some compilers don't like the macro htonl within an ASSERT_EQ.
uint32_t htonl_fnc(uint32_t x);

// Prints error string from BoringSSL
void dump_boringssl_error();

class Session;
// The prototype of the OEMCrypto function to prepare and sign a request.
typedef OEMCryptoResult (*PrepAndSignRequest_t)(
    OEMCrypto_SESSION session, uint8_t* message, size_t message_length,
    size_t* core_message_length, uint8_t* signature, size_t* signature_length);

// A RoundTrip helps generate and verify a request message, helps generate the
// corresponding response, and then helps verify loading the response.
template <class CoreRequest, PrepAndSignRequest_t PrepAndSignRequest,
          class CoreResponse, class ResponseData>
class RoundTrip {
 public:
  RoundTrip() = delete;
  RoundTrip(Session* session)
      : session_(session),
        core_request_(),
        core_response_(),
        response_data_(),
        encrypted_response_data_(),
        required_message_size_(0),
        required_core_message_size_(0),
        required_request_signature_size_(0),
        encrypted_response_length_(0),
        response_signature_length_(0) {}
  virtual ~RoundTrip() = default;

  // Have OEMCrypto sign a request message and then verify the signature and the
  // core message.
  virtual void SignAndVerifyRequest() {
    // Boolean true generates core request and verifies the request.
    // Custom message sizes are 0 by default, so the behavior of following
    // functions will be sign and verify request without any custom buffers
    // sizes.
    ASSERT_EQ(SignAndCreateRequestWithCustomBufferLengths(true),
              OEMCrypto_SUCCESS);
  }
  // Have OEMCrypto sign and call create request APIs. Buffer parameters in API
  // can be set to custom values to test with varying lengths of buffers.
  virtual OEMCryptoResult SignAndCreateRequestWithCustomBufferLengths(
      bool verify_request = false);
  // Used for OEMCrypto Fuzzing: Function to convert fuzzer data to valid
  // License/Provisioning/Renwal request data that can be serialized.
  virtual void InjectFuzzedRequestData(uint8_t* data, size_t size);
  // Create a default |response_data| and |core_response|.
  virtual void CreateDefaultResponse() = 0;
  // Copy fields from |response_data| to |padded_response_data|, encrypting
  // those that should be encrypted. Serialize the core message. Then sign the
  // response.
  virtual void EncryptAndSignResponse() = 0;
  // Attempt to load the response and return the error.  Short buffer errors are
  // handled by LoadResponse, not the caller.  All other errors should be
  // handled by the caller.
  virtual OEMCryptoResult LoadResponse() { return LoadResponse(session_); }
  // As with LoadResponse, but load into a different session.
  virtual OEMCryptoResult LoadResponse(Session* session) = 0;

  // Accessors are all read/write because tests modify default values.
  Session* session() { return session_; }
  void set_session(Session* session) { session_ = session; }
  CoreRequest& core_request() { return core_request_; }
  CoreResponse& core_response() { return core_response_; }
  ResponseData& response_data() { return response_data_; }
  std::vector<uint8_t>& encrypted_response_buffer() {
    return encrypted_response_;
  }

  // Set the size of the buffer used the encrypted license.
  void set_message_size(size_t size) { required_message_size_ = size; }
  // Set core message size to test OEMCrypto request APIs for varying core
  // message lengths.
  void set_core_message_size(size_t size) {
    required_core_message_size_ = size;
  }
  // Set signature size to test OEMCrypto request APIs for varying signature
  // lengths.
  void set_request_signature_size(size_t size) {
    required_request_signature_size_ = size;
  }
  std::vector<uint8_t>& response_signature() { return response_signature_; }
  const std::string& serialized_core_message() const {
    return serialized_core_message_;
  }

 protected:
  // Returns true if a nonce should be generated before signing the request.
  virtual bool RequestHasNonce() = 0;

  // ----------------------------------------------------------------------
  // Specialized functionality for each message type.

  // Verify the signature of the request.
  virtual void VerifyRequestSignature(
      const vector<uint8_t>& data, const vector<uint8_t>& generated_signature,
      size_t core_message_length) = 0;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) = 0;
  // Find the given pointer in the response_data_.
  virtual OEMCrypto_Substring FindSubstring(const void* pointer, size_t length);

  // Set EncryptAndSignResponse output lengths for later verification.
  void SetEncryptAndSignResponseLengths();
  // Verify EncryptAndSignResponse output lengths are unchanged.
  void VerifyEncryptAndSignResponseLengths() const;

  // ----------------------------------------------------------------------
  // Member variables.
  Session* session_;
  CoreRequest core_request_;
  CoreResponse core_response_;
  ResponseData response_data_, encrypted_response_data_;
  // Message buffers will be at least this big. Tests for loading and signing
  // messages will increase all buffers to this size.
  size_t required_message_size_;
  size_t required_core_message_size_;
  size_t required_request_signature_size_;
  std::vector<uint8_t> response_signature_;
  std::string serialized_core_message_;
  std::vector<uint8_t> encrypted_response_;

 private:
  // EncryptAndSignResponse output lengths.
  size_t encrypted_response_length_;
  size_t response_signature_length_;
};

class ProvisioningRoundTrip
    : public RoundTrip<
          /* CoreRequest */ oemcrypto_core_message::ODK_ProvisioningRequest,
          OEMCrypto_PrepAndSignProvisioningRequest,
          /* CoreResponse */ ODK_ParsedProvisioning,
          /* ResponseData */ RSAPrivateKeyMessage> {
 public:
  ProvisioningRoundTrip(Session* session,
                        const std::vector<uint8_t>& encoded_rsa_key)
      : RoundTrip(session),
        allowed_schemes_(kSign_RSASSA_PSS),
        keybox_(nullptr),
        encoded_rsa_key_(encoded_rsa_key) {}
  // Prepare the session for signing the request.
  virtual void PrepareSession(const wvoec::WidevineKeybox& keybox);
  void CreateDefaultResponse() override;
  void EncryptAndSignResponse() override;
  void EncryptAndSignResponseWithoutUpdatingEncPrivateKeyLength();
  void SignResponse();
  OEMCryptoResult LoadResponse() override { return LoadResponse(session_); }
  OEMCryptoResult LoadResponse(Session* session) override;
  void VerifyLoadFailed();
  const std::vector<uint8_t>& encoded_rsa_key() { return encoded_rsa_key_; }
  const std::vector<uint8_t>& wrapped_rsa_key() { return wrapped_rsa_key_; }
  void set_allowed_schemes(uint32_t allowed_schemes) {
    allowed_schemes_ = allowed_schemes;
  }
  // Used for OEMCrypto Fuzzing: Function to convert fuzzer data to valid
  // provisioning response data that can be parsed. Calculates signature for
  // data generated by fuzzer, so that signature validation passes when parsing
  // provisioning response.
  void InjectFuzzedResponseData(const uint8_t* data, size_t size);

 protected:
  bool RequestHasNonce() override { return true; }
  void VerifyRequestSignature(const vector<uint8_t>& data,
                              const vector<uint8_t>& generated_signature,
                              size_t core_message_length) override;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) override;
  // Load the response, without the retry. Called by LoadResponse.
  OEMCryptoResult LoadResponseNoRetry(Session* session,
                                      size_t* wrapped_key_length);
  // This takes a pointer in the response_data_ and remaps it to the same
  // pointer within the encrypted message. This is used for backwards
  // compatibility testing, so that a v15 oemcrypto will accept range checks.
  template <typename T>
  const T* RemapPointer(const T* response_pointer) const;

  uint32_t allowed_schemes_;
  Encryptor encryptor_;
  std::vector<uint8_t> request_;
  const wvoec::WidevineKeybox* keybox_;
  // The message key used for Prov 3.0.
  std::vector<uint8_t> encoded_rsa_key_;
  std::vector<uint8_t> wrapped_rsa_key_;
};

class Provisioning40RoundTrip
    : public RoundTrip<
          /* CoreRequest */ oemcrypto_core_message::ODK_Provisioning40Request,
          OEMCrypto_PrepAndSignProvisioningRequest,
          /* CoreResponse */ ODK_ParsedProvisioning,
          /* ResponseData */ Prov40CertMessage> {
 public:
  Provisioning40RoundTrip(Session* session)
      : RoundTrip(session), allowed_schemes_(kSign_RSASSA_PSS) {}
  void PrepareSession(bool is_oem_key);

  // Not used. Use Load*CertResponse() below to load OEM/DRM response
  // respectively.
  void CreateDefaultResponse() override {};
  void EncryptAndSignResponse() override {};
  OEMCryptoResult LoadResponse(Session* session) override {
    return OEMCrypto_ERROR_NOT_IMPLEMENTED;
  }

  OEMCryptoResult LoadOEMCertResponse();
  OEMCryptoResult LoadDRMCertResponse();

  const std::vector<uint8_t>& wrapped_oem_key() { return wrapped_oem_key_; }
  const std::vector<uint8_t>& oem_public_key() { return oem_public_key_; }
  OEMCrypto_PrivateKeyType oem_key_type() { return oem_key_type_; }
  const std::vector<uint8_t>& wrapped_drm_key() { return wrapped_drm_key_; }
  const std::vector<uint8_t>& drm_public_key() { return drm_public_key_; }
  OEMCrypto_PrivateKeyType drm_key_type() { return drm_key_type_; }
  void set_allowed_schemes(uint32_t allowed_schemes) {
    allowed_schemes_ = allowed_schemes;
  }

 protected:
  bool RequestHasNonce() override { return true; }
  void VerifyRequestSignature(const vector<uint8_t>& data,
                              const vector<uint8_t>& generated_signature,
                              size_t core_message_length) override;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) override;

  uint32_t allowed_schemes_;
  std::vector<uint8_t> wrapped_oem_key_;
  std::vector<uint8_t> oem_public_key_;
  OEMCrypto_PrivateKeyType oem_key_type_;
  std::vector<uint8_t> wrapped_drm_key_;
  std::vector<uint8_t> drm_public_key_;
  OEMCrypto_PrivateKeyType drm_key_type_;
};

class Provisioning40CastRoundTrip
    : public RoundTrip<
          /* CoreRequest */ oemcrypto_core_message::ODK_Provisioning40Request,
          OEMCrypto_PrepAndSignProvisioningRequest,
          /* CoreResponse */ ODK_ParsedProvisioning,
          /* ResponseData */ RSAPrivateKeyMessage> {
 public:
  Provisioning40CastRoundTrip(Session* session,
                              const std::vector<uint8_t>& encoded_rsa_key)
      : RoundTrip(session),
        allowed_schemes_(kSign_RSASSA_PSS),
        encryptor_(),
        encoded_rsa_key_(encoded_rsa_key) {}

  void PrepareSession();
  void LoadDRMPrivateKey();
  void CreateDefaultResponse() override;
  void SignResponse();
  void EncryptAndSignResponse() override;
  OEMCryptoResult LoadResponse() override { return LoadResponse(session_); }
  OEMCryptoResult LoadResponse(Session* session) override;
  OEMCryptoResult LoadResponseNoRetry(Session* session,
                                      size_t* wrapped_key_length);

  // Returned
  const std::vector<uint8_t>& wrapped_drm_key() { return wrapped_drm_key_; }
  const std::vector<uint8_t>& wrapped_rsa_key() { return wrapped_rsa_key_; }
  const std::vector<uint8_t>& drm_public_key() { return drm_public_key_; }
  OEMCrypto_PrivateKeyType drm_key_type() { return drm_key_type_; }
  void set_allowed_schemes(uint32_t allowed_schemes) {
    allowed_schemes_ = allowed_schemes;
  }

 protected:
  bool RequestHasNonce() override { return true; }
  void VerifyRequestSignature(const vector<uint8_t>& data,
                              const vector<uint8_t>& generated_signature,
                              size_t core_message_length) override;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) override;

  uint32_t allowed_schemes_;
  Encryptor encryptor_;
  std::vector<uint8_t> wrapped_oem_key_;
  std::vector<uint8_t> oem_public_key_;
  OEMCrypto_PrivateKeyType oem_key_type_;
  std::vector<uint8_t> wrapped_drm_key_;
  std::vector<uint8_t> drm_public_key_;
  OEMCrypto_PrivateKeyType drm_key_type_;
  std::vector<uint8_t> encoded_rsa_key_;
  std::vector<uint8_t> wrapped_rsa_key_;
};

class LicenseRoundTrip
    : public RoundTrip<
          /* CoreRequest */ oemcrypto_core_message::ODK_LicenseRequest,
          OEMCrypto_PrepAndSignLicenseRequest,
          /* CoreResponse */ ODK_Packing_ParsedLicense,
          /* ResponseData */ MessageData> {
 public:
  LicenseRoundTrip(Session* session)
      : RoundTrip(session),
        control_(wvoec::kControlNonceEnabled),
        num_keys_(4),
        pst_(),
        minimum_srm_version_(0),
        update_mac_keys_(true),
        api_version_(kCurrentAPI),
        expect_request_has_correct_nonce_(true),
        skip_request_hash_(global_features.api_version < 19),
        license_type_(OEMCrypto_ContentLicense),
        request_hash_() {}
  void CreateDefaultResponse() override;
  // Used for OEMCrypto Fuzzing: Function to inject fuzzed timer limits
  // into timer_limits field from core_response. We need to fuzz timer
  // limits in order to efficiently fuzz load renewal response API.
  void InjectFuzzedTimerLimits(OEMCrypto_Renewal_Response_Fuzz& fuzzed_data);
  // Used for OEMCrypto Fuzzing: Function to convert fuzzer data to valid
  // License response data that can be parsed. Calculates signature for data
  // generated by fuzzer, so that signature validation passes when parsing
  // license response.
  void InjectFuzzedResponseData(const uint8_t* data, size_t size);
  // Used for OEMCrypto Fuzzing: Convert boolean flags in parsed_license to
  // valid bytes to avoid errors from msan.
  void ConvertDataToValidBools(ODK_Packing_ParsedLicense* t);
  // Create a license with four keys. Each key is responsible for one of generic
  // encrypt (key 0), decrypt (key 1), sign (key 2) and verify (key 3). Each key
  // is allowed only one type of operation.
  void CreateResponseWithGenericCryptoKeys();
  // Fill the |core_response| substrings.
  virtual void FillCoreResponseSubstrings();
  void EncryptAndSignResponse() override;
  // Encrypt and sign license response created from a specific odk version.
  void EncryptAndSignResponseWithCoreMessageFeatures(
      const oemcrypto_core_message::features::CoreMessageFeatures& features,
      bool force_clear_kcb);
  // Encrypt license response. This is used in EncryptAndSignResponse().
  void EncryptResponse(bool force_clear_kcb = false);
  // Create core license response with a specific ODK version. This is used in
  // EncryptAndSignResponse().
  void CreateCoreLicenseResponseWithFeatures(
      const oemcrypto_core_message::features::CoreMessageFeatures& features);
  // Sign license response. This is used in EncryptAndSignResponse().
  void SignEncryptedResponse();
  OEMCryptoResult LoadResponse() override { return LoadResponse(session_); }
  OEMCryptoResult LoadResponse(Session* session) override;
  OEMCryptoResult LoadResponse(Session* session, bool verify_keys);
  // Reload an offline license into a different session. This derives new mac
  // keys and then calls LoadResponse.
  OEMCryptoResult ReloadResponse(Session* session);
  void VerifyTestKeys(Session* session);
  // Set the default key control block for all keys. This is used in
  // CreateDefaultResponse.  The key control block determines the restrictions
  // that OEMCrypto should place on a key's use. For example, it specifies the
  // minimum HDCP requirement and whether the key can only be used with a secure
  // video path. See the section "Key Control Block" in the document "Widevine
  // Modular DRM Security Integration Guide for CENC".
  void set_control(uint32_t control) { control_ = control; }
  uint32_t control() const { return control_; }
  // Set the number of keys to use in the license.
  void set_num_keys(uint32_t num_keys) { num_keys_ = num_keys; }
  uint32_t num_keys() const { return num_keys_; }
  // Get/Set the pst for the license and usage table entry.
  const std::string& pst() const { return pst_; }
  void set_pst(const std::string& pst) { pst_ = pst; }
  // Set the minimum SRM version for the license.
  void set_minimum_srm_version(uint32_t minimum_srm_version) {
    minimum_srm_version_ = minimum_srm_version;
  }
  // Change the hash of the core request. This should cause the response to be
  // rejected.
  void BreakRequestHash() { request_hash_[3] ^= 42; }
  // Set the API version for the license itself. This will be used in
  // CreateDefaultResponse.
  void set_api_version(uint32_t api_version) { api_version_ = api_version; }
  uint32_t api_version() const { return api_version_; }
  void set_update_mac_keys(bool update_mac_keys) {
    update_mac_keys_ = update_mac_keys;
  }
  void set_license_type(OEMCrypto_LicenseType license_type) {
    license_type_ = license_type;
  }
  // Skip the nonce check when verifying the license request.
  void skip_nonce_check() { expect_request_has_correct_nonce_ = false; }
  // Skip hashing license request before signing/KDF.
  void skip_request_hash() { skip_request_hash_ = true; }
  // This sets the key id of the specified key to the specified string.
  // This is used to test with different key id lengths.
  void SetKeyId(size_t index, const string& key_id);

 protected:
  bool RequestHasNonce() override { return true; }
  void VerifyRequestSignature(const vector<uint8_t>& data,
                              const vector<uint8_t>& generated_signature,
                              size_t core_message_length) override;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) override;

  // The default key control bits used with CreateDefaultResponse.
  uint32_t control_;
  // The number of keys in the license response.
  uint32_t num_keys_;
  // If non-empty, the license's provider session token.
  std::string pst_;
  // If non-zero, the minimum SRM version.
  uint32_t minimum_srm_version_;
  // If true, the license contains new mac keys for signing renewals.
  bool update_mac_keys_;
  // API version for the license itself. If this is 0 when the license request
  // is signed, it will be set to the same as OEMCrypto's API version.  It may
  // be set to a lower value in order to test backwards compatibility.
  uint32_t api_version_;
  // If true, then we expect the nonce in the core request to match that in
  // session. This is usually true, but when we are testing how OEMCrypto
  // handles a bad nonce, we don't want to.
  bool expect_request_has_correct_nonce_;
  // Whether to skip hashing the request before signing and KDF; this is used
  // for license protocol 2.2.
  bool skip_request_hash_;
  // Whether this is a content license or an entitlement license. Used in
  // CreateDefaultResponse.
  OEMCrypto_LicenseType license_type_;
  uint8_t request_hash_[ODK_SHA256_HASH_SIZE];
  // Used to hold and add/update key information to be transferred into the core
  // response later on.
  std::vector<OEMCrypto_KeyObject> key_array_;
};

class RenewalRoundTrip
    : public RoundTrip<
          /* CoreRequest */ oemcrypto_core_message::ODK_RenewalRequest,
          OEMCrypto_PrepAndSignRenewalRequest,
          //  Renewal response info is same as request:
          /* CoreResponse */ oemcrypto_core_message::ODK_RenewalRequest,
          /* ResponseData */ MessageData> {
 public:
  RenewalRoundTrip(LicenseRoundTrip* license_messages)
      : RoundTrip(license_messages->session()),
        license_messages_(license_messages),
        refresh_object_(),
        renewal_duration_seconds_(
            license_messages->core_response()
                .timer_limits.initial_renewal_duration_seconds),
        is_release_(false) {}
  void CreateDefaultResponse() override;
  void EncryptAndSignResponse() override;
  void InjectFuzzedResponseData(OEMCrypto_Renewal_Response_Fuzz& fuzzed_data,
                                const uint8_t* renewal_response,
                                size_t renewal_response_size);
  OEMCryptoResult LoadResponse() override { return LoadResponse(session_); }
  OEMCryptoResult LoadResponse(Session* session) override;
  uint64_t renewal_duration_seconds() const {
    return renewal_duration_seconds_;
  }
  void set_renewal_duration_seconds(uint64_t renewal_duration_seconds) {
    renewal_duration_seconds_ = renewal_duration_seconds;
  }
  void set_is_release(bool is_release) { is_release_ = is_release; }

 protected:
  bool RequestHasNonce() override { return false; }
  void VerifyRequestSignature(const vector<uint8_t>& data,
                              const vector<uint8_t>& generated_signature,
                              size_t core_message_length) override;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) override;
  LicenseRoundTrip* license_messages_;
  OEMCrypto_KeyRefreshObject refresh_object_;
  uint64_t renewal_duration_seconds_;
  bool is_release_;  // If this is a license release, and not a real renewal.
};

class ReleaseRoundTrip
    : public RoundTrip<
          /* CoreRequest */ oemcrypto_core_message::ODK_ReleaseRequest,
          OEMCrypto_PrepAndSignReleaseRequest,
          //  Release response info is same as request:
          /* CoreResponse */ oemcrypto_core_message::ODK_ReleaseRequest,
          /* ResponseData */ MessageData> {
 public:
  ReleaseRoundTrip(LicenseRoundTrip* license_messages)
      : RoundTrip(license_messages->session()),
        license_messages_(license_messages) {}
  void CreateDefaultResponse() override;
  void EncryptAndSignResponse() override;
  OEMCryptoResult LoadResponse() override { return LoadResponse(session_); }
  OEMCryptoResult LoadResponse(Session* session) override;
  int64_t seconds_since_license_received() const {
    return seconds_since_license_received_;
  }
  void set_seconds_since_license_received(
      int64_t seconds_since_license_received) {
    seconds_since_license_received_ = seconds_since_license_received;
  }
  int64_t seconds_since_first_decrypt() const {
    return seconds_since_first_decrypt_;
  }
  void set_seconds_since_first_decrypt(int64_t seconds_since_first_decrypt) {
    seconds_since_first_decrypt_ = seconds_since_first_decrypt;
  }

 protected:
  bool RequestHasNonce() override { return false; }
  void VerifyRequestSignature(const vector<uint8_t>& data,
                              const vector<uint8_t>& generated_signature,
                              size_t core_message_length) override;
  // Verify the values of the core response.
  virtual void FillAndVerifyCoreRequest(
      const std::string& core_message_string) override;
  LicenseRoundTrip* license_messages_;
  int64_t seconds_since_license_received_;
  int64_t seconds_since_first_decrypt_;
};

class EntitledMessage {
 public:
  EntitledMessage(LicenseRoundTrip* license_messages)
      : license_messages_(license_messages) {}
  void FillKeyArray();
  void MakeOneKey(size_t entitlement_key_index);
  void SetEntitledKeySession(uint32_t key_session) {
    entitled_key_session_ = key_session;
  }
  void LoadKeys(bool expected_success);
  OEMCryptoResult LoadKeys(const vector<uint8_t>& message);
  OEMCryptoResult LoadKeys();
  void EncryptContentKey();
  void LoadCasKeys(bool load_even, bool load_odd, OEMCryptoResult expected_sts);
  void set_num_keys(uint32_t num_keys) { num_keys_ = num_keys; }
  uint32_t num_keys() const { return num_keys_; }
  void SetEntitlementKeyId(unsigned int index, const std::string& key_id);
  void SetContentKeyId(unsigned int index, const std::string& key_id);
  OEMCrypto_EntitledContentKeyObject* entitled_key_array();
  // Returns entitled_key_data_ which is used as input message buffer to
  // load entitled content keys API.
  EntitledContentKeyData* entitled_key_data();
  size_t entitled_key_data_size();
  // Verify that key control blocks of the loaded keys.
  void VerifyEntitlementTestKeys();
  void VerifyEntitlementTestKey(size_t index);

 private:
  // Find the offset of the give pointer, relative to |entitled_key_data_|.
  OEMCrypto_Substring FindSubstring(const void* ptr, size_t size);
  // Verify that key control blocks of the loaded keys matches their entitlement
  // key.
  void VerifyKCBs();
  // Verify that decryption with the entitled keys works.
  void VerifyDecrypt();

  LicenseRoundTrip* license_messages_;
  uint32_t num_keys_ = 0;
  // Clear Entitlement key data. This is the backing data for
  // |entitled_key_array_|.
  EntitledContentKeyData entitled_key_data_[kMaxNumKeys] = {};
  // Entitled key object. Pointers are backed by |entitled_key_data_|.
  OEMCrypto_EntitledContentKeyObject entitled_key_array_[kMaxNumKeys] = {};
  uint32_t entitled_key_session_ = 0;
};

class Session {
 public:
  Session();
  ~Session();

  // Returns the most recently generated nonce.
  // Valid after call to GenerateNonce.
  uint32_t nonce() const { return nonce_; }
  // The nonce can be overridden.
  void set_nonce(uint32_t nonce) { nonce_ = nonce; }
  // Valid after call to open().
  uint32_t session_id() const { return (uint32_t)session_id_; }
  // Call OEMCrypto_OpenSession, with GTest ASSERTs.
  void open();
  // Call OEMCrypto_CloseSession, with GTest ASSERTs.
  void close();
  // Artifically set session id without calling OEMCrypto_OpenSession.  This is
  // used in core/test/generic_crypto_unittest.cpp.
  void SetSessionId(uint32_t session_id);
  uint32_t GetOecSessionId() { return session_id_; }
  // Generates one nonce.  If error_counter is null, this will sleep 1 second
  // and try again if a nonce flood has been detected.  If error_counter is
  // not null, it will be incremented when a nonce flood is detected.
  void GenerateNonce(int* error_counter = nullptr);
  // Fill the vector with test context which generate known mac and enc keys.
  std::vector<uint8_t> GetDefaultContext(bool do_hash = false);
  // Generate known mac and enc keys using OEMCrypto_GenerateDerivedKeys and
  // also fill out enc_key_, mac_key_server_, and mac_key_client_.
  void GenerateDerivedKeysFromKeybox(const wvoec::WidevineKeybox& keybox);
  void GenerateDerivedKeysFromKeybox(const wvoec::WidevineKeybox& keybox,
                                     const std::vector<uint8_t>& context);
  // Generate known mac and enc keys using OEMCrypto_DeriveKeysFromSessionKey
  // and also fill out enc_key_, mac_key_server_, and mac_key_client_.
  void GenerateDerivedKeysFromSessionKey();
  void GenerateDerivedKeysFromSessionKey(const std::vector<uint8_t>& context);
  // Encrypt some data and pass to OEMCrypto_DecryptCENC to verify decryption.
  void TestDecryptCTR(bool get_fresh_key_handle_first = true,
                      OEMCryptoResult expected_result = OEMCrypto_SUCCESS,
                      size_t key_index = 0);
  // Encrypt some data and pass to OEMCrypto_DecryptCENC to verify decryption
  // for entitled sessions.
  void TestDecryptEntitled(OEMCryptoResult expected_result = OEMCrypto_SUCCESS,
                           OEMCrypto_SESSION session = 0,
                           const uint8_t* content_key_id = nullptr,
                           size_t content_key_id_length = 0);
  // Verify that an attempt to get a key handle for an expired key either
  // succeeds or gives an actionable error code.
  void TestGetKeyHandleExpired(size_t key_index);
  // Calls OEMCrypto_GetOEMPublicCertificate and OEMCrypto_LoadOEMPrivateKey and
  // loads the OEM cert's public rsa key into public_rsa_.
  void LoadOEMCert(bool verify_cert = false);
  // Calls OEMCrypto_RewrapDeviceRSAKey with the given provisioning response
  // message.  If force is true, we assert that the key loads successfully.
  void RewrapRSAKey(const struct RSAPrivateKeyMessage& encrypted,
                    size_t message_size, const std::vector<uint8_t>& signature,
                    vector<uint8_t>* wrapped_key, bool force);
  // Loads the default test RSA public key into public_rsa_.
  void SetTestRsaPublicKey();
  // Loads the specified DRM or OEM public key into the appropriate key.
  // The provided key is serialized as an ASN.1 DER encoded PrivateKeyInfo.
  void SetPublicKeyFromPrivateKeyInfo(OEMCrypto_PrivateKeyType key_type,
                                      const uint8_t* buffer, size_t length);
  // Loads the specified RSA public key into public_rsa_.
  // The provided key is serialized as an ASN.1 DER encoded PrivateKeyInfo.
  void SetRsaPublicKeyFromPrivateKeyInfo(const uint8_t* buffer, size_t length);
  // Loads the specified EC public key into public_ec_.
  // The provided key is serialized as an ASN.1 DER encoded PrivateKeyInfo.
  void SetEccPublicKeyFromPrivateKeyInfo(const uint8_t* buffer, size_t length);

  // Loads the specified DRM or OEM public key into the appropriate key.
  // The provided key is serialized as an ASN.1 DER encoded SubjectPublicKey.
  void SetPublicKeyFromSubjectPublicKey(OEMCrypto_PrivateKeyType key_type,
                                        const uint8_t* buffer, size_t length);
  // Loads the specified RSA public key into public_rsa_.
  // The provided key is serialized as an ASN.1 DER encoded SubjectPublicKey.
  void SetRsaPublicKeyFromSubjectPublicKey(const uint8_t* buffer,
                                           size_t length);
  // Loads the specified EC public key into public_ec_.
  // The provided key is serialized as an ASN.1 DER encoded SubjectPublicKey.
  void SetEccPublicKeyFromSubjectPublicKey(const uint8_t* buffer,
                                           size_t length);

  // Verify that the message was signed by the private key associated with
  // |public_rsa_| using the specified padding scheme.
  void VerifyRsaSignature(const vector<uint8_t>& message,
                          const uint8_t* signature, size_t signature_length,
                          RSA_Padding_Scheme padding_scheme);
  // Verify that the message was signed by the private key associated with
  // |public_ecc_| using Widevine ECDSA.
  void VerifyEccSignature(const vector<uint8_t>& message,
                          const uint8_t* signature, size_t signature_length);
  // Verify RSA or ECC signature based on the key type installed. The
  // padding_scheme will be ignored in case of ECC key.
  void VerifySignature(const vector<uint8_t>& message, const uint8_t* signature,
                       size_t signature_length,
                       RSA_Padding_Scheme padding_scheme);

  // Encrypts a known session key with public_rsa_ for use in future calls to
  // OEMCrypto_DeriveKeysFromSessionKey or OEMCrypto_RewrapDeviceRSAKey30.
  // The unencrypted session key is stored in session_key.
  bool GenerateRsaSessionKey();
  // Derives a session key with public_ec_ and a ephemeral "server" ECC key
  // for use in future calls to OEMCrypto_DeriveKeysFromSessionKey.
  // The unencrypted session key is stored in session_key.
  bool GenerateEccSessionKey();
  // Based on the key type installed, call GenerateRsaSessionKey or
  // GenerateEccSessionKey.
  bool GenerateSessionKey();

  // Calls OEMCrypto_RewrapDeviceRSAKey30 with the given provisioning response
  // message. If force is true, we assert that the key loads successfully.
  void RewrapRSAKey30(const struct RSAPrivateKeyMessage& encrypted,
                      const std::vector<uint8_t>& encrypted_message_key,
                      vector<uint8_t>* wrapped_key, bool force);

  void LoadWrappedDrmKey(OEMCrypto_PrivateKeyType key_type,
                         const vector<uint8_t>& wrapped_drm_key);
  // Loads the specified wrapped_rsa_key into OEMCrypto.
  void LoadWrappedRsaDrmKey(const vector<uint8_t>& wrapped_rsa_key);
  // Loads the specified wrapped_ecc_key into OEMCrypto.
  void LoadWrappedEccDrmKey(const vector<uint8_t>& wrapped_ecc_key);

  // Creates a new usage entry, and keeps track of the index.
  // If status is null, we expect success, otherwise status is set to the
  // return value.
  void CreateNewUsageEntry(OEMCryptoResult* status = nullptr);
  // Copy encrypted usage entry from other session, and then load it.
  // This session must already be open.
  void LoadUsageEntry(uint32_t index, const vector<uint8_t>& buffer);
  // Copy encrypted usage entry from other session.
  // This session must already be open.
  void LoadUsageEntry(const Session& other) {
    LoadUsageEntry(other.usage_entry_number(), other.encrypted_usage_entry());
  }
  // Reload previously used usage entry.
  void ReloadUsageEntry() { LoadUsageEntry(*this); }
  // Update the usage entry and save the header to the specified buffer.
  void UpdateUsageEntry(std::vector<uint8_t>* header_buffer);
  // The usage entry number for this session's usage entry.
  uint32_t usage_entry_number() const { return usage_entry_number_; }
  void set_usage_entry_number(uint32_t v) { usage_entry_number_ = v; }
  // The encrypted buffer holding the recently updated and saved usage entry.
  const vector<uint8_t>& encrypted_usage_entry() const {
    return encrypted_usage_entry_;
  }
  // Generates a usage report for the specified pst.  If there is success,
  // the report's signature is verified, and several fields are given sanity
  // checks.  If |other| is not null, then the mac keys are copied from other in
  // order to verify signatures.
  void GenerateReport(const std::string& pst,
                      OEMCryptoResult expected_result = OEMCrypto_SUCCESS,
                      Session* other = nullptr);
  // Move this usage entry to a new index.
  void MoveUsageEntry(uint32_t new_index, std::vector<uint8_t>* header_buffer,
                      OEMCryptoResult expect_result = OEMCrypto_SUCCESS);
  // PST used in FillSimpleMesage.
  string pst() const { return pst_; }
  // Returns a pointer-like thing to the usage report generated by the previous
  // call to GenerateReport.
  wvutil::Unpacked_PST_Report pst_report() {
    return wvutil::Unpacked_PST_Report(&pst_report_buffer_[0]);
  }
  // Verify the values in the PST report.  The signature should have been
  // verified in GenerateReport, above.
  void VerifyPST(const Test_PST_Report& report);
  // Verify the Usage Report. If any time is greater than 10 minutes, it is
  // assumed to be an absolute time, and time_since will be computed relative to
  // now.
  void VerifyReport(Test_PST_Report report, int64_t time_license_received = 0,
                    int64_t time_first_decrypt = 0,
                    int64_t time_last_decrypt = 0);
  // Create an entry in the old usage table based on the given report.
  void CreateOldEntry(const Test_PST_Report& report);
  // Create a new entry and copy the old entry into it.  Then very the report
  // is right.
  void CopyAndVerifyOldEntry(const Test_PST_Report& report,
                             std::vector<uint8_t>* header_buffer);

  // The unencrypted license response or license renewal response.
  MessageData& license() { return license_; }

  void set_license(const MessageData& license) { license_ = license; }

  // Gives access to the last key handle used by the Session. Useful if a test
  // case uses one of the decrypt test functions above and then wants to perform
  // further crypto operations using the same key handle. Also useful if a test
  // case needs to inject a specific erroneous key handle into the decrypt test
  // functions.
  vector<uint8_t>& key_handle() { return key_handle_; }

  const std::vector<uint8_t>& session_key() const { return session_key_; }
  const std::vector<uint8_t>& enc_session_key() const {
    return enc_session_key_;
  }

  const KeyDeriver& key_deriver() const { return key_deriver_; }
  void set_mac_keys(const uint8_t* mac_keys) {
    key_deriver_.set_mac_keys(mac_keys);
  }
  bool IsPublicKeySet() {
    return public_rsa_ != nullptr || public_ec_ != nullptr;
  }

 private:
  OEMCryptoResult GetKeyHandle(
      vector<uint8_t>& key_handle, size_t key_index = 0,
      OEMCryptoCipherMode cipher_mode = OEMCrypto_CipherMode_CENC);
  // This compares the actual result with the expected result.  If OEMCrypto is
  // an older version, we allow it to report an equivalent error code.
  void TestDecryptResult(OEMCryptoResult expected_result,
                         OEMCryptoResult actual_getkeyhandle_result,
                         OEMCryptoResult actual_decryt_result);

  bool open_ = false;
  bool forced_session_id_ = false;
  OEMCrypto_SESSION session_id_ = 0;
  KeyDeriver key_deriver_;
  uint32_t nonce_ = 0;
  // Only one of RSA or EC should be set.
  std::unique_ptr<util::RsaPublicKey> public_rsa_;
  std::unique_ptr<util::EccPublicKey> public_ec_;
  // In provisioning 4.0, the shared session key is derived from either
  //    1. (client side) client private key + server ephemeral public key, or
  //    2. (server side) server ephemeral private key + client public key
  // Encryption key and mac keys are derived from the shared session key, and
  // are inserted in to the default license response which simulates the
  // response from a license server. In order for these keys to be deterministic
  // across multiple test calls of GenerateDerivedKeysFromSessionKey(), which
  // simulates how the server derives keys, the ephemeral keys used by the
  // "server" need to be stored for re-use.
  static std::unordered_map<
      util::EccCurve, std::unique_ptr<util::EccPrivateKey>, std::hash<int>>
      server_ephemeral_keys_;
  static std::mutex ephemeral_key_map_lock_;
  vector<uint8_t> pst_report_buffer_;
  MessageData license_ = {};
  vector<uint8_t> key_handle_;
  std::vector<uint8_t> session_key_;
  std::vector<uint8_t> enc_session_key_;

  vector<uint8_t> encrypted_usage_entry_;
  uint32_t usage_entry_number_ = 0;
  string pst_;
};  // class Session

// Used for OEMCrypto Fuzzing: Convert byte to a valid boolean to avoid errors
// generated by msan.
bool ConvertByteToValidBoolean(const bool* in);
// Used for OEMCrypto Fuzzing: Generates corpus for request APIs.
template <class CoreRequest>
void WriteRequestApiCorpus(size_t signature_length, size_t core_message_length,
                           vector<uint8_t>& data);
template <PrepAndSignRequest_t PrepAndSignRequest>
void GetDefaultRequestSignatureAndCoreMessageLengths(
    uint32_t& session_id, const size_t& small_size,
    size_t* gen_signature_length, size_t* core_message_length);
// Loads the key matching the given |key_id| into the |session| in OEMCrypto for
// the given |cipher_mode| and returns a handle to that key. This function
// handles negotiating the size of the |key_handle| buffer. For non-bypassing
// systems, this is equivalent to the old SelectKey call and will deselect any
// previous keys selected in the session. For bypassing systems, multiple key
// handles may be valid simultaneously, but this call may invalidate any
// previous handles.
OEMCryptoResult GetKeyHandleIntoVector(OEMCrypto_SESSION session,
                                       const uint8_t* key_id,
                                       size_t key_id_length,
                                       OEMCryptoCipherMode cipher_mode,
                                       vector<uint8_t>& key_handle);

}  // namespace wvoec

#endif  // CDM_OEC_SESSION_UTIL_H_