Refactor OperationMap interface.

Test: runtest --path cts/tests/tests/keystore/src/android/keystore/cts
Change-Id: If1cdbe8e5f1f32fe04bb2f6083221e2c00840585
diff --git a/keystore/key_store_service.cpp b/keystore/key_store_service.cpp
index 0a8648f..c95b06a 100644
--- a/keystore/key_store_service.cpp
+++ b/keystore/key_store_service.cpp
@@ -813,7 +813,7 @@
     }
 
     if (containsTag(params.getParameters(), Tag::INCLUDE_UNIQUE_ID)) {
-        //TODO(jbires): remove uid checking upon implementation of b/25646100
+        // TODO(jbires): remove uid checking upon implementation of b/25646100
         if (!checkBinderPermission(P_GEN_UNIQUE_ID) ||
             originalUid != IPCThreadState::self()->getCallingUid()) {
             *aidl_return = static_cast<int32_t>(ResponseCode::PERMISSION_DENIED);
@@ -1363,74 +1363,65 @@
         result->resultCode = ErrorCode::INVALID_ARGUMENT;
         return Status::ok();
     }
-    km_device_t dev;
-    uint64_t handle;
-    KeyPurpose purpose;
-    km_id_t keyid;
-    const KeyCharacteristics* characteristics;
-    if (!mOperationMap.getOperation(token, &handle, &keyid, &purpose, &dev, &characteristics)) {
+
+    auto getOpResult = mOperationMap.getOperation(token);
+    if (!getOpResult.isOk()) {
         result->resultCode = ErrorCode::INVALID_OPERATION_HANDLE;
         return Status::ok();
     }
+    const auto& op = getOpResult.value();
+
     AuthorizationSet opParams = params.getParameters();
     result->resultCode = addOperationAuthTokenIfNeeded(token, &opParams);
-    if (!result->resultCode.isOk()) {
-        return Status::ok();
-    }
+    if (!result->resultCode.isOk()) return Status::ok();
 
     // Check that all key authorization policy requirements are met.
-    AuthorizationSet key_auths(characteristics->teeEnforced);
-    key_auths.append(&characteristics->softwareEnforced[0],
-                     &characteristics->softwareEnforced[characteristics->softwareEnforced.size()]);
+    AuthorizationSet key_auths(op.characteristics.teeEnforced);
+    key_auths.append(op.characteristics.softwareEnforced.begin(),
+                     op.characteristics.softwareEnforced.end());
+
     result->resultCode = enforcement_policy.AuthorizeOperation(
-        purpose, keyid, key_auths, opParams, handle, false /* is_begin_operation */);
-    if (!result->resultCode.isOk()) {
-        return Status::ok();
-    }
+        op.purpose, op.keyid, key_auths, opParams, op.handle, false /* is_begin_operation */);
+    if (!result->resultCode.isOk()) return Status::ok();
 
     auto hidlCb = [&](ErrorCode ret, uint32_t inputConsumed,
                       const hidl_vec<KeyParameter>& outParams,
                       const ::std::vector<uint8_t>& output) {
         result->resultCode = ret;
-        if (!result->resultCode.isOk()) {
-            return;
+        if (result->resultCode.isOk()) {
+            result->inputConsumed = inputConsumed;
+            result->outParams = outParams;
+            result->data = output;
         }
-        result->inputConsumed = inputConsumed;
-        result->outParams = outParams;
-        result->data = output;
     };
 
     KeyStoreServiceReturnCode rc =
-        KS_HANDLE_HIDL_ERROR(dev->update(handle, opParams.hidl_data(), data, hidlCb));
+        KS_HANDLE_HIDL_ERROR(op.device->update(op.handle, opParams.hidl_data(), data, hidlCb));
+
     // just a reminder: on success result->resultCode was set in the callback. So we only overwrite
     // it if there was a communication error indicated by the ErrorCode.
-    if (!rc.isOk()) {
-        result->resultCode = rc;
-    }
+    if (!rc.isOk()) result->resultCode = rc;
+
     return Status::ok();
 }
 
 Status KeyStoreService::finish(const sp<IBinder>& token, const KeymasterArguments& params,
                                const ::std::vector<uint8_t>& signature,
                                const ::std::vector<uint8_t>& entropy, OperationResult* result) {
+    auto getOpResult = mOperationMap.getOperation(token);
+    if (!getOpResult.isOk()) {
+        result->resultCode = ErrorCode::INVALID_OPERATION_HANDLE;
+        return Status::ok();
+    }
+    const auto& op = std::move(getOpResult.value());
     if (!checkAllowedOperationParams(params.getParameters())) {
         result->resultCode = ErrorCode::INVALID_ARGUMENT;
         return Status::ok();
     }
-    km_device_t dev;
-    uint64_t handle;
-    KeyPurpose purpose;
-    km_id_t keyid;
-    const KeyCharacteristics* characteristics;
-    if (!mOperationMap.getOperation(token, &handle, &keyid, &purpose, &dev, &characteristics)) {
-        result->resultCode = ErrorCode::INVALID_OPERATION_HANDLE;
-        return Status::ok();
-    }
+
     AuthorizationSet opParams = params.getParameters();
     result->resultCode = addOperationAuthTokenIfNeeded(token, &opParams);
-    if (!result->resultCode.isOk()) {
-        return Status::ok();
-    }
+    if (!result->resultCode.isOk()) return Status::ok();
 
     if (entropy.size()) {
         int resultCode;
@@ -1442,29 +1433,29 @@
     }
 
     // Check that all key authorization policy requirements are met.
-    AuthorizationSet key_auths(characteristics->teeEnforced);
-    key_auths.append(&characteristics->softwareEnforced[0],
-                     &characteristics->softwareEnforced[characteristics->softwareEnforced.size()]);
+    AuthorizationSet key_auths(op.characteristics.teeEnforced);
+    key_auths.append(op.characteristics.softwareEnforced.begin(),
+                     op.characteristics.softwareEnforced.end());
+
     result->resultCode = enforcement_policy.AuthorizeOperation(
-        purpose, keyid, key_auths, opParams, handle, false /* is_begin_operation */);
+        op.purpose, op.keyid, key_auths, opParams, op.handle, false /* is_begin_operation */);
     if (!result->resultCode.isOk()) return Status::ok();
 
     auto hidlCb = [&](ErrorCode ret, const hidl_vec<KeyParameter>& outParams,
                       const ::std::vector<uint8_t>& output) {
         result->resultCode = ret;
-        if (!result->resultCode.isOk()) {
+        if (result->resultCode.isOk()) {
+            result->outParams = outParams;
+            result->data = output;
         }
-        result->outParams = outParams;
-        result->data = output;
     };
 
     KeyStoreServiceReturnCode rc = KS_HANDLE_HIDL_ERROR(
-        dev->finish(handle, opParams.hidl_data(),
-                    ::std::vector<uint8_t>() /* TODO(swillden): wire up input to finish() */,
-                    signature, hidlCb));
-    // Remove the operation regardless of the result
+        op.device->finish(op.handle, opParams.hidl_data(),
+                          ::std::vector<uint8_t>() /* TODO(swillden): wire up input to finish() */,
+                          signature, hidlCb));
     mOperationMap.removeOperation(token);
-    mAuthTokenTable.MarkCompleted(handle);
+    mAuthTokenTable.MarkCompleted(op.handle);
 
     // just a reminder: on success result->resultCode was set in the callback. So we only overwrite
     // it if there was a communication error indicated by the ErrorCode.
@@ -1475,35 +1466,20 @@
 }
 
 Status KeyStoreService::abort(const sp<IBinder>& token, int32_t* aidl_return) {
-    km_device_t dev;
-    uint64_t handle;
-    KeyPurpose purpose;
-    km_id_t keyid;
-    if (!mOperationMap.getOperation(token, &handle, &keyid, &purpose, &dev, NULL)) {
-        *aidl_return =
-            static_cast<int32_t>(KeyStoreServiceReturnCode(ErrorCode::INVALID_OPERATION_HANDLE));
+    auto getOpResult = mOperationMap.removeOperation(token);
+    if (!getOpResult.isOk()) {
+        *aidl_return = static_cast<int32_t>(ErrorCode::INVALID_OPERATION_HANDLE);
         return Status::ok();
     }
-    mOperationMap.removeOperation(token);
+    auto op = std::move(getOpResult.value());
+    mAuthTokenTable.MarkCompleted(op.handle);
 
-    ErrorCode error_code = KS_HANDLE_HIDL_ERROR(dev->abort(handle));
-    mAuthTokenTable.MarkCompleted(handle);
+    ErrorCode error_code = KS_HANDLE_HIDL_ERROR(op.device->abort(op.handle));
     *aidl_return = static_cast<int32_t>(KeyStoreServiceReturnCode(error_code));
     return Status::ok();
 }
 
 Status KeyStoreService::isOperationAuthorized(const sp<IBinder>& token, bool* aidl_return) {
-    km_device_t dev;
-    uint64_t handle;
-    const KeyCharacteristics* characteristics;
-    KeyPurpose purpose;
-    km_id_t keyid;
-    if (!mOperationMap.getOperation(token, &handle, &keyid, &purpose, &dev, &characteristics)) {
-        *aidl_return = false;
-        return Status::ok();
-    }
-    const HardwareAuthToken* authToken = NULL;
-    mOperationMap.getOperationAuthToken(token, &authToken);
     AuthorizationSet ignored;
     auto authResult = addOperationAuthTokenIfNeeded(token, &ignored);
     *aidl_return = authResult.isOk();
@@ -1940,26 +1916,18 @@
  */
 KeyStoreServiceReturnCode KeyStoreService::addOperationAuthTokenIfNeeded(const sp<IBinder>& token,
                                                                          AuthorizationSet* params) {
-    const HardwareAuthToken* authToken = nullptr;
-    mOperationMap.getOperationAuthToken(token, &authToken);
-    if (!authToken) {
-        km_device_t dev;
-        uint64_t handle;
-        const KeyCharacteristics* characteristics = nullptr;
-        KeyPurpose purpose;
-        km_id_t keyid;
-        if (!mOperationMap.getOperation(token, &handle, &keyid, &purpose, &dev, &characteristics)) {
-            return ErrorCode::INVALID_OPERATION_HANDLE;
-        }
-        auto result = getAuthToken(*characteristics, handle, purpose, &authToken);
-        if (!result.isOk()) {
-            return result;
-        }
-        if (authToken) {
-            mOperationMap.setOperationAuthToken(token, *authToken);
-        }
+    auto getOpResult = mOperationMap.getOperation(token);
+    if (!getOpResult.isOk()) return ErrorCode::INVALID_OPERATION_HANDLE;
+    const auto& op = getOpResult.value();
+
+    if (!op.authToken) {
+        const HardwareAuthToken* found = nullptr;
+        auto result = getAuthToken(op.characteristics, op.handle, op.purpose, &found);
+        if (!result.isOk()) return result;
+        if (found) mOperationMap.setOperationAuthToken(token, *found);
+        assert(*op.authToken == *found);
     }
-    addAuthTokenToParams(params, authToken);
+    addAuthTokenToParams(params, op.authToken.get());
     return ResponseCode::NO_ERROR;
 }
 
diff --git a/keystore/operation.cpp b/keystore/operation.cpp
index e1ba713..66fcee2 100644
--- a/keystore/operation.cpp
+++ b/keystore/operation.cpp
@@ -20,7 +20,6 @@
 #include <algorithm>
 
 namespace keystore {
-using namespace android;
 
 OperationMap::OperationMap(IBinder::DeathRecipient* deathRecipient)
     : mDeathRecipient(deathRecipient) {}
@@ -29,38 +28,22 @@
                                        const OperationMap::km_device_t& dev,
                                        const sp<IBinder>& appToken,
                                        KeyCharacteristics&& characteristics, bool pruneable) {
-    sp<IBinder> token = new BBinder();
-    mMap[token] = Operation(handle, keyid, purpose, dev, std::move(characteristics), appToken);
-    if (pruneable) {
-        mLru.push_back(token);
-    }
-    if (mAppTokenMap.find(appToken) == mAppTokenMap.end()) {
-        appToken->linkToDeath(mDeathRecipient);
-    }
+    sp<IBinder> token = new ::android::BBinder();
+    mMap.emplace(token,
+                 Operation(handle, keyid, purpose, dev, std::move(characteristics), appToken));
+    if (pruneable) mLru.push_back(token);
+    if (mAppTokenMap.find(appToken) == mAppTokenMap.end()) appToken->linkToDeath(mDeathRecipient);
     mAppTokenMap[appToken].push_back(token);
+
     return token;
 }
 
-bool OperationMap::getOperation(const sp<IBinder>& token, uint64_t* outHandle, uint64_t* outKeyid,
-                                KeyPurpose* outPurpose, km_device_t* outDevice,
-                                const KeyCharacteristics** outCharacteristics) {
-    if (!outHandle || !outDevice) {
-        return false;
-    }
+NullOr<const OperationMap::Operation&> OperationMap::getOperation(const sp<IBinder>& token) {
     auto entry = mMap.find(token);
-    if (entry == mMap.end()) {
-        return false;
-    }
-    updateLru(token);
+    if (entry == mMap.end()) return {};
 
-    *outHandle = entry->second.handle;
-    *outKeyid = entry->second.keyid;
-    *outPurpose = entry->second.purpose;
-    *outDevice = entry->second.device;
-    if (outCharacteristics) {
-        *outCharacteristics = &entry->second.characteristics;
-    }
-    return true;
+    updateLru(token);
+    return entry->second;
 }
 
 void OperationMap::updateLru(const sp<IBinder>& token) {
@@ -71,19 +54,18 @@
     }
 }
 
-bool OperationMap::removeOperation(const sp<IBinder>& token) {
+NullOr<OperationMap::Operation> OperationMap::removeOperation(const sp<IBinder>& token) {
     auto entry = mMap.find(token);
-    if (entry == mMap.end()) {
-        return false;
-    }
-    sp<IBinder> appToken = entry->second.appToken;
+    if (entry == mMap.end()) return {};
+
+    Operation op = std::move(entry->second);
     mMap.erase(entry);
+
     auto lruEntry = std::find(mLru.begin(), mLru.end(), token);
-    if (lruEntry != mLru.end()) {
-        mLru.erase(lruEntry);
-    }
-    removeOperationTracking(token, appToken);
-    return true;
+    if (lruEntry != mLru.end()) mLru.erase(lruEntry);
+    removeOperationTracking(token, op.appToken);
+
+    return op;
 }
 
 void OperationMap::removeOperationTracking(const sp<IBinder>& token, const sp<IBinder>& appToken) {
@@ -102,7 +84,7 @@
 }
 
 bool OperationMap::hasPruneableOperation() const {
-    return mLru.size() != 0;
+    return !mLru.empty();
 }
 
 size_t OperationMap::getPruneableOperationCount() const {
@@ -110,38 +92,22 @@
 }
 
 sp<IBinder> OperationMap::getOldestPruneableOperation() {
-    if (!hasPruneableOperation()) {
-        return sp<IBinder>(NULL);
-    }
-    return mLru[0];
-}
-
-bool OperationMap::getOperationAuthToken(const sp<IBinder>& token,
-                                         const HardwareAuthToken** outToken) {
-    auto entry = mMap.find(token);
-    if (entry == mMap.end()) {
-        return false;
-    }
-    *outToken = entry->second.authToken.get();
-    return true;
+    if (!hasPruneableOperation()) return sp<IBinder>(nullptr);
+    return mLru.front();
 }
 
 bool OperationMap::setOperationAuthToken(const sp<IBinder>& token, HardwareAuthToken authToken) {
     auto entry = mMap.find(token);
-    if (entry == mMap.end()) {
-        return false;
-    }
+    if (entry == mMap.end()) return false;
+
     entry->second.authToken = std::make_unique<HardwareAuthToken>(std::move(authToken));
     return true;
 }
 
 std::vector<sp<IBinder>> OperationMap::getOperationsForToken(const sp<IBinder>& appToken) {
     auto appEntry = mAppTokenMap.find(appToken);
-    if (appEntry != mAppTokenMap.end()) {
-        return appEntry->second;
-    } else {
-        return std::vector<sp<IBinder>>();
-    }
+    if (appEntry == mAppTokenMap.end()) return {};
+    return appEntry->second;
 }
 
 OperationMap::Operation::Operation(uint64_t handle_, uint64_t keyid_, KeyPurpose purpose_,
@@ -150,7 +116,4 @@
     : handle(handle_), keyid(keyid_), purpose(purpose_), device(device_),
       characteristics(characteristics_), appToken(appToken_) {}
 
-OperationMap::Operation::Operation()
-    : handle(0), keyid(0), device(nullptr), characteristics(), appToken(nullptr) {}
-
-}  // namespace android
+}  // namespace keystore
diff --git a/keystore/operation.h b/keystore/operation.h
index 9249437..ac8e945 100644
--- a/keystore/operation.h
+++ b/keystore/operation.h
@@ -38,51 +38,45 @@
  */
 
 class OperationMap {
-    typedef ::android::sp<::android::hardware::keymaster::V3_0::IKeymasterDevice> km_device_t;
+    typedef sp<::android::hardware::keymaster::V3_0::IKeymasterDevice> km_device_t;
 
   public:
-    explicit OperationMap(IBinder::DeathRecipient* deathRecipient);
-    android::sp<android::IBinder> addOperation(uint64_t handle, uint64_t keyid, KeyPurpose purpose,
-                                               const km_device_t& dev,
-                                               const android::sp<android::IBinder>& appToken,
-                                               KeyCharacteristics&& characteristics,
-                                               bool pruneable);
-    bool getOperation(const android::sp<android::IBinder>& token, uint64_t* outHandle,
-                      uint64_t* outKeyid, KeyPurpose* outPurpose, km_device_t* outDev,
-                      const KeyCharacteristics** outCharacteristics);
-    bool removeOperation(const android::sp<android::IBinder>& token);
-    bool hasPruneableOperation() const;
-    size_t getOperationCount() const { return mMap.size(); }
-    size_t getPruneableOperationCount() const;
-    bool getOperationAuthToken(const android::sp<android::IBinder>& token,
-                               const HardwareAuthToken** outToken);
-    bool setOperationAuthToken(const android::sp<android::IBinder>& token,
-                               HardwareAuthToken authToken);
-    android::sp<android::IBinder> getOldestPruneableOperation();
-    std::vector<android::sp<android::IBinder>>
-    getOperationsForToken(const android::sp<android::IBinder>& appToken);
-
-  private:
-    void updateLru(const android::sp<android::IBinder>& token);
-    void removeOperationTracking(const android::sp<android::IBinder>& token,
-                                 const android::sp<android::IBinder>& appToken);
     struct Operation {
-        Operation();
+        Operation() = default;
         Operation(uint64_t handle, uint64_t keyid, KeyPurpose purpose, const km_device_t& device,
-                  KeyCharacteristics&& characteristics, android::sp<android::IBinder> appToken);
+                  KeyCharacteristics&& characteristics, sp<IBinder> appToken);
+        Operation(Operation&&) = default;
+        Operation(const Operation&) = delete;
+
         uint64_t handle;
         uint64_t keyid;
         KeyPurpose purpose;
         km_device_t device;
         KeyCharacteristics characteristics;
-        android::sp<android::IBinder> appToken;
+        sp<IBinder> appToken;
         std::unique_ptr<HardwareAuthToken> authToken;
     };
-    std::map<android::sp<android::IBinder>, Operation> mMap;
-    std::vector<android::sp<android::IBinder>> mLru;
-    std::map<android::sp<android::IBinder>, std::vector<android::sp<android::IBinder>>>
-        mAppTokenMap;
-    android::IBinder::DeathRecipient* mDeathRecipient;
+
+    explicit OperationMap(IBinder::DeathRecipient* deathRecipient);
+    sp<IBinder> addOperation(uint64_t handle, uint64_t keyid, KeyPurpose purpose,
+                             const km_device_t& dev, const sp<IBinder>& appToken,
+                             KeyCharacteristics&& characteristics, bool pruneable);
+    NullOr<const Operation&> getOperation(const sp<IBinder>& token);
+    NullOr<Operation> removeOperation(const sp<IBinder>& token);
+    bool hasPruneableOperation() const;
+    size_t getOperationCount() const { return mMap.size(); }
+    size_t getPruneableOperationCount() const;
+    bool setOperationAuthToken(const sp<IBinder>& token, HardwareAuthToken authToken);
+    sp<IBinder> getOldestPruneableOperation();
+    std::vector<sp<IBinder>> getOperationsForToken(const sp<IBinder>& appToken);
+
+  private:
+    void updateLru(const sp<IBinder>& token);
+    void removeOperationTracking(const sp<IBinder>& token, const sp<IBinder>& appToken);
+    std::map<sp<IBinder>, Operation> mMap;
+    std::vector<sp<IBinder>> mLru;
+    std::map<sp<IBinder>, std::vector<sp<IBinder>>> mAppTokenMap;
+    IBinder::DeathRecipient* mDeathRecipient;
 };
 
 }  // namespace keystore