Fix memory leak in keystore

When operations are aborted by an app or an app dies, tokens are not
removed from the device token map. This patch moves the this map from
key_store_service to KeyStore so that it can be accessed by the
keymaster workers. It also adds calls to removeOperationDevice to the
binderDied hook of the keymaster workers as well as to keystore service
abort.

Add a call to removeOperationDevice() inside pruneOperation() function on
keystore/keymaster_worker.cpp

Bug: 139383076
Bug: 141317862
Test: atest keystore_unit_tests (passed)

Merged-In: I90d4dc9d4510f4ac250022c89240a742b9e8d4b4
Change-Id: I90d4dc9d4510f4ac250022c89240a742b9e8d4b4
diff --git a/keystore/KeyStore.h b/keystore/KeyStore.h
index 69a02ae..a7fbab4 100644
--- a/keystore/KeyStore.h
+++ b/keystore/KeyStore.h
@@ -143,6 +143,23 @@
     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 @@
     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 e1b1a66..7c42d0b 100644
--- a/keystore/key_store_service.cpp
+++ b/keystore/key_store_service.cpp
@@ -888,7 +888,7 @@
                [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 @@
         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,7 +930,7 @@
         return AIDL_RETURN(ErrorCode::INVALID_ARGUMENT);
     }
 
-    auto dev = getOperationDevice(token);
+    auto dev = mKeyStore->getOperationDevice(token);
     if (!dev) {
         return AIDL_RETURN(ErrorCode::INVALID_OPERATION_HANDLE);
     }
@@ -938,7 +938,7 @@
     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 +950,15 @@
                               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 2fdc3dd..96d0c07 100644
--- a/keystore/key_store_service.h
+++ b/keystore/key_store_service.h
@@ -243,25 +243,6 @@
      */
     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 922ef0a..728e607 100644
--- a/keystore/keymaster_worker.cpp
+++ b/keystore/keymaster_worker.cpp
@@ -321,6 +321,7 @@
     // 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 @@
         auto operations = operationMap_.getOperationsForToken(who.unsafe_get());
         for (const auto& token : operations) {
             abort(token);
+            keyStore_->removeOperationDevice(token);
         }
     });
 }