Fixed race condition in closeSession

Merged from http://go/wvgerrit/165061
poc: http://go/ag/20978750

Fix race that corrupts mCryptoSessions std::map,
and race that occurs when CryptoSessions are used after free.

Test: poc
Test: atest MediaDrmParameterizedTests
Test: atest GtsMediaTestCases

Bug: 258189255
Change-Id: I298d3e0770ace9cd590dfaacaa4c52a0732c2fe3
Merged-In: I298d3e0770ace9cd590dfaacaa4c52a0732c2fe3

libwvdrmengine/mediadrm/aidl_include/WVDrmPlugin.h

@@ -14,6 +14,7 @@
 
 #include <list>
 #include <map>
+#include <mutex>
 
 #include "OEMCryptoCENC.h"
 #include "WVGenericCryptoInterface.h"
@@ -395,9 +396,46 @@ struct WVDrmPlugin : public ::aidl::android::hardware::drm::BnDrmPlugin,
     android::Mutex mLock;
   };
 
+  class CryptoSessionMap {
+   public:
+    std::map<CdmSessionId, std::shared_ptr<CryptoSession>> clear() {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      auto copy = mMap;
+      mMap.clear();
+      return copy;
+    }
+
+    std::shared_ptr<CryptoSession> get(const CdmSessionId& sid) {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      if (mMap.count(sid)) {
+        return mMap[sid];
+      }
+      return nullptr;
+    }
+
+    bool empty() {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      return mMap.empty();
+    }
+
+    void erase(const CdmSessionId& sid) {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      mMap.erase(sid);
+    }
+
+    void insert(const CdmSessionId& sid, OEMCrypto_SESSION osid) {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      mMap[sid] = std::make_shared<CryptoSession>(osid);
+    }
+
+   private:
+     std::mutex mLock;
+     std::map<CdmSessionId, std::shared_ptr<CryptoSession>> mMap;
+  };
+
   android::sp<wvcdm::WvContentDecryptionModule> const mCDM;
   WVGenericCryptoInterface* mCrypto;
-  std::map<CdmSessionId, CryptoSession> mCryptoSessions;
+  CryptoSessionMap mCryptoSessions;
   std::shared_ptr<::aidl::android::hardware::drm::IDrmPluginListener> mListener;
 
   std::string mProvisioningServiceCertificate;

libwvdrmengine/mediadrm/aidl_src/WVDrmPlugin.cpp

