Snap for 6079661 from 0a62361d9ffd460afcaa32d8550f357e9c1e0812 to qt-aml-media-releaseandroid-mainline-10.0.0_r9android-mainline-10.0.0_r10android10-mainline-media-release

Change-Id: Ibf605f1c3c72503d74acae37f43fe514ac90d237

4 files changed, 32 insertions, 28 deletions

diff --git a/keystore/KeyStore.h b/keystore/KeyStore.h
index 69a02aea..a7fbab46 100644
--- a/keystore/KeyStore.h
+++ b/keystore/KeyStore.h
@@ -143,6 +143,23 @@ class KeyStore : public ::android::IBinder::DeathRecipient {
     KeystoreKeymasterEnforcement& getEnforcementPolicy() { return mEnforcementPolicy; }
     ConfirmationManager& getConfirmationManager() { return *mConfirmationManager; }
 
+    void addOperationDevice(sp<IBinder> token, std::shared_ptr<KeymasterWorker> dev) {
+        std::lock_guard<std::mutex> lock(operationDeviceMapMutex_);
+        operationDeviceMap_.emplace(std::move(token), std::move(dev));
+    }
+    std::shared_ptr<KeymasterWorker> getOperationDevice(const sp<IBinder>& token) {
+        std::lock_guard<std::mutex> lock(operationDeviceMapMutex_);
+        auto it = operationDeviceMap_.find(token);
+        if (it != operationDeviceMap_.end()) {
+            return it->second;
+        }
+        return {};
+    }
+    void removeOperationDevice(const sp<IBinder>& token) {
+        std::lock_guard<std::mutex> lock(operationDeviceMapMutex_);
+        operationDeviceMap_.erase(token);
+    }
+
   private:
     static const char* kOldMasterKey;
     static const char* kMetaDataFile;
@@ -173,6 +190,9 @@ class KeyStore : public ::android::IBinder::DeathRecipient {
     void writeMetaData();
 
     bool upgradeKeystore();
+
+    std::mutex operationDeviceMapMutex_;
+    std::map<sp<IBinder>, std::shared_ptr<KeymasterWorker>> operationDeviceMap_;
 };
 
 }  // namespace keystore
diff --git a/keystore/key_store_service.cpp b/keystore/key_store_service.cpp
index e1b1a661..2336e570 100644
--- a/keystore/key_store_service.cpp
+++ b/keystore/key_store_service.cpp
@@ -888,7 +888,7 @@ Status KeyStoreService::begin(const sp<IKeystoreOperationResultCallback>& cb,
                [this, cb, dev](OperationResult result_) {
                    if (result_.resultCode.isOk() ||
                        result_.resultCode == ResponseCode::OP_AUTH_NEEDED) {
-                       addOperationDevice(result_.token, dev);
+                       mKeyStore->addOperationDevice(result_.token, dev);
                    }
                    cb->onFinished(result_);
                });
@@ -905,14 +905,14 @@ Status KeyStoreService::update(const ::android::sp<IKeystoreOperationResultCallb
         return AIDL_RETURN(ErrorCode::INVALID_ARGUMENT);
     }
 
-    auto dev = getOperationDevice(token);
+    auto dev = mKeyStore->getOperationDevice(token);
     if (!dev) {
         return AIDL_RETURN(ErrorCode::INVALID_OPERATION_HANDLE);
     }
 
     dev->update(token, params.getParameters(), input, [this, cb, token](OperationResult result_) {
         if (!result_.resultCode.isOk()) {
-            removeOperationDevice(token);
+            mKeyStore->removeOperationDevice(token);
         }
         cb->onFinished(result_);
     });
@@ -930,16 +930,14 @@ Status KeyStoreService::finish(const ::android::sp<IKeystoreOperationResultCallb
         return AIDL_RETURN(ErrorCode::INVALID_ARGUMENT);
     }
 
