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 e0ee937..b6b7295 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,16 +930,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->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 @@
                               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);
 }
@@ -1039,6 +1040,10 @@
     std::tie(rc, keyBlob, charBlob, lockedEntry) =
         mKeyStore->getKeyForName(name8, callingUid, TYPE_KEYMASTER_10);
 
+    if (!rc.isOk()) {
+        return AIDL_RETURN(rc);
+    }
+
     auto dev = mKeyStore->getDevice(keyBlob);
     auto hidlKey = blob2hidlVec(keyBlob);
     dev->attestKey(
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);
         }
     });
 }