From 6a8bd096f4085570da33fb528a570fb0582aef3f Mon Sep 17 00:00:00 2001
From: Srujan Gaddam <srujzs@google.com>
Date: Thu, 21 Feb 2019 16:51:38 -0800
Subject: [PATCH] Separate key session SelectKey methods

Test: CE CDM tests, Android unit/integration tests
Bug: b/124773017
Merge of http://go/wvgerrit/73144

Since OEMCrypto supports one content key per entitlement key,
LoadEntitledContentKeys must be called every time we want to select a
key. EntitlementKeySession::SelectKey calls ContentKeySession::SelectKey
after loading the keys, which caches the key id from the previous call,
and if the key id hasn't changed, doesn't call SelectKey. This caching
is fine for content keys since we don't call LoadKeys every time, but
not fine for entitled keys since we do call LoadEntitledContentKeys
every time. So, we instead cache the current content key id per
entitlement key and only call LoadEntitledContentKeys and SelectKey if
the entitled content key id changes. Furthermore, the test
HandlesKeyRotationWithOnlyOneLicenseRequest is modified to complete
multiple decrypts per key to test this behavior.

Change-Id: I9d0d94e49da0fe1965beadbddec99d8dff744d73
---
 .../core/include/entitlement_key_session.h    |  5 +-
 .../cdm/core/src/entitlement_key_session.cpp  | 46 ++++++++++++-------
 2 files changed, 33 insertions(+), 18 deletions(-)

diff --git a/libwvdrmengine/cdm/core/include/entitlement_key_session.h b/libwvdrmengine/cdm/core/include/entitlement_key_session.h
index ff065fde..2c2f9b8c 100644
--- a/libwvdrmengine/cdm/core/include/entitlement_key_session.h
+++ b/libwvdrmengine/cdm/core/include/entitlement_key_session.h
@@ -44,7 +44,10 @@ class EntitlementKeySession : public ContentKeySession {
   OEMCrypto_EntitledContentKeyObject MakeOecEntitledKey(
       const CryptoKey& input_key, std::string& message);
 
-  std::map<std::string, CryptoKey> entitled_keys_;
+  // Find the CryptoKey for the given entitled content key id.
+  std::map<KeyId, CryptoKey> entitled_keys_;
+  // Find the current entitled content key id for the given entitlement key id.
+  std::map<KeyId, KeyId> current_loaded_content_keys_;
 };
 
 }  // namespace wvcdm
diff --git a/libwvdrmengine/cdm/core/src/entitlement_key_session.cpp b/libwvdrmengine/cdm/core/src/entitlement_key_session.cpp
index cd8650ae..2b4ec9ad 100644
--- a/libwvdrmengine/cdm/core/src/entitlement_key_session.cpp
+++ b/libwvdrmengine/cdm/core/src/entitlement_key_session.cpp
@@ -5,8 +5,10 @@
 #include "entitlement_key_session.h"
 
 #include "crypto_key.h"
+#include "crypto_session.h"
 #include "log.h"
 
+
 namespace wvcdm {
 EntitlementKeySession::EntitlementKeySession(CryptoSessionId oec_session_id,
                                              metrics::CryptoMetrics* metrics)
@@ -43,30 +45,40 @@ OEMCryptoResult EntitlementKeySession::LoadEntitledContentKeys(
 
 OEMCryptoResult EntitlementKeySession::SelectKey(const std::string& key_id,
                                                  CdmCipherMode cipher_mode) {
-  // Before the key can be selected, it must be loaded under its associated
-  // entitlement key. This could, in theory, be done ahead of time during
-  // LoadEntitledContentKeys(), but OEMCrypto v14 only supports one content key
-  // per entitlement key at a time, (b/110266851) so we must swap out for the
-  // correct key every time SelectKey() is called.
   if (entitled_keys_.find(key_id) == entitled_keys_.end()) {
     LOGE("Unknown entitled key ID selected.");
     return OEMCrypto_ERROR_NO_CONTENT_KEY;
   }
 
-  std::string message;
-  OEMCrypto_EntitledContentKeyObject entitled_key =
-      MakeOecEntitledKey(entitled_keys_[key_id], message);
+  CryptoKey entitled_content_key = entitled_keys_[key_id];
+  if (current_loaded_content_keys_[entitled_content_key
+                                        .entitlement_key_id()] != key_id) {
+    // Before the key can be selected, it must be loaded under its associated
+    // entitlement key. This could, in theory, be done ahead of time during
+    // LoadEntitledContentKeys(), but OEMCrypto v14 only supports one content
+    // key per entitlement key at a time, (b/110266851) so we must swap out for
+    // the correct key every time the current content key id has changed for an
+    // entitlement key.
+    std::string message;
+    OEMCrypto_EntitledContentKeyObject entitled_key =
+        MakeOecEntitledKey(entitled_content_key, message);
+    OEMCryptoResult result = OEMCrypto_SUCCESS;
+    M_TIME(
+        result = OEMCrypto_LoadEntitledContentKeys(
+            oec_session_id_, reinterpret_cast<const uint8_t*>(message.data()),
+            message.size(), 1, &entitled_key),
+        metrics_, oemcrypto_load_entitled_keys_, result);
+    if (result != OEMCrypto_SUCCESS) {
+      return result;
+    }
 
-  OEMCryptoResult result = OEMCrypto_SUCCESS;
-  M_TIME(result = OEMCrypto_LoadEntitledContentKeys(
-             oec_session_id_, reinterpret_cast<const uint8_t*>(message.data()),
-             message.size(), 1, &entitled_key),
-         metrics_, oemcrypto_load_entitled_keys_, result);
-  if (result != OEMCrypto_SUCCESS) {
-    return result;
+    current_loaded_content_keys_[entitled_content_key.entitlement_key_id()] =
+        key_id;
+
+    return ContentKeySession::SelectKey(key_id, cipher_mode);
+  } else {
+    return OEMCrypto_SUCCESS;
   }
-
-  return ContentKeySession::SelectKey(key_id, cipher_mode);
 }
 
 OEMCrypto_EntitledContentKeyObject EntitlementKeySession::MakeOecEntitledKey(