From 08acec9c3a4d2ba722175483fc07399df11e62ae Mon Sep 17 00:00:00 2001
From: Edwin Wong <edwinwong@google.com>
Date: Wed, 4 Jan 2023 01:56:05 +0000
Subject: [PATCH] Fixed race condition in closeSession

Merged from http://go/wvgerrit/165498
poc: http://go/ag/20978761

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/include/WVDrmPlugin.h | 48 +++++++++-
 .../mediadrm/include_hidl/WVDrmPlugin.h       | 40 +++++++-
 libwvdrmengine/mediadrm/src/WVDrmPlugin.cpp   | 93 +++++++++---------
 .../mediadrm/src_hidl/WVDrmPlugin.cpp         | 95 +++++++++----------
 4 files changed, 173 insertions(+), 103 deletions(-)

diff --git a/libwvdrmengine/mediadrm/include/WVDrmPlugin.h b/libwvdrmengine/mediadrm/include/WVDrmPlugin.h
index f61d5e37..a24a5c80 100644
--- a/libwvdrmengine/mediadrm/include/WVDrmPlugin.h
+++ b/libwvdrmengine/mediadrm/include/WVDrmPlugin.h
@@ -7,15 +7,17 @@
 #ifndef WV_DRM_PLUGIN_H_
 #define WV_DRM_PLUGIN_H_
 
-#include <stdint.h>
 #include <map>
+#include <mutex>
+#include <stdint.h>
 
+#include "OEMCryptoCENC.h"
+#include "WVGenericCryptoInterface.h"
 #include "cdm_client_property_set.h"
 #include "cdm_identifier.h"
 #include "media/drm/DrmAPI.h"
 #include "media/stagefright/foundation/ABase.h"
 #include "media/stagefright/foundation/AString.h"
-#include "OEMCryptoCENC.h"
 #include "utils/Errors.h"
 #include "utils/KeyedVector.h"
 #include "utils/List.h"
@@ -24,7 +26,6 @@
 #include "utils/Vector.h"
 #include "wv_cdm_event_listener.h"
 #include "wv_content_decryption_module.h"
-#include "WVGenericCryptoInterface.h"
 
 namespace wvdrm {
 
@@ -36,8 +37,8 @@ using android::Vector;
 using std::map;
 using wvcdm::CdmIdentifier;
 using wvcdm::CdmKeyStatusMap;
-using wvcdm::CdmSessionId;
 using wvcdm::CdmResponseType;
+using wvcdm::CdmSessionId;
 using wvcdm::WvContentDecryptionModule;
 
 const OEMCrypto_Algorithm kInvalidCryptoAlgorithm =
@@ -265,9 +266,46 @@ class WVDrmPlugin : public android::DrmPlugin,
     const std::string mEmptyString;
   } mPropertySet;
 
+  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;
-  map<CdmSessionId, CryptoSession> mCryptoSessions;
+  CryptoSessionMap mCryptoSessions;
 
   CdmIdentifier mCdmIdentifier;
 
diff --git a/libwvdrmengine/mediadrm/include_hidl/WVDrmPlugin.h b/libwvdrmengine/mediadrm/include_hidl/WVDrmPlugin.h
index fb22e34b..1a45cec6 100644
--- a/libwvdrmengine/mediadrm/include_hidl/WVDrmPlugin.h
+++ b/libwvdrmengine/mediadrm/include_hidl/WVDrmPlugin.h
@@ -8,6 +8,7 @@
 #define WV_DRM_PLUGIN_H_
 
 #include <map>
+#include <mutex>
 
 #include "cdm_client_property_set.h"
 #include "cdm_identifier.h"
@@ -450,9 +451,46 @@ struct WVDrmPlugin : public ::drm::V1_4::IDrmPlugin, IDrmPluginListener,
     uint32_t getNextUniqueId();
   } mCdmIdentifierBuilder;
 
+  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;
 