@@ -47,6 +47,7 @@ namespace hardware {
 namespace drm {
 namespace widevine {
 
+using std::shared_ptr;
 using std::string;
 using std::vector;
 
@@ -170,15 +171,18 @@ WVDrmPlugin::~WVDrmPlugin() {
 }
 
 void WVDrmPlugin::Close() {
-  typedef std::map<CdmSessionId, CryptoSession>::iterator mapIterator;
+  typedef std::map<CdmSessionId, shared_ptr<CryptoSession>>::iterator mapIterator;
-  for (mapIterator iter = mCryptoSessions.begin();
+  auto cryptoSessions = mCryptoSessions.clear();
-       iter != mCryptoSessions.end(); ++iter) {
+  for (mapIterator iter = cryptoSessions.begin();
+       iter != cryptoSessions.end(); ++iter) {
     CdmResponseType res = mCDM->CloseSession(iter->first);
     if (!isCdmResponseTypeSuccess(res)) {
       ALOGE("Failed to close session while destroying WVDrmPlugin");
     }
   }
-  mCryptoSessions.clear();
+  // clear local copy of cryptoSessions map
+  cryptoSessions.clear();
+
   if (mCdmIdentifierBuilder.is_sealed()) {
     CdmIdentifier identifier;
     Status status = mCdmIdentifierBuilder.getCdmIdentifier(&identifier);
@@ -221,7 +225,7 @@ Status WVDrmPlugin::openSessionCommon(vector<uint8_t>& sessionId) {
       info.count(wvcdm::QUERY_KEY_OEMCRYPTO_SESSION_ID)) {
     OEMCrypto_SESSION oecSessionId =
         std::stoul(info[wvcdm::QUERY_KEY_OEMCRYPTO_SESSION_ID]);
-    mCryptoSessions[cdmSessionId] = CryptoSession(oecSessionId);
+    mCryptoSessions.insert(cdmSessionId, oecSessionId);
     success = true;
   } else {
     ALOGE("Unable to query key control info.");
@@ -1225,7 +1229,7 @@ Status WVDrmPlugin::unprovisionDevice() {
   std::string _value(in_value.c_str());
 
   if (name == "securityLevel") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       if (_value == wvcdm::QUERY_VALUE_SECURITY_LEVEL_L3.c_str()) {
         mPropertySet.set_security_level(wvcdm::QUERY_VALUE_SECURITY_LEVEL_L3);
       } else if (_value == wvcdm::QUERY_VALUE_SECURITY_LEVEL_L1.c_str()) {
@@ -1265,7 +1269,7 @@ Status WVDrmPlugin::unprovisionDevice() {
       return toNdkScopedAStatus(Status::BAD_VALUE);
     }
   } else if (name == "sessionSharing") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       if (_value == kEnable) {
         mPropertySet.set_is_session_sharing_enabled(true);
       } else if (_value == kDisable) {
@@ -1280,7 +1284,7 @@ Status WVDrmPlugin::unprovisionDevice() {
       return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
     }
   } else if (name == "appId") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       mPropertySet.set_app_id(_value.c_str());
     } else {
       ALOGE("App tried to set the application id while sessions are opened.");
@@ -1288,7 +1292,7 @@ Status WVDrmPlugin::unprovisionDevice() {
       return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
     }
   } else if (name == "origin") {
-    if (mCryptoSessions.size() != 0) {
+    if (!mCryptoSessions.empty()) {
       ALOGE("App tried to set the origin while sessions are opened.");
       ALOGW("Returns UNKNOWN error for legacy status kErrorSessionIsOpen");
       return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
@@ -1389,14 +1393,13 @@ Status WVDrmPlugin::unprovisionDevice() {
   std::string algo(in_algorithm.c_str());
 
   CdmSessionId cdmSessionId(in_sessionId.begin(), in_sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return toNdkScopedAStatus(Status::ERROR_DRM_SESSION_NOT_OPENED);
   }
 
-  CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
   if (algo == "AES/CBC/NoPadding") {
-    cryptoSession.setCipherAlgorithm(OEMCrypto_AES_CBC_128_NO_PADDING);
+    cryptoSession->setCipherAlgorithm(OEMCrypto_AES_CBC_128_NO_PADDING);
   } else {
     return toNdkScopedAStatus(Status::ERROR_DRM_CANNOT_HANDLE);
   }
@@ -1412,14 +1415,13 @@ Status WVDrmPlugin::unprovisionDevice() {
   std::string algo(in_algorithm.c_str());
 
   CdmSessionId cdmSessionId(in_sessionId.begin(), in_sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return toNdkScopedAStatus(Status::ERROR_DRM_SESSION_NOT_OPENED);
   }
 
-  CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
   if (algo == "HmacSHA256") {
-    cryptoSession.setMacAlgorithm(OEMCrypto_HMAC_SHA256);
+    cryptoSession->setMacAlgorithm(OEMCrypto_HMAC_SHA256);
   } else {
     return toNdkScopedAStatus(Status::ERROR_DRM_CANNOT_HANDLE);
   }
@@ -1436,18 +1438,17 @@ Status WVDrmPlugin::unprovisionDevice() {
 
   *_aidl_return = output;
   CdmSessionId cdmSessionId(in_sessionId.begin(), in_sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return toNdkScopedAStatus(Status::ERROR_DRM_SESSION_NOT_OPENED);
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->cipherAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.cipherAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            in_keyId.data(), in_keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1458,9 +1459,9 @@ Status WVDrmPlugin::unprovisionDevice() {
   output.resize(in_input.size());
 
   Status status = Status::OK;
-  res = mCrypto->encrypt(cryptoSession.oecSessionId(), in_input.data(),
+  res = mCrypto->encrypt(cryptoSession->oecSessionId(), in_input.data(),
                          in_input.size(), in_iv.data(),
-                         cryptoSession.cipherAlgorithm(), output.data());
+                         cryptoSession->cipherAlgorithm(), output.data());
 
   *_aidl_return = output;
   if (res == OEMCrypto_SUCCESS) {
@@ -1481,18 +1482,17 @@ Status WVDrmPlugin::unprovisionDevice() {
   *_aidl_return = output;
 
   CdmSessionId cdmSessionId(in_sessionId.begin(), in_sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return toNdkScopedAStatus(Status::ERROR_DRM_SESSION_NOT_OPENED);
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->cipherAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.cipherAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            in_keyId.data(), in_keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1503,9 +1503,9 @@ Status WVDrmPlugin::unprovisionDevice() {
   output.resize(in_input.size());
 
   Status status = Status::OK;
-  res = mCrypto->decrypt(cryptoSession.oecSessionId(), in_input.data(),
+  res = mCrypto->decrypt(cryptoSession->oecSessionId(), in_input.data(),
                          in_input.size(), in_iv.data(),
-                         cryptoSession.cipherAlgorithm(), output.data());
+                         cryptoSession->cipherAlgorithm(), output.data());
 
   *_aidl_return = output;
   if (res == OEMCrypto_SUCCESS) {
@@ -1525,18 +1525,17 @@ Status WVDrmPlugin::unprovisionDevice() {
   *_aidl_return = signature;
 
   CdmSessionId cdmSessionId(in_sessionId.begin(), in_sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return toNdkScopedAStatus(Status::ERROR_DRM_SESSION_NOT_OPENED);
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->macAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.macAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            in_keyId.data(), in_keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1546,8 +1545,8 @@ Status WVDrmPlugin::unprovisionDevice() {
 
   size_t signatureSize = 0;
 
-  res = mCrypto->sign(cryptoSession.oecSessionId(), in_message.data(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), in_message.data(),
-                      in_message.size(), cryptoSession.macAlgorithm(), nullptr,
+                      in_message.size(), cryptoSession->macAlgorithm(), nullptr,
                       &signatureSize);
 
   Status status = Status::OK;
@@ -1566,8 +1565,8 @@ Status WVDrmPlugin::unprovisionDevice() {
 
   signature.resize(signatureSize);
 
-  res = mCrypto->sign(cryptoSession.oecSessionId(), in_message.data(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), in_message.data(),
-                      in_message.size(), cryptoSession.macAlgorithm(),
+                      in_message.size(), cryptoSession->macAlgorithm(),
                       signature.data(), &signatureSize);
 
   *_aidl_return = signature;
@@ -1589,18 +1588,17 @@ Status WVDrmPlugin::unprovisionDevice() {
   *_aidl_return = match;
 
   CdmSessionId cdmSessionId(in_sessionId.begin(), in_sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return toNdkScopedAStatus(Status::ERROR_DRM_SESSION_NOT_OPENED);
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->macAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.macAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     return toNdkScopedAStatus(Status::ERROR_DRM_UNKNOWN);
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            in_keyId.data(), in_keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1608,8 +1606,8 @@ Status WVDrmPlugin::unprovisionDevice() {
     return toNdkScopedAStatus(mapAndNotifyOfOEMCryptoResult(in_sessionId, res));
   }
 
-  res = mCrypto->verify(cryptoSession.oecSessionId(), in_message.data(),
+  res = mCrypto->verify(cryptoSession->oecSessionId(), in_message.data(),
-                        in_message.size(), cryptoSession.macAlgorithm(),
+                        in_message.size(), cryptoSession->macAlgorithm(),
                         in_signature.data(), in_signature.size());
 
   Status status = Status::OK;
@@ -1678,7 +1676,7 @@ Status WVDrmPlugin::unprovisionDevice() {
 }
 
 ::ndk::ScopedAStatus WVDrmPlugin::setListener(
-    const std::shared_ptr<::aidl::android::hardware::drm::IDrmPluginListener>&
+    const shared_ptr<::aidl::android::hardware::drm::IDrmPluginListener>&
         in_listener) {
   mListener = in_listener;
   ::ndk::ScopedAStatus _aidl_status;

libwvdrmengine/mediadrm/include_hidl/WVDrmPlugin.h

@@ -11,6 +11,7 @@
 
 #include <list>
 #include <map>
+#include <mutex>
 
 #include "HidlTypes.h"
 #include "OEMCryptoCENC.h"
@@ -450,9 +451,46 @@ struct WVDrmPlugin : public ::drm::V1_4::IDrmPlugin,
     Mutex mLock;
   };
 
+  class CryptoSessionMap {
+   public:
+    std::map<CdmSessionId, std::shared_ptr<CryptoSession>> clear() {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      auto copy = mMap;
+      mMap.clear();
+      return copy;
+    }
+
+    std::shared_ptr<CryptoSession> get(const CdmSessionId& sid) {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      if (mMap.count(sid)) {
+        return mMap[sid];
+      }
+      return nullptr;
+    }
+
+    bool empty() {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      return mMap.empty();
+    }
+
+    void erase(const CdmSessionId& sid) {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      mMap.erase(sid);
+    }
+
+    void insert(const CdmSessionId& sid, OEMCrypto_SESSION osid) {
+      std::unique_lock<std::mutex> auto_lock(mLock);
+      mMap[sid] = std::make_shared<CryptoSession>(osid);
+    }
+
+   private:
+     std::mutex mLock;
+     std::map<CdmSessionId, std::shared_ptr<CryptoSession>> mMap;
+  };
+
   sp<wvcdm::WvContentDecryptionModule> const mCDM;
   WVGenericCryptoInterface* mCrypto;
-  map<CdmSessionId, CryptoSession> mCryptoSessions;
+  CryptoSessionMap mCryptoSessions;
   sp<IDrmPluginListener> mListener;
   sp<IDrmPluginListener_V1_2> mListenerV1_2;
 

libwvdrmengine/mediadrm/src_hidl/WVDrmPlugin.cpp

@@ -46,6 +46,8 @@ namespace drm {
 namespace V1_4 {
 namespace widevine {
 
+using std::shared_ptr;
+
 using android::hardware::drm::V1_2::widevine::toHidlVec;
 using android::hardware::drm::V1_2::widevine::toVector;
 using wvcdm::CdmAppParameterMap;
@@ -208,15 +210,18 @@ WVDrmPlugin::~WVDrmPlugin() {
 }
 
 void WVDrmPlugin::Close() {
-  typedef map<CdmSessionId, CryptoSession>::iterator mapIterator;
+  typedef map<CdmSessionId, shared_ptr<CryptoSession>>::iterator mapIterator;
-  for (mapIterator iter = mCryptoSessions.begin();
+  auto cryptoSessions = mCryptoSessions.clear();
-       iter != mCryptoSessions.end(); ++iter) {
+  for (mapIterator iter = cryptoSessions.begin();
+       iter != cryptoSessions.end(); ++iter) {
     CdmResponseType res = mCDM->CloseSession(iter->first);
     if (!isCdmResponseTypeSuccess(res)) {
       ALOGE("Failed to close session while destroying WVDrmPlugin");
     }
   }
-  mCryptoSessions.clear();
+  // clear local copy of cryptoSessions map
+  cryptoSessions.clear();
+
   if (mCdmIdentifierBuilder.is_sealed()) {
     CdmIdentifier identifier;
     Status status = mCdmIdentifierBuilder.getCdmIdentifier(&identifier);
@@ -259,7 +264,7 @@ Status WVDrmPlugin::openSessionCommon(std::vector<uint8_t>& sessionId) {
       info.count(wvcdm::QUERY_KEY_OEMCRYPTO_SESSION_ID)) {
     OEMCrypto_SESSION oecSessionId =
         std::stoul(info[wvcdm::QUERY_KEY_OEMCRYPTO_SESSION_ID]);
-    mCryptoSessions[cdmSessionId] = CryptoSession(oecSessionId);
+    mCryptoSessions.insert(cdmSessionId, oecSessionId);
     success = true;
   } else {
     ALOGE("Unable to query key control info.");
@@ -1333,7 +1338,7 @@ Return<Status> WVDrmPlugin::setPropertyString(const hidl_string& propertyName,
   std::string _value(value.c_str());
 
   if (name == "securityLevel") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       if (_value == wvcdm::QUERY_VALUE_SECURITY_LEVEL_L3.c_str()) {
         mPropertySet.set_security_level(wvcdm::QUERY_VALUE_SECURITY_LEVEL_L3);
       } else if (_value == wvcdm::QUERY_VALUE_SECURITY_LEVEL_L1.c_str()) {
@@ -1373,7 +1378,7 @@ Return<Status> WVDrmPlugin::setPropertyString(const hidl_string& propertyName,
       return Status::BAD_VALUE;
     }
   } else if (name == "sessionSharing") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       if (_value == kEnable) {
         mPropertySet.set_is_session_sharing_enabled(true);
       } else if (_value == kDisable) {
@@ -1388,7 +1393,7 @@ Return<Status> WVDrmPlugin::setPropertyString(const hidl_string& propertyName,
       return Status::ERROR_DRM_UNKNOWN;
     }
   } else if (name == "appId") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       mPropertySet.set_app_id(_value.c_str());
     } else {
       ALOGE("App tried to set the application id while sessions are opened.");
@@ -1396,7 +1401,7 @@ Return<Status> WVDrmPlugin::setPropertyString(const hidl_string& propertyName,
       return Status::ERROR_DRM_UNKNOWN;
     }
   } else if (name == "origin") {
-    if (mCryptoSessions.size() != 0) {
+    if (!mCryptoSessions.empty()) {
       ALOGE("App tried to set the origin while sessions are opened.");
       ALOGW("Returns UNKNOWN error for legacy status kErrorSessionIsOpen");
       return Status::ERROR_DRM_UNKNOWN;
@@ -1511,14 +1516,13 @@ Return<Status> WVDrmPlugin::setCipherAlgorithm(
   std::vector<uint8_t> sId = toVector(sessionId);
 
   CdmSessionId cdmSessionId(sId.begin(), sId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return Status::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
   if (algo == "AES/CBC/NoPadding") {
-    cryptoSession.setCipherAlgorithm(OEMCrypto_AES_CBC_128_NO_PADDING);
+    cryptoSession->setCipherAlgorithm(OEMCrypto_AES_CBC_128_NO_PADDING);
   } else {
     return Status::ERROR_DRM_CANNOT_HANDLE;
   }
@@ -1535,14 +1539,13 @@ Return<Status> WVDrmPlugin::setMacAlgorithm(const hidl_vec<uint8_t>& sessionId,
   std::vector<uint8_t> sId = toVector(sessionId);
 
   CdmSessionId cdmSessionId(sId.begin(), sId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return Status::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
   if (algo == "HmacSHA256") {
-    cryptoSession.setMacAlgorithm(OEMCrypto_HMAC_SHA256);
+    cryptoSession->setMacAlgorithm(OEMCrypto_HMAC_SHA256);
   } else {
     return Status::ERROR_DRM_CANNOT_HANDLE;
   }
@@ -1559,21 +1562,20 @@ Return<void> WVDrmPlugin::encrypt(const hidl_vec<uint8_t>& sessionId,
   std::vector<uint8_t> output;
 
   CdmSessionId cdmSessionId(sId.begin(), sId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     _hidl_cb(Status::ERROR_DRM_SESSION_NOT_OPENED, toHidlVec(output));
     return Void();
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->cipherAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.cipherAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     _hidl_cb(Status::ERROR_DRM_UNKNOWN, toHidlVec(output));
     return Void();
   }
 
   const std::vector<uint8_t> _keyId = toVector(keyId);
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            _keyId.data(), _keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1586,9 +1588,9 @@ Return<void> WVDrmPlugin::encrypt(const hidl_vec<uint8_t>& sessionId,
   const std::vector<uint8_t> _iv = toVector(iv);
   output.resize(_input.size());
 
-  res = mCrypto->encrypt(cryptoSession.oecSessionId(), _input.data(),
+  res = mCrypto->encrypt(cryptoSession->oecSessionId(), _input.data(),
                          _input.size(), _iv.data(),
-                         cryptoSession.cipherAlgorithm(), output.data());
+                         cryptoSession->cipherAlgorithm(), output.data());
 
   if (res == OEMCrypto_SUCCESS) {
     _hidl_cb(Status::OK, toHidlVec(output));
@@ -1608,21 +1610,20 @@ Return<void> WVDrmPlugin::decrypt(const hidl_vec<uint8_t>& sessionId,
   std::vector<uint8_t> output;
 
   CdmSessionId cdmSessionId(sId.begin(), sId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     _hidl_cb(Status::ERROR_DRM_SESSION_NOT_OPENED, toHidlVec(output));
     return Void();
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->cipherAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.cipherAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     _hidl_cb(Status::ERROR_DRM_UNKNOWN, toHidlVec(output));
     return Void();
   }
 
   const std::vector<uint8_t> _keyId = toVector(keyId);
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            _keyId.data(), _keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1635,9 +1636,9 @@ Return<void> WVDrmPlugin::decrypt(const hidl_vec<uint8_t>& sessionId,
   const std::vector<uint8_t> _iv = toVector(iv);
   output.resize(_input.size());
 
-  res = mCrypto->decrypt(cryptoSession.oecSessionId(), _input.data(),
+  res = mCrypto->decrypt(cryptoSession->oecSessionId(), _input.data(),
                          _input.size(), _iv.data(),
-                         cryptoSession.cipherAlgorithm(), output.data());
+                         cryptoSession->cipherAlgorithm(), output.data());
 
   if (res == OEMCrypto_SUCCESS) {
     _hidl_cb(Status::OK, toHidlVec(output));
@@ -1656,21 +1657,20 @@ Return<void> WVDrmPlugin::sign(const hidl_vec<uint8_t>& sessionId,
   std::vector<uint8_t> signature;
 
   CdmSessionId cdmSessionId(sId.begin(), sId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     _hidl_cb(Status::ERROR_DRM_SESSION_NOT_OPENED, toHidlVec(signature));
     return Void();
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->macAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.macAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     _hidl_cb(Status::ERROR_DRM_UNKNOWN, toHidlVec(signature));
     return Void();
   }
 
   const std::vector<uint8_t> _keyId = toVector(keyId);
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            _keyId.data(), _keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1682,8 +1682,8 @@ Return<void> WVDrmPlugin::sign(const hidl_vec<uint8_t>& sessionId,
   size_t signatureSize = 0;
 
   const std::vector<uint8_t> msg = toVector(message);
-  res = mCrypto->sign(cryptoSession.oecSessionId(), msg.data(), msg.size(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), msg.data(), msg.size(),
-                      cryptoSession.macAlgorithm(), NULL, &signatureSize);
+                      cryptoSession->macAlgorithm(), NULL, &signatureSize);
 
   if (res != OEMCrypto_ERROR_SHORT_BUFFER) {
     ALOGE(
@@ -1700,8 +1700,8 @@ Return<void> WVDrmPlugin::sign(const hidl_vec<uint8_t>& sessionId,
 
   signature.resize(signatureSize);
 
-  res = mCrypto->sign(cryptoSession.oecSessionId(), msg.data(), msg.size(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), msg.data(), msg.size(),
-                      cryptoSession.macAlgorithm(), signature.data(),
+                      cryptoSession->macAlgorithm(), signature.data(),
                       &signatureSize);
 
   if (res == OEMCrypto_SUCCESS) {
@@ -1722,21 +1722,20 @@ Return<void> WVDrmPlugin::verify(const hidl_vec<uint8_t>& sessionId,
   const std::vector<uint8_t> sId = toVector(sessionId);
 
   CdmSessionId cdmSessionId(sId.begin(), sId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     _hidl_cb(Status::ERROR_DRM_SESSION_NOT_OPENED, match);
     return Void();
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
+  if (cryptoSession->macAlgorithm() == kInvalidCryptoAlgorithm) {
-
-  if (cryptoSession.macAlgorithm() == kInvalidCryptoAlgorithm) {
     ALOGW("Returns UNKNOWN error for legacy status NO_INIT");
     _hidl_cb(Status::ERROR_DRM_UNKNOWN, match);
     return Void();
   }
 
   const std::vector<uint8_t> _keyId = toVector(keyId);
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            _keyId.data(), _keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -1747,8 +1746,8 @@ Return<void> WVDrmPlugin::verify(const hidl_vec<uint8_t>& sessionId,
 
   const std::vector<uint8_t> _message = toVector(message);
   const std::vector<uint8_t> _signature = toVector(signature);
-  res = mCrypto->verify(cryptoSession.oecSessionId(), _message.data(),
+  res = mCrypto->verify(cryptoSession->oecSessionId(), _message.data(),
-                        _message.size(), cryptoSession.macAlgorithm(),
+                        _message.size(), cryptoSession->macAlgorithm(),
                         _signature.data(), _signature.size());
 
   if (res == OEMCrypto_SUCCESS) {