Add Entitlement License to OEMCrypto

This CL adds entitlement license features and moves cipher mode from LoadKeys to SelectKeys. Merge from Widevine repo of http://go/wvgerrit/41660 bug: 70334840 Entitlement License - cdm layer bug: 70334345 Entitlement License - reference code and unit tests test: Entitlement license unit tests pass. Change-Id: Ic7d7f42c15e6d83ef7fcfd8a866c778adc4c8095
2018-01-23 15:55:43 -08:00
parent 95fa4ffca9
commit 979ed70c7b
11 changed files with 1136 additions and 595 deletions
--- a/libwvdrmengine/oemcrypto/mock/src/oemcrypto_key_mock.h
+++ b/libwvdrmengine/oemcrypto/mock/src/oemcrypto_key_mock.h
@@ -69,15 +69,15 @@ class Key {
 public:
  Key(const Key& key)
      : value_(key.value_), control_(key.control_), ctr_mode_(key.ctr_mode_) {}
-  Key(const std::vector<uint8_t>& key_string, const KeyControlBlock& control,
-      bool ctr_mode)
-      : value_(key_string), control_(control), ctr_mode_(ctr_mode){};
+  Key(const std::vector<uint8_t>& key_string, const KeyControlBlock& control)
+      : value_(key_string), control_(control), ctr_mode_(true){};

  virtual ~Key(){};
  void UpdateDuration(const KeyControlBlock& control);
-  const std::vector<uint8_t>& value() const { return value_; }
+  virtual const std::vector<uint8_t>& value() const { return value_; }
  const KeyControlBlock& control() const { return control_; }
  bool ctr_mode() const { return ctr_mode_; }
+  void set_ctr_mode(bool ctr_mode) { ctr_mode_ = ctr_mode; }

 private:
  std::vector<uint8_t> value_;
@@ -85,6 +85,28 @@ class Key {
  bool ctr_mode_;
 };

+// AES-256 entitlement key. |Key| holds the entitlement key. |EntitlementKey|
+// holds the content key.
+class EntitlementKey : public Key {
+ public:
+  EntitlementKey(const Key& key) : Key(key) {}
+  virtual ~EntitlementKey() {}
+  virtual const std::vector<uint8_t>& value() const { return content_key_; }
+  const std::vector<uint8_t>& content_key() { return content_key_; }
+  const std::vector<uint8_t>& content_key_id() { return content_key_id_; }
+  const std::vector<uint8_t>& entitlement_key() { return Key::value(); }
+  bool SetContentKey(const std::vector<uint8_t>& content_key,
+                     const std::vector<uint8_t>& content_key_id) {
+    content_key_.assign(content_key.begin(), content_key.end());
+    content_key_id_.assign(content_key_id.begin(), content_key_id.end());
+    return true;
+  }
+
+ private:
+  std::vector<uint8_t> content_key_;
+  std::vector<uint8_t> content_key_id_;
+};
+
 }  // namespace wvoec_mock

 #endif  // OEMCRYPTO_KEY_MOCK_H_
--- a/libwvdrmengine/oemcrypto/mock/src/oemcrypto_mock.cpp
+++ b/libwvdrmengine/oemcrypto/mock/src/oemcrypto_mock.cpp
@@ -31,7 +31,7 @@ const uint8_t kBakedInCertificateMagicBytes[] = {0xDE, 0xAD, 0xBE, 0xEF};

 // Return uint32 referenced through a potentially unaligned pointer.
 // If the pointer is NULL, return 0.