diff --git a/libwvdrmengine/mediadrm/src/WVDrmPlugin.cpp b/libwvdrmengine/mediadrm/src/WVDrmPlugin.cpp
index 71ec8a31..a5d41337 100644
--- a/libwvdrmengine/mediadrm/src/WVDrmPlugin.cpp
+++ b/libwvdrmengine/mediadrm/src/WVDrmPlugin.cpp
@@ -31,6 +31,8 @@ namespace {
 
 namespace wvdrm {
 
+using std::shared_ptr;
+
 using namespace android;
 using namespace std;
 using namespace wvcdm;
@@ -85,16 +87,17 @@ WVDrmPlugin::WVDrmPlugin(const sp<WvContentDecryptionModule>& cdm,
     mCdmIdentifier(kDefaultCdmIdentifier) {}
 
 WVDrmPlugin::~WVDrmPlugin() {
-  typedef map<CdmSessionId, CryptoSession>::iterator mapIterator;
-  for (mapIterator iter = mCryptoSessions.begin();
-       iter != mCryptoSessions.end();
-       ++iter) {
+  typedef map<CdmSessionId, shared_ptr<CryptoSession>>::iterator mapIterator;
+  auto cryptoSessions = mCryptoSessions.clear();
+  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();
 }
 
 status_t WVDrmPlugin::openSession(Vector<uint8_t>& sessionId) {
@@ -117,7 +120,7 @@ status_t WVDrmPlugin::openSession(Vector<uint8_t>& sessionId) {
       info.count(QUERY_KEY_OEMCRYPTO_SESSION_ID)) {
     OEMCrypto_SESSION oecSessionId =
         std::stoul(info[QUERY_KEY_OEMCRYPTO_SESSION_ID]);
-    mCryptoSessions[cdmSessionId] = CryptoSession(oecSessionId);
+    mCryptoSessions.insert(cdmSessionId, oecSessionId);
     success = true;
   } else {
     ALOGE("Unable to query key control info.");
@@ -603,7 +606,7 @@ status_t WVDrmPlugin::getPropertyByteArray(const String8& name,
 status_t WVDrmPlugin::setPropertyString(const String8& name,
                                         const String8& value) {
   if (name == "securityLevel") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       if (value == QUERY_VALUE_SECURITY_LEVEL_L3.c_str()) {
         mPropertySet.set_security_level(QUERY_VALUE_SECURITY_LEVEL_L3);
       } else if (value == QUERY_VALUE_SECURITY_LEVEL_L1.c_str()) {
@@ -640,7 +643,7 @@ status_t WVDrmPlugin::setPropertyString(const String8& name,
       return android::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) {
@@ -654,14 +657,14 @@ status_t WVDrmPlugin::setPropertyString(const String8& name,
       return kErrorSessionIsOpen;
     }
   } else if (name == "appId") {
-    if (mCryptoSessions.size() == 0) {
+    if (mCryptoSessions.empty()) {
       mPropertySet.set_app_id(value.string());
     } else {
       ALOGE("App tried to set the application id while sessions are opened.");
       return kErrorSessionIsOpen;
     }
   } else if (name == "origin") {
-    if (mCryptoSessions.size() != 0) {
+    if (!mCryptoSessions.empty()) {
       ALOGE("App tried to set the origin while sessions are opened.");
       return kErrorSessionIsOpen;
     } else {
@@ -724,14 +727,13 @@ status_t WVDrmPlugin::setCipherAlgorithm(const Vector<uint8_t>& sessionId,
     return android::BAD_VALUE;
   }
   CdmSessionId cdmSessionId(sessionId.begin(), sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return android::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
   if (algorithm == "AES/CBC/NoPadding") {
-    cryptoSession.setCipherAlgorithm(OEMCrypto_AES_CBC_128_NO_PADDING);
+    cryptoSession->setCipherAlgorithm(OEMCrypto_AES_CBC_128_NO_PADDING);
   } else {
     return android::ERROR_DRM_CANNOT_HANDLE;
   }
@@ -745,14 +747,13 @@ status_t WVDrmPlugin::setMacAlgorithm(const Vector<uint8_t>& sessionId,
     return android::BAD_VALUE;
   }
   CdmSessionId cdmSessionId(sessionId.begin(), sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return android::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
   if (algorithm == "HmacSHA256") {
-    cryptoSession.setMacAlgorithm(OEMCrypto_HMAC_SHA256);
+    cryptoSession->setMacAlgorithm(OEMCrypto_HMAC_SHA256);
   } else {
     return android::ERROR_DRM_CANNOT_HANDLE;
   }
@@ -766,17 +767,16 @@ status_t WVDrmPlugin::encrypt(const Vector<uint8_t>& sessionId,
                               const Vector<uint8_t>& iv,
                               Vector<uint8_t>& output) {
   CdmSessionId cdmSessionId(sessionId.begin(), sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return android::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
-  if (cryptoSession.cipherAlgorithm() == kInvalidCryptoAlgorithm) {
+  if (cryptoSession->cipherAlgorithm() == kInvalidCryptoAlgorithm) {
     return android::NO_INIT;
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            keyId.array(), keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -786,9 +786,9 @@ status_t WVDrmPlugin::encrypt(const Vector<uint8_t>& sessionId,
 
   output.resize(input.size());
 
-  res = mCrypto->encrypt(cryptoSession.oecSessionId(), input.array(),
+  res = mCrypto->encrypt(cryptoSession->oecSessionId(), input.array(),
                          input.size(), iv.array(),
-                         cryptoSession.cipherAlgorithm(), output.editArray());
+                         cryptoSession->cipherAlgorithm(), output.editArray());
 
   if (res == OEMCrypto_SUCCESS) {
     return android::OK;
@@ -804,17 +804,16 @@ status_t WVDrmPlugin::decrypt(const Vector<uint8_t>& sessionId,
                               const Vector<uint8_t>& iv,
                               Vector<uint8_t>& output) {
   CdmSessionId cdmSessionId(sessionId.begin(), sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return android::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
-  if (cryptoSession.cipherAlgorithm() == kInvalidCryptoAlgorithm) {
+  if (cryptoSession->cipherAlgorithm() == kInvalidCryptoAlgorithm) {
     return android::NO_INIT;
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            keyId.array(), keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -824,9 +823,9 @@ status_t WVDrmPlugin::decrypt(const Vector<uint8_t>& sessionId,
 
   output.resize(input.size());
 
-  res = mCrypto->decrypt(cryptoSession.oecSessionId(), input.array(),
+  res = mCrypto->decrypt(cryptoSession->oecSessionId(), input.array(),
                          input.size(), iv.array(),
-                         cryptoSession.cipherAlgorithm(), output.editArray());
+                         cryptoSession->cipherAlgorithm(), output.editArray());
 
   if (res == OEMCrypto_SUCCESS) {
     return android::OK;
@@ -841,17 +840,16 @@ status_t WVDrmPlugin::sign(const Vector<uint8_t>& sessionId,
                            const Vector<uint8_t>& message,
                            Vector<uint8_t>& signature) {
   CdmSessionId cdmSessionId(sessionId.begin(), sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return android::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
-  if (cryptoSession.macAlgorithm() == kInvalidCryptoAlgorithm) {
+  if (cryptoSession->macAlgorithm() == kInvalidCryptoAlgorithm) {
     return android::NO_INIT;
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            keyId.array(), keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -861,8 +859,8 @@ status_t WVDrmPlugin::sign(const Vector<uint8_t>& sessionId,
 
   size_t signatureSize = 0;
 
-  res = mCrypto->sign(cryptoSession.oecSessionId(), message.array(),
-                      message.size(), cryptoSession.macAlgorithm(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), message.array(),
+                      message.size(), cryptoSession->macAlgorithm(),
                       NULL, &signatureSize);
 
   if (res != OEMCrypto_ERROR_SHORT_BUFFER) {
@@ -877,8 +875,8 @@ status_t WVDrmPlugin::sign(const Vector<uint8_t>& sessionId,
 
   signature.resize(signatureSize);
 
-  res = mCrypto->sign(cryptoSession.oecSessionId(), message.array(),
-                      message.size(), cryptoSession.macAlgorithm(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), message.array(),
+                      message.size(), cryptoSession->macAlgorithm(),
                       signature.editArray(), &signatureSize);
 
   if (res == OEMCrypto_SUCCESS) {
@@ -895,17 +893,16 @@ status_t WVDrmPlugin::verify(const Vector<uint8_t>& sessionId,
                              const Vector<uint8_t>& signature,
                              bool& match) {
   CdmSessionId cdmSessionId(sessionId.begin(), sessionId.end());
-  if (!mCryptoSessions.count(cdmSessionId)) {
+  const shared_ptr<CryptoSession> cryptoSession = mCryptoSessions.get(cdmSessionId);
+  if (cryptoSession == nullptr) {
     return android::ERROR_DRM_SESSION_NOT_OPENED;
   }
 
-  const CryptoSession& cryptoSession = mCryptoSessions[cdmSessionId];
-
-  if (cryptoSession.macAlgorithm() == kInvalidCryptoAlgorithm) {
+  if (cryptoSession->macAlgorithm() == kInvalidCryptoAlgorithm) {
     return android::NO_INIT;
   }
 
-  OEMCryptoResult res = mCrypto->selectKey(cryptoSession.oecSessionId(),
+  OEMCryptoResult res = mCrypto->selectKey(cryptoSession->oecSessionId(),
                                            keyId.array(), keyId.size());
 
   if (res != OEMCrypto_SUCCESS) {
@@ -913,8 +910,8 @@ status_t WVDrmPlugin::verify(const Vector<uint8_t>& sessionId,
     return mapAndNotifyOfOEMCryptoResult(sessionId, res);
   }
 
-  res = mCrypto->verify(cryptoSession.oecSessionId(), message.array(),
-                        message.size(), cryptoSession.macAlgorithm(),
+  res = mCrypto->verify(cryptoSession->oecSessionId(), message.array(),
+                        message.size(), cryptoSession->macAlgorithm(),
                         signature.array(), signature.size());
 
   if (res == OEMCrypto_SUCCESS) {
diff --git a/libwvdrmengine/mediadrm/src_hidl/WVDrmPlugin.cpp b/libwvdrmengine/mediadrm/src_hidl/WVDrmPlugin.cpp
index abb2af61..5e8d077d 100644
--- a/libwvdrmengine/mediadrm/src_hidl/WVDrmPlugin.cpp
+++ b/libwvdrmengine/mediadrm/src_hidl/WVDrmPlugin.cpp
@@ -45,6 +45,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::kDefaultCdmIdentifier;
@@ -199,17 +201,18 @@ WVDrmPlugin::WVDrmPlugin(const sp<WvContentDecryptionModule>& cdm,
     mAppPackageName(appPackageName) {}
 
 WVDrmPlugin::~WVDrmPlugin() {
-  wvcdm::SetLoggingUid(mCdmIdentifierBuilder.user_id());
-  typedef map<CdmSessionId, CryptoSession>::iterator mapIterator;
-  for (mapIterator iter = mCryptoSessions.begin();
-       iter != mCryptoSessions.end();
-       ++iter) {
+  typedef map<CdmSessionId, shared_ptr<CryptoSession>>::iterator mapIterator;
+  auto cryptoSessions = mCryptoSessions.clear();
+  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);
@@ -253,7 +256,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.");
@@ -1352,7 +1355,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()) {
@@ -1392,7 +1395,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) {
@@ -1407,7 +1410,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.");
@@ -1415,7 +1418,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;
@@ -1497,14 +1500,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;
   }
@@ -1522,14 +1524,13 @@ Return<Status> WVDrmPlugin::setMacAlgorithm(
   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;
   }
@@ -1548,21 +1549,20 @@ Return<void> WVDrmPlugin::encrypt(
   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) {
@@ -1576,9 +1576,9 @@ Return<void> WVDrmPlugin::encrypt(
   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));
@@ -1601,21 +1601,20 @@ Return<void> WVDrmPlugin::decrypt(
   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) {
@@ -1629,9 +1628,9 @@ Return<void> WVDrmPlugin::decrypt(
   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));
@@ -1653,21 +1652,20 @@ Return<void> WVDrmPlugin::sign(
   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) {
@@ -1680,8 +1678,8 @@ Return<void> WVDrmPlugin::sign(
   size_t signatureSize = 0;
 
   const std::vector<uint8_t> msg = toVector(message);
-  res = mCrypto->sign(cryptoSession.oecSessionId(), msg.data(),
-                      msg.size(), cryptoSession.macAlgorithm(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), msg.data(),
+                      msg.size(), cryptoSession->macAlgorithm(),
                       NULL, &signatureSize);
 
   if (res != OEMCrypto_ERROR_SHORT_BUFFER) {
@@ -1698,8 +1696,8 @@ Return<void> WVDrmPlugin::sign(
 
   signature.resize(signatureSize);
 
-  res = mCrypto->sign(cryptoSession.oecSessionId(), msg.data(),
-                      msg.size(), cryptoSession.macAlgorithm(),
+  res = mCrypto->sign(cryptoSession->oecSessionId(), msg.data(),
+                      msg.size(), cryptoSession->macAlgorithm(),
                       signature.data(), &signatureSize);
 
   if (res == OEMCrypto_SUCCESS) {
@@ -1723,21 +1721,20 @@ Return<void> WVDrmPlugin::verify(
   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) {
@@ -1748,8 +1745,8 @@ Return<void> WVDrmPlugin::verify(
 
   const std::vector<uint8_t> _message = toVector(message);
   const std::vector<uint8_t> _signature = toVector(signature);
-  res = mCrypto->verify(cryptoSession.oecSessionId(), _message.data(),
-                        _message.size(), cryptoSession.macAlgorithm(),
+  res = mCrypto->verify(cryptoSession->oecSessionId(), _message.data(),
+                        _message.size(), cryptoSession->macAlgorithm(),
                         _signature.data(), _signature.size());
 
   if (res == OEMCrypto_SUCCESS) {