-    auto dev = getOperationDevice(token);
+    auto dev = mKeyStore->getOperationDevice(token);
     if (!dev) {
         return AIDL_RETURN(ErrorCode::INVALID_OPERATION_HANDLE);
     }
 
     dev->finish(token, params.getParameters(), {}, signature, entropy,
                 [this, cb, token](OperationResult result_) {
-                    if (!result_.resultCode.isOk()) {
-                        removeOperationDevice(token);
-                    }
+                    mKeyStore->removeOperationDevice(token);
                     cb->onFinished(result_);
                 });
 
@@ -950,12 +948,15 @@ Status KeyStoreService::abort(const ::android::sp<IKeystoreResponseCallback>& cb
                               const ::android::sp<::android::IBinder>& token,
                               int32_t* _aidl_return) {
     KEYSTORE_SERVICE_LOCK;
-    auto dev = getOperationDevice(token);
+    auto dev = mKeyStore->getOperationDevice(token);
     if (!dev) {
         return AIDL_RETURN(ErrorCode::INVALID_OPERATION_HANDLE);
     }
 
-    dev->abort(token, [cb](KeyStoreServiceReturnCode rc) { cb->onFinished(rc); });
+    dev->abort(token, [this, cb, token](KeyStoreServiceReturnCode rc) {
+        mKeyStore->removeOperationDevice(token);
+        cb->onFinished(rc);
+    });
 
     return AIDL_RETURN(ResponseCode::NO_ERROR);
 }
diff --git a/keystore/key_store_service.h b/keystore/key_store_service.h
index 2fdc3dd2..96d0c079 100644
--- a/keystore/key_store_service.h
+++ b/keystore/key_store_service.h
@@ -243,25 +243,6 @@ class KeyStoreService : public android::security::keystore::BnKeystoreService {
      */
     std::mutex keystoreServiceMutex_;
 
-    std::mutex operationDeviceMapMutex_;
-    std::map<sp<IBinder>, std::shared_ptr<KeymasterWorker>> operationDeviceMap_;
-
-    void addOperationDevice(sp<IBinder> token, std::shared_ptr<KeymasterWorker> dev) {
-        std::lock_guard<std::mutex> lock(operationDeviceMapMutex_);
-        operationDeviceMap_.emplace(std::move(token), std::move(dev));
-    }
-    std::shared_ptr<KeymasterWorker> getOperationDevice(const sp<IBinder>& token) {
-        std::lock_guard<std::mutex> lock(operationDeviceMapMutex_);
-        auto it = operationDeviceMap_.find(token);
-        if (it != operationDeviceMap_.end()) {
-            return it->second;
-        }
-        return {};
-    }
-    void removeOperationDevice(const sp<IBinder>& token) {
-        std::lock_guard<std::mutex> lock(operationDeviceMapMutex_);
-        operationDeviceMap_.erase(token);
-    }
 };
 
 };  // namespace keystore
diff --git a/keystore/keymaster_worker.cpp b/keystore/keymaster_worker.cpp
index 922ef0a4..728e607d 100644
--- a/keystore/keymaster_worker.cpp
+++ b/keystore/keymaster_worker.cpp
@@ -321,6 +321,7 @@ bool KeymasterWorker::pruneOperation() {
     // We mostly ignore errors from abort() because all we care about is whether at least
     // one operation has been removed.
     auto rc = abort(oldest);
+    keyStore_->removeOperationDevice(oldest);
     if (operationMap_.getOperationCount() >= op_count_before_abort) {
         ALOGE("Failed to abort pruneable operation %p, error: %d", oldest.get(), rc.getErrorCode());
         return false;
@@ -1091,6 +1092,7 @@ void KeymasterWorker::binderDied(android::wp<IBinder> who) {
         auto operations = operationMap_.getOperationsForToken(who.unsafe_get());
         for (const auto& token : operations) {
             abort(token);
+            keyStore_->removeOperationDevice(token);
         }
     });
 }