-uint32_t unaligned_dereference_uint32(const uint32_t* unaligned_ptr) {
+uint32_t unaligned_dereference_uint32(const void* unaligned_ptr) {
  if (unaligned_ptr == NULL) return 0;
  uint32_t value;
  const uint8_t* src = reinterpret_cast<const uint8_t*>(unaligned_ptr);
@@ -299,7 +299,7 @@ extern "C" OEMCryptoResult OEMCrypto_LoadKeys(
    const uint8_t* signature, size_t signature_length,
    const uint8_t* enc_mac_key_iv, const uint8_t* enc_mac_keys, size_t num_keys,
    const OEMCrypto_KeyObject* key_array, const uint8_t* pst, size_t pst_length,
-    const uint8_t* srm_requirement) {
+    const uint8_t* srm_requirement, OEMCrypto_LicenseType license_type) {
  if (!crypto_engine) {
    LOGE("OEMCrypto_LoadKeys: OEMCrypto Not Initialized.");
    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
@@ -326,9 +326,6 @@ extern "C" OEMCryptoResult OEMCrypto_LoadKeys(
                        key_array[i].key_control_iv, wvcdm::KEY_IV_SIZE);
        dump_array_part("key_array", i, "key_control", key_array[i].key_control,
                        wvcdm::KEY_IV_SIZE);
-        LOGV("key_array[%zu].cipher_mode=%s;\n", i,
-             key_array[i].cipher_mode == OEMCrypto_CipherMode_CTR ? "CTR"
-                                                                  : "CBC");
      }
    }
  }
@@ -389,7 +386,32 @@ extern "C" OEMCryptoResult OEMCrypto_LoadKeys(
  return session_ctx->LoadKeys(message, message_length, signature,
                               signature_length, enc_mac_key_iv, enc_mac_keys,
                               num_keys, key_array, pst, pst_length,
-                               srm_requirement);
+                               srm_requirement, license_type);
+}
+
+extern "C" OEMCryptoResult OEMCrypto_LoadEntitledContentKeys(
+    OEMCrypto_SESSION session,
+    size_t num_keys,
+    const OEMCrypto_EntitledContentKeyObject* key_array) {
+  if (num_keys == 0) {
+    LOGE("[OEMCrypto_LoadEntitledContentKeys(): key_array is empty.");
+    return OEMCrypto_SUCCESS;
+  }
+  if (!key_array) {
+    LOGE("[OEMCrypto_LoadEntitledContentKeys(): missing key_array.");
+    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+  }
+  if (!crypto_engine) {
+    LOGE("OEMCrypto_LoadEntitledContentKeys: OEMCrypto Not Initialized.");
+    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+  }
+  SessionContext* session_ctx = crypto_engine->FindSession(session);
+  if (!session_ctx || !session_ctx->isValid()) {
+    LOGE("[OEMCrypto_LoadEntitledContentKeys(): ERROR_INVALID_SESSION]");
+    return OEMCrypto_ERROR_INVALID_SESSION;
+  }
+
+  return session_ctx->LoadEntitledContentKeys(num_keys, key_array);
 }

 extern "C" OEMCryptoResult OEMCrypto_RefreshKeys(
@@ -521,9 +543,9 @@ extern "C" OEMCryptoResult OEMCrypto_QueryKeyControl(
  return OEMCrypto_SUCCESS;
 }

-extern "C" OEMCryptoResult OEMCrypto_SelectKey(const OEMCrypto_SESSION session,
-                                               const uint8_t* key_id,
-                                               size_t key_id_length) {
+extern "C" OEMCryptoResult OEMCrypto_SelectKey(
+    const OEMCrypto_SESSION session, const uint8_t* key_id,
+    size_t key_id_length, OEMCryptoCipherMode cipher_mode) {
  if (LogCategoryEnabled(kLoggingTraceDecryptCalls)) {
    LOGI("-- OEMCryptoResult OEMCrypto_SelectKey(%d, id=%s)", session,
         wvcdm::HexEncode(key_id, key_id_length).c_str());
@@ -543,7 +565,7 @@ extern "C" OEMCryptoResult OEMCrypto_SelectKey(const OEMCrypto_SESSION session,

  const std::vector<uint8_t> key_id_str =
      std::vector<uint8_t>(key_id, key_id + key_id_length);
-  return session_ctx->SelectContentKey(key_id_str);
+  return session_ctx->SelectContentKey(key_id_str, cipher_mode);
 }

 extern "C" OEMCryptoResult OEMCrypto_DecryptCENC(
--- a/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session.cpp
+++ b/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session.cpp
@@ -56,11 +56,124 @@ void dump_boringssl_error() {

 namespace wvoec_mock {

+/***************************************/
+
+class ContentKeysContext : public SessionContextKeys {
+ public:
+  explicit ContentKeysContext() {}
+  virtual ~ContentKeysContext() {}
+  virtual size_t size() { return session_keys_.size(); }
+  bool Insert(const KeyId& key_id, const Key& key_data);
+  virtual Key* Find(const KeyId& key_id);
+  virtual void Remove(const KeyId& key_id);
+  virtual void UpdateDuration(const KeyControlBlock& control);
+
+  virtual OEMCrypto_LicenseType type() { return OEMCrypto_ContentLicense; }
+
+  virtual bool SetContentKey(const KeyId& entitlement_id,
+                             const KeyId& content_key_id,
+                             const std::vector<uint8_t>& content_key);
+  virtual bool GetEntitlementKey(
+      const KeyId& entitlement_id,
+      const std::vector<uint8_t>** entitlement_key);
+
+ private:
+  SessionKeyTable session_keys_;
+  CORE_DISALLOW_COPY_AND_ASSIGN(ContentKeysContext);
+};
+
+bool ContentKeysContext::Insert(const KeyId& key_id, const Key& key_data) {
+  return session_keys_.Insert(key_id, key_data);
+}
+
+Key* ContentKeysContext::Find(const KeyId& key_id) {
+  return session_keys_.Find(key_id);
+}
+
+void ContentKeysContext::Remove(const KeyId& key_id) {
+  session_keys_.Remove(key_id);
+}
+
+void ContentKeysContext::UpdateDuration(const KeyControlBlock& control) {
+  session_keys_.UpdateDuration(control);
+}
+
+bool ContentKeysContext::SetContentKey(
+    const KeyId& entitlement_id, const KeyId& content_key_id,
+    const std::vector<uint8_t>& content_key) {
+  // Unsupported action for this type.
+  return false;
+}
+
+bool ContentKeysContext::GetEntitlementKey(
+    const KeyId& entitlement_id, const std::vector<uint8_t>** key) {
+  // Unsupported action for this type.
+  return false;
+};
+
+/***************************************/
+
+class EntitlementKeysContext : public SessionContextKeys {
+ public:
+  EntitlementKeysContext() {}
+  virtual ~EntitlementKeysContext() {}
+  virtual size_t size() { return session_keys_.size(); }
+  bool Insert(const KeyId& key_id, const Key& key_data);
+  virtual Key* Find(const KeyId& key_id);
+  virtual void Remove(const KeyId& key_id);
+  virtual void UpdateDuration(const KeyControlBlock& control);
+  virtual bool SetContentKey(const KeyId& entitlement_id,
+                             const KeyId& content_key_id,
+                             const std::vector<uint8_t>& content_key);
+  virtual bool GetEntitlementKey(const KeyId& entitlement_id,
+                                 const std::vector<uint8_t>** key);
+
+  virtual OEMCrypto_LicenseType type() { return OEMCrypto_EntitlementLicense; }
+
+ private:
+  EntitlementKeyTable session_keys_;
+  CORE_DISALLOW_COPY_AND_ASSIGN(EntitlementKeysContext);
+};
+
+bool EntitlementKeysContext::Insert(const KeyId& key_id, const Key& key_data) {
+  return session_keys_.Insert(key_id, key_data);
+}
+
+Key* EntitlementKeysContext::Find(const KeyId& key_id) {
+  return session_keys_.Find(key_id);
+}
+
+void EntitlementKeysContext::Remove(const KeyId& key_id) {
+  session_keys_.Remove(key_id);
+}
+
+void EntitlementKeysContext::UpdateDuration(const KeyControlBlock& control) {
+  session_keys_.UpdateDuration(control);
+}
+
+bool EntitlementKeysContext::SetContentKey(
+    const KeyId& entitlement_id, const KeyId& content_key_id,
+    const std::vector<uint8_t>& content_key) {
+  return session_keys_.SetContentKey(entitlement_id, content_key_id,
+                                     content_key);
+}
+
+bool EntitlementKeysContext::GetEntitlementKey(
+    const KeyId& entitlement_id, const std::vector<uint8_t>** out_key) {
+  return session_keys_.GetEntitlementKey(entitlement_id, out_key);
+}
+
+/***************************************/
+
 SessionContext::~SessionContext() {
  if (usage_entry_) {
    delete usage_entry_;
    usage_entry_ = NULL;
  }
+  if (session_keys_) {
+    delete session_keys_;
+    session_keys_ = NULL;
+  }
 }

 // Internal utility function to derive key using CMAC-128
@@ -398,7 +511,7 @@ OEMCryptoResult SessionContext::CheckNonceOrEntry(
      return CheckStatusOnline(key_control_block.nonce(),
                               key_control_block.control_bits());
      break;
-    case kControlNonceOrEntry: // Offline license. Nonce required on first use.
+    case kControlNonceOrEntry:  // Offline license. Nonce required on first use.
      return CheckStatusOffline(key_control_block.nonce(),
                                key_control_block.control_bits());
      break;
@@ -423,13 +536,32 @@ OEMCryptoResult SessionContext::LoadKeys(
    const uint8_t* message, size_t message_length, const uint8_t* signature,
    size_t signature_length, const uint8_t* enc_mac_key_iv,
    const uint8_t* enc_mac_keys, size_t num_keys,
-    const OEMCrypto_KeyObject* key_array, const uint8_t* pst,
-    size_t pst_length, const uint8_t* srm_requirement) {
+    const OEMCrypto_KeyObject* key_array, const uint8_t* pst, size_t pst_length,
+    const uint8_t* srm_requirement, OEMCrypto_LicenseType license_type) {
  // Validate message signature
  if (!ValidateMessage(message, message_length, signature, signature_length)) {
    return OEMCrypto_ERROR_SIGNATURE_FAILURE;
  }

+  if (!session_keys_) {
+    switch (license_type) {
+      case OEMCrypto_ContentLicense:
+        session_keys_ = new ContentKeysContext();
+        break;
+
+      case OEMCrypto_EntitlementLicense:
+        session_keys_ = new EntitlementKeysContext();
+        break;
+
+      default:
+        return OEMCrypto_ERROR_INVALID_CONTEXT;
+    }
+  } else {
+    if (session_keys_->type() != license_type) {
+      return OEMCrypto_ERROR_INVALID_CONTEXT;
+    }
+  }
+
  StartTimer();

  if (srm_requirement) {
@@ -454,15 +586,15 @@ OEMCryptoResult SessionContext::LoadKeys(
      LOGW("[LoadKeys: SRM blacklisted device attached]");
      srm_requirements_status_ = InvalidSRMVersion;
    } else {
-      LOGI("[LoadKeys: SRM Versions is %d, required: %d]",
-           current_version, minimum_version);
+      LOGI("[LoadKeys: SRM Versions is %d, required: %d]", current_version,
+           minimum_version);
      srm_requirements_status_ = ValidSRMVersion;
    }
  }

  // If there are already keys installed in this session, then we can load
  // a shared license.
-  bool second_license = (session_keys_.size() > 0);
+  bool second_license = (session_keys_->size() > 0);

  // Decrypt and install keys in key object
  // Each key will have a key control block.  They will all have the same nonce.
@@ -490,7 +622,7 @@ OEMCryptoResult SessionContext::LoadKeys(

    OEMCryptoResult result = InstallKey(
        key_id, enc_key_data, key_data_iv, key_control, key_control_iv,
-        key_array[i].cipher_mode == OEMCrypto_CipherMode_CTR, second_license);
+        second_license);
    if (result != OEMCrypto_SUCCESS) {
      status = result;
      break;
@@ -550,11 +682,57 @@ OEMCryptoResult SessionContext::LoadKeys(
  return OEMCrypto_SUCCESS;
 }

+OEMCryptoResult SessionContext::LoadEntitledContentKeys(
+      size_t num_keys,
+      const OEMCrypto_EntitledContentKeyObject* key_array) {
+  if (!key_array) {
+    return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+  }
+  if (!session_keys_ || session_keys_->type() != OEMCrypto_EntitlementLicense) {
+    return OEMCrypto_ERROR_INVALID_CONTEXT;
+  }
+  for (size_t i = 0; i < num_keys; ++i) {
+    const OEMCrypto_EntitledContentKeyObject* key_data = &key_array[i];
+    std::vector<uint8_t> entitlement_key_id;
+    entitlement_key_id.assign(key_data->entitlement_key_id,
+                              key_data->entitlement_key_id +
+                              key_data->entitlement_key_id_length);
+
+    const std::vector<uint8_t>* entitlement_key = NULL;
+    if (!session_keys_->GetEntitlementKey(entitlement_key_id,
+                                          &entitlement_key)) {
+      return OEMCrypto_KEY_NOT_ENTITLED;
+    }
+    std::vector<uint8_t> content_key;
+    std::vector<uint8_t> iv;
+    std::vector<uint8_t> encrypted_content_key;
+    std::vector<uint8_t> content_key_id;
+
+    iv.assign(key_data->content_key_data_iv,
+              key_data->content_key_data_iv + 16);
+    encrypted_content_key.assign(
+        key_data->content_key_data,
+        key_data->content_key_data + key_data->content_key_data_length);
+    content_key_id.assign(
+        key_data->content_key_id,
+        key_data->content_key_id + key_data->content_key_id_length);
+    if (!DecryptEntitlement(*entitlement_key, iv,
+                            encrypted_content_key, &content_key)) {
+      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+    }
+    if (!session_keys_->SetContentKey(
+        entitlement_key_id, content_key_id, content_key)) {
+      return OEMCrypto_ERROR_UNKNOWN_FAILURE;
+    }
+  }
+  return OEMCrypto_SUCCESS;
+}
+
 OEMCryptoResult SessionContext::InstallKey(
    const KeyId& key_id, const std::vector<uint8_t>& key_data,
    const std::vector<uint8_t>& key_data_iv,
    const std::vector<uint8_t>& key_control,
-    const std::vector<uint8_t>& key_control_iv, bool ctr_mode,
+    const std::vector<uint8_t>& key_control_iv,
    bool second_license) {
  // Decrypt encrypted key_data using derived encryption key and offered iv
  std::vector<uint8_t> content_key;
@@ -566,12 +744,11 @@ OEMCryptoResult SessionContext::InstallKey(
  }

  if (LogCategoryEnabled(kLoggingDumpContentKeys)) {
-    LOGI(("  InstallKey: key_id      = " +
-          wvcdm::b2a_hex(key_id)).c_str());
-    LOGI(("  InstallKey: content_key = " +
-          wvcdm::b2a_hex(content_key)).c_str());
-    LOGI(("  InstallKey: key_control = " +
-          wvcdm::b2a_hex(key_control_str)).c_str());
+    LOGI(("  InstallKey: key_id      = " + wvcdm::b2a_hex(key_id)).c_str());
+    LOGI(
+        ("  InstallKey: content_key = " + wvcdm::b2a_hex(content_key)).c_str());
+    LOGI(("  InstallKey: key_control = " + wvcdm::b2a_hex(key_control_str))
+             .c_str());
  }

  // Key control must be supplied by license server
@@ -633,8 +810,11 @@ OEMCryptoResult SessionContext::InstallKey(
    }
  }

-  Key key(content_key, key_control_block, ctr_mode);
-  session_keys_.Insert(key_id, key);
+  Key key(content_key, key_control_block);
+  if (!session_keys_) {
+    return OEMCrypto_ERROR_INVALID_CONTEXT;
+  }
+  session_keys_->Insert(key_id, key);
  return OEMCrypto_SUCCESS;
 }

@@ -663,6 +843,9 @@ bool SessionContext::InstallRSAEncryptedKey(
 OEMCryptoResult SessionContext::RefreshKey(
    const KeyId& key_id, const std::vector<uint8_t>& key_control,
    const std::vector<uint8_t>& key_control_iv) {
+  if (!session_keys_) {
+    return OEMCrypto_ERROR_INVALID_CONTEXT;
+  }
  if (key_id.empty()) {
    // Key control is not encrypted if key id is NULL
    KeyControlBlock key_control_block(key_control);
@@ -676,11 +859,11 @@ OEMCryptoResult SessionContext::RefreshKey(
      return OEMCrypto_ERROR_INVALID_NONCE;
    }
    // Apply duration to all keys in this session
-    session_keys_.UpdateDuration(key_control_block);
+    session_keys_->UpdateDuration(key_control_block);
    return OEMCrypto_SUCCESS;
  }

-  Key* content_key = session_keys_.Find(key_id);
+  Key* content_key = session_keys_->Find(key_id);

  if (NULL == content_key) {
    if (LogCategoryEnabled(kLoggingDumpKeyControlBlocks)) {
@@ -1014,7 +1197,10 @@ bool SessionContext::UpdateMacKeys(const std::vector<uint8_t>& enc_mac_keys,
 }

 bool SessionContext::QueryKeyControlBlock(const KeyId& key_id, uint32_t* data) {
-  const Key* content_key = session_keys_.Find(key_id);
+  if (!session_keys_) {
+    return false;
+  }
+  const Key* content_key = session_keys_->Find(key_id);
  if (LogCategoryEnabled(kLoggingTraceDecryption)) {
    LOGI(("Select Key: key_id = " + wvcdm::b2a_hex(key_id)).c_str());
    if (content_key) {
@@ -1035,25 +1221,34 @@ bool SessionContext::QueryKeyControlBlock(const KeyId& key_id, uint32_t* data) {
  return true;
 }

-OEMCryptoResult SessionContext::SelectContentKey(const KeyId& key_id) {
-  const Key* content_key = session_keys_.Find(key_id);
-
+OEMCryptoResult SessionContext::SelectContentKey(
+    const KeyId& key_id, OEMCryptoCipherMode cipher_mode) {
  if (LogCategoryEnabled(kLoggingTraceDecryption)) {
-    LOGI(("  Select Key: key_id      = " + wvcdm::b2a_hex(key_id)).c_str());
-    LOGI(("  Select Key: key = " + wvcdm::b2a_hex(content_key->value()))
-             .c_str());
+    LOGI("  Select Key: key_id      = %s", wvcdm::b2a_hex(key_id).c_str());
+    LOGI("  Select Key: cipher_mode = %s",
+         (cipher_mode == OEMCrypto_CipherMode_CTR) ? "CTR" : "CBC");
  }
-
+  if (!session_keys_) {
+    LOGE("Select Key: no session keys.");
+    return OEMCrypto_ERROR_INVALID_CONTEXT;
+  }
+  Key* content_key = session_keys_->Find(key_id);
  if (NULL == content_key) {
    LOGE("[SelectContentKey(): No key matches key id]");
    return OEMCrypto_ERROR_NO_CONTENT_KEY;
  }
+  if (LogCategoryEnabled(kLoggingTraceDecryption)) {
+    LOGI(("  Select Key: key = " + wvcdm::b2a_hex(content_key->value()))
+             .c_str());
+  }
+  content_key->set_ctr_mode(cipher_mode == OEMCrypto_CipherMode_CTR);
  current_content_key_ = content_key;
  const KeyControlBlock& control = current_content_key()->control();
+
  if (control.duration() > 0) {
    if (control.duration() < CurrentTimer()) {
-      LOGE("[SelectContentKey(): KEY_EXPIRED %d versus %d]",
-           control.duration(), CurrentTimer());
+      LOGE("[SelectContentKey(): KEY_EXPIRED %d versus %d]", control.duration(),
+           CurrentTimer());
      return OEMCrypto_ERROR_KEY_EXPIRED;
    }
  }
@@ -1131,7 +1326,6 @@ OEMCryptoResult SessionContext::CopyOldUsageEntry(
  return usage_entry_->CopyOldUsageEntry(pst);
 }

-
 // Internal utility function to decrypt the message
 bool SessionContext::DecryptMessage(const std::vector<uint8_t>& key,
                                    const std::vector<uint8_t>& iv,
@@ -1151,6 +1345,25 @@ bool SessionContext::DecryptMessage(const std::vector<uint8_t>& key,
  return true;
 }

+bool SessionContext::DecryptEntitlement(
+    const std::vector<uint8_t>& key,
+    const std::vector<uint8_t>& iv,
+    const std::vector<uint8_t>& message,
+    std::vector<uint8_t>* decrypted) {
+  if (key.empty() || iv.empty() || message.empty() || !decrypted) {
+    LOGE("[DecryptMessage(): OEMCrypto_ERROR_INVALID_CONTEXT]");
+    return false;
+  }
+  decrypted->resize(message.size());
+  uint8_t iv_buffer[16];
+  memcpy(iv_buffer, &iv[0], 16);
+  AES_KEY aes_key;
+  AES_set_decrypt_key(&key[0], 256, &aes_key);
+  AES_cbc_encrypt(&message[0], &(decrypted->front()), message.size(), &aes_key,
+                  iv_buffer, AES_DECRYPT);
+  return true;
+}
+
 OEMCryptoResult SessionContext::DecryptCENC(
    const uint8_t* iv, size_t block_offset,
    const OEMCrypto_CENCEncryptPatternDesc* pattern, const uint8_t* cipher_data,
--- a/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session.h
+++ b/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session.h
@@ -30,6 +30,34 @@ typedef uint32_t SessionId;

 enum SRMVersionStatus { NoSRMVersion, ValidSRMVersion, InvalidSRMVersion };

+// TODO(jfore): Is there a better name?
+class SessionContextKeys {
+ public:
+  virtual OEMCrypto_LicenseType type() = 0;
+  virtual size_t size() = 0;
+  virtual bool Insert(const KeyId& key_id, const Key& key_data) = 0;
+  virtual Key* Find(const KeyId& key_id) = 0;
+  virtual void Remove(const KeyId& key_id) = 0;
+  virtual void UpdateDuration(const KeyControlBlock& control) = 0;
+
+  // Methods supported exclusively for entitlement keys. Returns false if
+  // entitlement keys are not found or not supported by the current key table.
+  // It is the caller's responsibility to check the context.
+  virtual bool SetContentKey(const KeyId& entitlement_id,
+                             const KeyId& content_key_id,
+                             const std::vector<uint8_t>& content_key) = 0;
+  virtual bool GetEntitlementKey(const KeyId& entitlement_id,
+                                 const std::vector<uint8_t>** key) = 0;
+
+  virtual ~SessionContextKeys() {}
+
+ protected:
+  SessionContextKeys() {}
+
+ private:
+  CORE_DISALLOW_COPY_AND_ASSIGN(SessionContextKeys);
+};
+
 class SessionContext {
 private:
  SessionContext() {}
@@ -40,6 +68,7 @@ class SessionContext {
        ce_(ce),
        id_(sid),
        current_content_key_(NULL),
+        session_keys_(NULL),
        rsa_key_(rsa_key),
        allowed_schemes_(kSign_RSASSA_PSS),
        usage_entry_(NULL),
@@ -89,19 +118,22 @@ class SessionContext {
                                 size_t signature_length);
  void StartTimer();
  uint32_t CurrentTimer();  // (seconds).
-  OEMCryptoResult LoadKeys(const uint8_t* message, size_t message_length,
-                           const uint8_t* signature, size_t signature_length,
-                           const uint8_t* enc_mac_key_iv,
-                           const uint8_t* enc_mac_keys, size_t num_keys,
-                           const OEMCrypto_KeyObject* key_array,
-                           const uint8_t* pst, size_t pst_length,
-                           const uint8_t* srm_requirement);
+  OEMCryptoResult LoadKeys(
+      const uint8_t* message, size_t message_length, const uint8_t* signature,
+      size_t signature_length, const uint8_t* enc_mac_key_iv,
+      const uint8_t* enc_mac_keys, size_t num_keys,
+      const OEMCrypto_KeyObject* key_array, const uint8_t* pst,
+      size_t pst_length, const uint8_t* srm_requirement,
+      OEMCrypto_LicenseType license_type);
+  OEMCryptoResult LoadEntitledContentKeys(
+      size_t num_keys,
+      const OEMCrypto_EntitledContentKeyObject* key_array);
  OEMCryptoResult InstallKey(const KeyId& key_id,
                             const std::vector<uint8_t>& key_data,
                             const std::vector<uint8_t>& key_data_iv,
                             const std::vector<uint8_t>& key_control,
                             const std::vector<uint8_t>& key_control_iv,
-                             bool ctr_mode, bool second_license);
+                             bool second_license);
  bool InstallRSAEncryptedKey(const uint8_t* encrypted_message_key,
                              size_t encrypted_message_key_length);
  bool DecryptRSAKey(const uint8_t* enc_rsa_key, size_t enc_rsa_key_length,
@@ -115,7 +147,8 @@ class SessionContext {
  bool UpdateMacKeys(const std::vector<uint8_t>& mac_keys,
                     const std::vector<uint8_t>& iv);
  bool QueryKeyControlBlock(const KeyId& key_id, uint32_t* data);
-  OEMCryptoResult SelectContentKey(const KeyId& key_id);
+  OEMCryptoResult SelectContentKey(const KeyId& key_id,
+                                   OEMCryptoCipherMode cipher_mode);
  const Key* current_content_key(void) { return current_content_key_; }
  void set_mac_key_server(const std::vector<uint8_t>& mac_key_server) {
    mac_key_server_ = mac_key_server;
@@ -157,6 +190,10 @@ class SessionContext {
                      const std::vector<uint8_t>& iv,
                      const std::vector<uint8_t>& message,
                      std::vector<uint8_t>* decrypted);
+  bool DecryptEntitlement(const std::vector<uint8_t>& key,
+                          const std::vector<uint8_t>& iv,
+                          const std::vector<uint8_t>& message,
+                          std::vector<uint8_t>* decrypted);
  // Either verify the nonce or usage entry, as required by the key control
  // block.
  OEMCryptoResult CheckNonceOrEntry(const KeyControlBlock& key_control_block);
@@ -193,7 +230,7 @@ class SessionContext {
  std::vector<uint8_t> encryption_key_;
  std::vector<uint8_t> session_key_;
  const Key* current_content_key_;
-  SessionKeyTable session_keys_;
+  SessionContextKeys* session_keys_;
  NonceTable nonce_table_;
  RSA_shared_ptr rsa_key_;
  uint32_t allowed_schemes_;  // for RSA signatures.
--- a/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session_key_table.cpp
+++ b/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session_key_table.cpp
@@ -43,4 +43,75 @@ void SessionKeyTable::UpdateDuration(const KeyControlBlock& control) {
  }
 }

+bool EntitlementKeyTable::Insert(const KeyId key_id, const Key& key_data) {
+  // |key_id| and |key_data| are for an entitlement key. Insert a new
+  // entitlement key entry.
+  if (keys_.find(key_id) != keys_.end()) return false;
+  keys_[key_id] = new EntitlementKey(key_data);
+  // If this is a new insertion, we don't have a content key assigned yet.
+  return true;
+}
+
+Key* EntitlementKeyTable::Find(const KeyId key_id) {
+  // |key_id| refers to a content key.
+  ContentIdToEntitlementIdMap::iterator it =
+      contentid_to_entitlementid_.find(key_id);
+  if (it == contentid_to_entitlementid_.end()) {
+    return NULL;
+  }
+
+  if (keys_.find(it->second) == keys_.end()) {
+    return NULL;
+  }
+  return keys_[it->second];
+}
+
+void EntitlementKeyTable::Remove(const KeyId key_id) {
+  // |key_id| refers to a content key. No one currently calls Remove so this
+  // method is free to change if needed.
+  ContentIdToEntitlementIdMap::iterator it =
+      contentid_to_entitlementid_.find(key_id);
+  if (it == contentid_to_entitlementid_.end()) {
+    return;
+  }
+  keys_.erase(it->second);
+  contentid_to_entitlementid_.erase(key_id);
+}
+
+void EntitlementKeyTable::UpdateDuration(const KeyControlBlock& control) {
+  for (EntitlementKeyMap::iterator it = keys_.begin(); it != keys_.end();
+       ++it) {
+    it->second->UpdateDuration(control);
+  }
+}
+
+bool EntitlementKeyTable::SetContentKey(
+    const KeyId& entitlement_id, const KeyId& content_key_id,
+    const std::vector<uint8_t> content_key) {
+  EntitlementKeyMap::iterator it = keys_.find(entitlement_id);
+  if (it == keys_.end()) {
+    return false;
+  }
+  contentid_to_entitlementid_.erase(it->second->content_key_id());
+  if (!it->second->SetContentKey(content_key_id, content_key)) {
+    return false;
+  }
+  contentid_to_entitlementid_[content_key_id] = entitlement_id;
+  return true;
+}
+
+bool EntitlementKeyTable::GetEntitlementKey(
+    const KeyId& entitlement_id,
+    const std::vector<uint8_t>** entitlement_key) {
+  if (!entitlement_key) {
+    return false;
+  }
+  EntitlementKeyMap::iterator it = keys_.find(entitlement_id);
+  if (it == keys_.end()) {
+    return false;
+  }
+  *entitlement_key = &it->second->entitlement_key();
+  return true;
+}
+
 }  // namespace wvoec_mock
--- a/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session_key_table.h
+++ b/libwvdrmengine/oemcrypto/mock/src/oemcrypto_session_key_table.h
@@ -19,8 +19,9 @@ class CryptoEngine;
 class UsageTable;
 class UsageTableEntry;

-typedef std::vector<uint8_t>  KeyId;
+typedef std::vector<uint8_t> KeyId;
 typedef std::map<KeyId, Key*> KeyMap;
+typedef std::map<KeyId, EntitlementKey*> EntitlementKeyMap;

 // SessionKeyTable holds the keys for the current session
 class SessionKeyTable {
@@ -40,6 +41,28 @@ class SessionKeyTable {
  CORE_DISALLOW_COPY_AND_ASSIGN(SessionKeyTable);
 };

+class EntitlementKeyTable {
+  typedef std::map<KeyId, KeyId> ContentIdToEntitlementIdMap;
+
+ public:
+  EntitlementKeyTable() {}
+  ~EntitlementKeyTable() {}
+  bool Insert(const KeyId key_id, const Key& key_data);
+  Key* Find(const KeyId key_id);
+  void Remove(const KeyId key_id);
+  void UpdateDuration(const KeyControlBlock& control);
+  size_t size() const { return contentid_to_entitlementid_.size(); }
+  bool SetContentKey(const KeyId& entitlement_id, const KeyId& content_key_id,
+                     const std::vector<uint8_t> content_key);
+  bool GetEntitlementKey(const KeyId& entitlement_id,
+                         const std::vector<uint8_t>** entitlement_key);
+
+ private:
+  EntitlementKeyMap keys_;
+  ContentIdToEntitlementIdMap contentid_to_entitlementid_;
+  CORE_DISALLOW_COPY_AND_ASSIGN(EntitlementKeyTable);
+};
+
 }  // namespace wvoec_mock

 #endif  // MOCK_OEMCRYPTO_SESSION_KEY_TABLE_H_