Make vold use keystore2 instead of keymaster

Make vold use keystore2 for all its operations instead of directly using
keymaster. This way, we won't have any clients that bypass keystore2,
and we'll no longer need to reserve a keymaster operation for vold.

Note that we now hardcode "SecurityLevel::TRUSTED_ENVIRONMENT" (TEE)
when talking to Keystore2 since Keystore2 only allows TEE and STRONGBOX.
Keystore2 presents any SOFTWARE implementation as a TEE to callers when
no "real" TEE is present. As far as storage encryption is concerned,
there's no advantage to using a STRONGBOX when a "real" TEE is present,
and a STRONGBOX can't be present if a "real" TEE isn't, so asking
Keystore2 for a TEE is the best we can do in any situation.

The difference in behaviour only really affects the full disk encryption
code in cryptfs.cpp, which used to explicitly check that the keymaster
device is a "real" TEE (as opposed to a SOFTWARE implementation) before
using it (it can no longer do so since Keystore2 doesn't provide a way
to do this).

A little code history digging (7c49ab0a0b in particular) shows that
cryptfs.cpp cared about two things when using a keymaster.
 - 1) that the keys generated by the keymaster were "standalone" keys -
      i.e. that the keymaster could operate on those keys without
      requiring /data or any other service to be available.
 - 2) that the keymaster was a non-SOFTWARE implementation so that things
      would still work in case a "real" TEE keymaster was ever somehow
      added to the device after first boot.

Today, all "real" TEE keymasters always generate "standalone" keys, and
a TEE has been required in Android devices since at least Android N. The
only two exceptions are Goldfish and ARC++, which have SOFTWARE
keymasters, but both those keymasters also generate "standalone" keys.

We're also no longer worried about possibly adding a "real" TEE KM to
either of those devices after first boot. So there's no longer a reason
cryptfs.cpp can't use the SOFTWARE keymaster on those devices.

There's also already an upgrade path in place (see
test_mount_encrypted_fs() in cryptfs.cpp) to upgrade the kdf that's
being used once a TEE keymaster is added to the device. So it's safe for
cryptfs.cpp to ask for a TEE keymaster from Keystore2 and use it
blindly, without checking whether or not it's a "real" TEE, which is why
Keymaster::isSecure() just returns true now. A future patch will remove
that function and simplify its callers.

Bug: 181910578
Test: cuttlefish and bramble boot. Adding, switching between, stopping
      and removing users work.
Change-Id: Iaebfef082eca0da8a305043fafb6d85e5de14cf8
diff --git a/Android.bp b/Android.bp
index f2a1a37..84809a3 100644
--- a/Android.bp
+++ b/Android.bp
@@ -46,9 +46,6 @@
         "libvold_binder",
     ],
     shared_libs: [
-        "android.hardware.keymaster@3.0",
-        "android.hardware.keymaster@4.0",
-        "android.hardware.keymaster@4.1",
         "android.hardware.boot@1.0",
         "libbase",
         "libbinder",
@@ -63,8 +60,6 @@
         "libhardware_legacy",
         "libincfs",
         "libhidlbase",
-        "libkeymaster4support",
-        "libkeymaster4_1support",
         "libkeyutils",
         "liblog",
         "liblogwrap",
@@ -171,7 +166,10 @@
     shared_libs: [
         "android.hardware.health.storage@1.0",
         "android.hardware.health.storage-V1-ndk_platform",
+        "android.system.keystore2-V1-ndk_platform",
+        "android.security.maintenance-ndk_platform",
         "libbinder_ndk",
+        "libkeymint_support",
     ],
     whole_static_libs: [
         "com.android.sysprop.apex",
@@ -202,7 +200,10 @@
     shared_libs: [
         "android.hardware.health.storage@1.0",
         "android.hardware.health.storage-V1-ndk_platform",
+        "android.system.keystore2-V1-ndk_platform",
+        "android.security.maintenance-ndk_platform",
         "libbinder_ndk",
+        "libkeymint_support",
     ],
 
     product_variables: {
@@ -245,15 +246,14 @@
     shared_libs: [
         "libbase",
         "libbinder",
+        "libbinder_ndk",
 
-        "android.hardware.keymaster@3.0",
-        "android.hardware.keymaster@4.0",
-        "android.hardware.keymaster@4.1",
+        "android.system.keystore2-V1-ndk_platform",
+        "android.security.maintenance-ndk_platform",
         "libhardware",
         "libhardware_legacy",
         "libhidlbase",
-        "libkeymaster4support",
-        "libkeymaster4_1support",
+        "libkeymint_support",
         "libutils",
     ],
 }
diff --git a/KeyStorage.cpp b/KeyStorage.cpp
index 11e85fd..8d518de 100644
--- a/KeyStorage.cpp
+++ b/KeyStorage.cpp
@@ -46,8 +46,6 @@
 #include <cutils/properties.h>
 
 #include <hardware/hw_auth_token.h>
-#include <keymasterV4_1/authorization_set.h>
-#include <keymasterV4_1/keymaster_utils.h>
 
 extern "C" {
 
@@ -151,12 +149,11 @@
 
 static bool generateKeyStorageKey(Keymaster& keymaster, const std::string& appId,
                                   std::string* key) {
-    auto paramBuilder =
-            km::AuthorizationSetBuilder()
-                    .AesEncryptionKey(AES_KEY_BYTES * 8)
-                    .GcmModeMinMacLen(GCM_MAC_BYTES * 8)
-                    .Authorization(km::TAG_APPLICATION_ID, km::support::blob2hidlVec(appId))
-                    .Authorization(km::TAG_NO_AUTH_REQUIRED);
+    auto paramBuilder = km::AuthorizationSetBuilder()
+                                .AesEncryptionKey(AES_KEY_BYTES * 8)
+                                .GcmModeMinMacLen(GCM_MAC_BYTES * 8)
+                                .Authorization(km::TAG_APPLICATION_ID, appId)
+                                .Authorization(km::TAG_NO_AUTH_REQUIRED);
     LOG(DEBUG) << "Generating \"key storage\" key that doesn't need auth token";
     return generateKeymasterKey(keymaster, paramBuilder, key);
 }
@@ -187,7 +184,7 @@
 static km::AuthorizationSet beginParams(const std::string& appId) {
     return km::AuthorizationSetBuilder()
             .GcmModeMacLen(GCM_MAC_BYTES * 8)
-            .Authorization(km::TAG_APPLICATION_ID, km::support::blob2hidlVec(appId));
+            .Authorization(km::TAG_APPLICATION_ID, appId);
 }
 
 static bool readFileToString(const std::string& filename, std::string* result) {
@@ -320,7 +317,6 @@
 
 // Begins a Keymaster operation using the key stored in |dir|.
 static KeymasterOperation BeginKeymasterOp(Keymaster& keymaster, const std::string& dir,
-                                           km::KeyPurpose purpose,
                                            const km::AuthorizationSet& keyParams,
                                            const km::AuthorizationSet& opParams,
                                            km::AuthorizationSet* outParams) {
@@ -344,9 +340,11 @@
         if (!readFileToString(blob_file, &blob)) return KeymasterOperation();
     }
 
-    auto opHandle = keymaster.begin(purpose, blob, inParams, outParams);
-    if (opHandle) return opHandle;
-    if (opHandle.errorCode() != km::ErrorCode::KEY_REQUIRES_UPGRADE) return opHandle;
+    auto opHandle = keymaster.begin(blob, inParams, outParams);
+    if (!opHandle) return opHandle;
+
+    // If key blob wasn't upgraded, nothing left to do.
+    if (!opHandle.getUpgradedBlob()) return opHandle;
 
     if (already_upgraded) {
         LOG(ERROR) << "Unexpected case; already-upgraded key " << upgraded_blob_file
@@ -354,8 +352,8 @@
         return KeymasterOperation();
     }
     LOG(INFO) << "Upgrading key: " << blob_file;
-    if (!keymaster.upgradeKey(blob, keyParams, &blob)) return KeymasterOperation();
-    if (!writeStringToFile(blob, upgraded_blob_file)) return KeymasterOperation();
+    if (!writeStringToFile(*opHandle.getUpgradedBlob(), upgraded_blob_file))
+        return KeymasterOperation();
     if (cp_needsCheckpoint()) {
         LOG(INFO) << "Wrote upgraded key to " << upgraded_blob_file
                   << "; delaying commit due to checkpoint";
@@ -364,26 +362,24 @@
         if (!CommitUpgradedKey(keymaster, dir)) return KeymasterOperation();
         LOG(INFO) << "Key upgraded: " << blob_file;
     }
-
-    return keymaster.begin(purpose, blob, inParams, outParams);
+    return opHandle;
 }
 
 static bool encryptWithKeymasterKey(Keymaster& keymaster, const std::string& dir,
                                     const km::AuthorizationSet& keyParams,
                                     const KeyBuffer& message, std::string* ciphertext) {
-    km::AuthorizationSet opParams;
+    km::AuthorizationSet opParams =
+            km::AuthorizationSetBuilder().Authorization(km::TAG_PURPOSE, km::KeyPurpose::ENCRYPT);
     km::AuthorizationSet outParams;
-    auto opHandle = BeginKeymasterOp(keymaster, dir, km::KeyPurpose::ENCRYPT, keyParams, opParams,
-                                     &outParams);
+    auto opHandle = BeginKeymasterOp(keymaster, dir, keyParams, opParams, &outParams);
     if (!opHandle) return false;
     auto nonceBlob = outParams.GetTagValue(km::TAG_NONCE);
-    if (!nonceBlob.isOk()) {
+    if (!nonceBlob) {
         LOG(ERROR) << "GCM encryption but no nonce generated";
         return false;
     }
     // nonceBlob here is just a pointer into existing data, must not be freed
-    std::string nonce(reinterpret_cast<const char*>(&nonceBlob.value()[0]),
-                      nonceBlob.value().size());
+    std::string nonce(nonceBlob.value().get().begin(), nonceBlob.value().get().end());
     if (!checkSize("nonce", nonce.size(), GCM_NONCE_BYTES)) return false;
     std::string body;
     if (!opHandle.updateCompletely(message, &body)) return false;
@@ -398,12 +394,12 @@
 static bool decryptWithKeymasterKey(Keymaster& keymaster, const std::string& dir,
                                     const km::AuthorizationSet& keyParams,
                                     const std::string& ciphertext, KeyBuffer* message) {
-    auto nonce = ciphertext.substr(0, GCM_NONCE_BYTES);
+    const std::string nonce = ciphertext.substr(0, GCM_NONCE_BYTES);
     auto bodyAndMac = ciphertext.substr(GCM_NONCE_BYTES);
-    auto opParams = km::AuthorizationSetBuilder().Authorization(km::TAG_NONCE,
-                                                                km::support::blob2hidlVec(nonce));
-    auto opHandle =
-            BeginKeymasterOp(keymaster, dir, km::KeyPurpose::DECRYPT, keyParams, opParams, nullptr);
+    auto opParams = km::AuthorizationSetBuilder()
+                            .Authorization(km::TAG_NONCE, nonce)
+                            .Authorization(km::TAG_PURPOSE, km::KeyPurpose::DECRYPT);
+    auto opHandle = BeginKeymasterOp(keymaster, dir, keyParams, opParams, nullptr);
     if (!opHandle) return false;
     if (!opHandle.updateCompletely(bodyAndMac, message)) return false;
     if (!opHandle.finish(nullptr)) return false;
diff --git a/Keymaster.cpp b/Keymaster.cpp
index e217ad4..5a68630 100644
--- a/Keymaster.cpp
+++ b/Keymaster.cpp
@@ -17,368 +17,235 @@
 #include "Keymaster.h"
 
 #include <android-base/logging.h>
-#include <keymasterV4_1/authorization_set.h>
-#include <keymasterV4_1/keymaster_utils.h>
+
+#include <aidl/android/hardware/security/keymint/SecurityLevel.h>
+#include <aidl/android/security/maintenance/IKeystoreMaintenance.h>
+#include <aidl/android/system/keystore2/Domain.h>
+#include <aidl/android/system/keystore2/KeyDescriptor.h>
+
+// Keep these in sync with system/security/keystore2/src/keystore2_main.rs
+static constexpr const char keystore2_service_name[] =
+        "android.system.keystore2.IKeystoreService/default";
+static constexpr const char maintenance_service_name[] = "android.security.maintenance";
+
+/*
+ * Keep this in sync with the description for update() in
+ * system/hardware/interfaces/keystore2/aidl/android/system/keystore2/IKeystoreOperation.aidl
+ */
+static constexpr const size_t UPDATE_INPUT_MAX_SIZE = 32 * 1024;  // 32 KiB
+
+// Keep this in sync with system/sepolicy/private/keystore2_key_contexts
+static constexpr const int VOLD_NAMESPACE = 100;
 
 namespace android {
 namespace vold {
 
-using ::android::hardware::hidl_string;
-using ::android::hardware::hidl_vec;
-using ::android::hardware::keymaster::V4_0::SecurityLevel;
+namespace ks2_maint = ::aidl::android::security::maintenance;
 
 KeymasterOperation::~KeymasterOperation() {
-    if (mDevice) mDevice->abort(mOpHandle);
+    if (ks2Operation) ks2Operation->abort();
+}
+
+static void zeroize_vector(std::vector<uint8_t>& vec) {
+    memset_s(vec.data(), 0, vec.size());
+}
+
+static bool logKeystore2ExceptionIfPresent(::ndk::ScopedAStatus& rc, const std::string& func_name) {
+    if (rc.isOk()) return false;
+
+    auto exception_code = rc.getExceptionCode();
+    if (exception_code == EX_SERVICE_SPECIFIC) {
+        LOG(ERROR) << "keystore2 Keystore " << func_name
+                   << " returned service specific error: " << rc.getServiceSpecificError();
+    } else {
+        LOG(ERROR) << "keystore2 Communication with Keystore " << func_name
+                   << " failed error: " << exception_code;
+    }
+    return true;
 }
 
 bool KeymasterOperation::updateCompletely(const char* input, size_t inputLen,
                                           const std::function<void(const char*, size_t)> consumer) {
-    uint32_t inputConsumed = 0;
+    if (!ks2Operation) return false;
 
-    km::ErrorCode km_error;
-    auto hidlCB = [&](km::ErrorCode ret, uint32_t inputConsumedDelta,
-                      const hidl_vec<km::KeyParameter>& /*ignored*/,
-                      const hidl_vec<uint8_t>& _output) {
-        km_error = ret;
-        if (km_error != km::ErrorCode::OK) return;
-        inputConsumed += inputConsumedDelta;
-        consumer(reinterpret_cast<const char*>(&_output[0]), _output.size());
-    };
+    while (inputLen != 0) {
+        size_t currLen = std::min(inputLen, UPDATE_INPUT_MAX_SIZE);
+        std::vector<uint8_t> input_vec(input, input + currLen);
+        inputLen -= currLen;
+        input += currLen;
 
-    while (inputConsumed != inputLen) {
-        size_t toRead = static_cast<size_t>(inputLen - inputConsumed);
-        auto inputBlob = km::support::blob2hidlVec(
-            reinterpret_cast<const uint8_t*>(&input[inputConsumed]), toRead);
-        auto error = mDevice->update(mOpHandle, hidl_vec<km::KeyParameter>(), inputBlob,
-                                     km::HardwareAuthToken(), km::VerificationToken(), hidlCB);
-        if (!error.isOk()) {
-            LOG(ERROR) << "update failed: " << error.description();
-            mDevice = nullptr;
+        std::optional<std::vector<uint8_t>> output;
+        auto rc = ks2Operation->update(input_vec, &output);
+        zeroize_vector(input_vec);
+        if (logKeystore2ExceptionIfPresent(rc, "update")) {
+            ks2Operation = nullptr;
             return false;
         }
-        if (km_error != km::ErrorCode::OK) {
-            LOG(ERROR) << "update failed, code " << int32_t(km_error);
-            mDevice = nullptr;
+
+        if (!output) {
+            LOG(ERROR) << "Keystore2 operation update didn't return output.";
+            ks2Operation = nullptr;
             return false;
         }
-        if (inputConsumed > inputLen) {
-            LOG(ERROR) << "update reported too much input consumed";
-            mDevice = nullptr;
-            return false;
-        }
+
+        consumer((const char*)output->data(), output->size());
     }
     return true;
 }
 
 bool KeymasterOperation::finish(std::string* output) {
-    km::ErrorCode km_error;
-    auto hidlCb = [&](km::ErrorCode ret, const hidl_vec<km::KeyParameter>& /*ignored*/,
-                      const hidl_vec<uint8_t>& _output) {
-        km_error = ret;
-        if (km_error != km::ErrorCode::OK) return;
-        if (output) output->assign(reinterpret_cast<const char*>(&_output[0]), _output.size());
-    };
-    auto error = mDevice->finish(mOpHandle, hidl_vec<km::KeyParameter>(), hidl_vec<uint8_t>(),
-                                 hidl_vec<uint8_t>(), km::HardwareAuthToken(),
-                                 km::VerificationToken(), hidlCb);
-    mDevice = nullptr;
-    if (!error.isOk()) {
-        LOG(ERROR) << "finish failed: " << error.description();
+    std::optional<std::vector<uint8_t>> out_vec;
+
+    if (!ks2Operation) return false;
+
+    auto rc = ks2Operation->finish(std::nullopt, std::nullopt, &out_vec);
+    if (logKeystore2ExceptionIfPresent(rc, "finish")) {
+        ks2Operation = nullptr;
         return false;
     }
-    if (km_error != km::ErrorCode::OK) {
-        LOG(ERROR) << "finish failed, code " << int32_t(km_error);
-        return false;
-    }
+
+    if (output) *output = std::string(out_vec->begin(), out_vec->end());
+
     return true;
 }
 
-/* static */ bool Keymaster::hmacKeyGenerated = false;
-
 Keymaster::Keymaster() {
-    auto devices = KmDevice::enumerateAvailableDevices();
-    if (!hmacKeyGenerated) {
-        KmDevice::performHmacKeyAgreement(devices);
-        hmacKeyGenerated = true;
+    ::ndk::SpAIBinder binder(AServiceManager_getService(keystore2_service_name));
+    auto keystore2Service = ks2::IKeystoreService::fromBinder(binder);
+
+    if (!keystore2Service) {
+        LOG(ERROR) << "Vold unable to connect to keystore2.";
+        return;
     }
-    for (auto& dev : devices) {
-        // Do not use StrongBox for device encryption / credential encryption.  If a security chip
-        // is present it will have Weaver, which already strengthens CE.  We get no additional
-        // benefit from using StrongBox here, so skip it.
-        if (dev->halVersion().securityLevel != SecurityLevel::STRONGBOX) {
-            mDevice = std::move(dev);
-            break;
-        }
-    }
-    if (!mDevice) return;
-    auto& version = mDevice->halVersion();
-    LOG(INFO) << "Using " << version.keymasterName << " from " << version.authorName
-              << " for encryption.  Security level: " << toString(version.securityLevel)
-              << ", HAL: " << mDevice->descriptor() << "/" << mDevice->instanceName();
+
+    /*
+     * There are only two options available to vold for the SecurityLevel: TRUSTED_ENVIRONMENT (TEE)
+     * and STRONGBOX. We don't use STRONGBOX because if a TEE is present it will have Weaver, which
+     * already strengthens CE, so there's no additional benefit from using StrongBox.
+     *
+     * The picture is slightly more complicated because Keystore2 reports a SOFTWARE instance as
+     * a TEE instance when there isn't a TEE instance available, but in that case, a STRONGBOX
+     * instance won't be available either, so we'll still be doing the best we can.
+     */
+    auto rc = keystore2Service->getSecurityLevel(km::SecurityLevel::TRUSTED_ENVIRONMENT,
+                                                 &securityLevel);
+    if (logKeystore2ExceptionIfPresent(rc, "getSecurityLevel"))
+        LOG(ERROR) << "Vold unable to get security level from keystore2.";
 }
 
 bool Keymaster::generateKey(const km::AuthorizationSet& inParams, std::string* key) {
-    km::ErrorCode km_error;
-    auto hidlCb = [&](km::ErrorCode ret, const hidl_vec<uint8_t>& keyBlob,
-                      const km::KeyCharacteristics& /*ignored*/) {
-        km_error = ret;
-        if (km_error != km::ErrorCode::OK) return;
-        if (key) key->assign(reinterpret_cast<const char*>(&keyBlob[0]), keyBlob.size());
+    ks2::KeyDescriptor in_key = {
+            .domain = ks2::Domain::BLOB,
+            .alias = std::nullopt,
+            .nspace = VOLD_NAMESPACE,
+            .blob = std::nullopt,
     };
+    ks2::KeyMetadata keyMetadata;
+    auto rc = securityLevel->generateKey(in_key, std::nullopt, inParams.vector_data(), 0, {},
+                                         &keyMetadata);
 
-    auto error = mDevice->generateKey(inParams.hidl_data(), hidlCb);
-    if (!error.isOk()) {
-        LOG(ERROR) << "generate_key failed: " << error.description();
+    if (logKeystore2ExceptionIfPresent(rc, "generateKey")) return false;
+
+    if (keyMetadata.key.blob == std::nullopt) {
+        LOG(ERROR) << "keystore2 generated key blob was null";
         return false;
     }
-    if (km_error != km::ErrorCode::OK) {
-        LOG(ERROR) << "generate_key failed, code " << int32_t(km_error);
-        return false;
-    }
+    if (key) *key = std::string(keyMetadata.key.blob->begin(), keyMetadata.key.blob->end());
+
+    zeroize_vector(keyMetadata.key.blob.value());
     return true;
 }
 
 bool Keymaster::exportKey(const KeyBuffer& kmKey, std::string* key) {
-    auto kmKeyBlob = km::support::blob2hidlVec(std::string(kmKey.data(), kmKey.size()));
-    km::ErrorCode km_error;
-    auto hidlCb = [&](km::ErrorCode ret, const hidl_vec<uint8_t>& exportedKeyBlob) {
-        km_error = ret;
-        if (km_error != km::ErrorCode::OK) return;
-        if (key)
-            key->assign(reinterpret_cast<const char*>(&exportedKeyBlob[0]), exportedKeyBlob.size());
+    bool ret = false;
+    ks2::KeyDescriptor storageKey = {
+            .domain = ks2::Domain::BLOB,
+            .alias = std::nullopt,
+            .nspace = VOLD_NAMESPACE,
     };
-    auto error = mDevice->exportKey(km::KeyFormat::RAW, kmKeyBlob, {}, {}, hidlCb);
-    if (!error.isOk()) {
-        LOG(ERROR) << "export_key failed: " << error.description();
-        return false;
-    }
-    if (km_error != km::ErrorCode::OK) {
-        LOG(ERROR) << "export_key failed, code " << int32_t(km_error);
-        return false;
-    }
-    return true;
+    storageKey.blob = std::make_optional<std::vector<uint8_t>>(kmKey.begin(), kmKey.end());
+    std::vector<uint8_t> ephemeral_key;
+    auto rc = securityLevel->convertStorageKeyToEphemeral(storageKey, &ephemeral_key);
+
+    if (logKeystore2ExceptionIfPresent(rc, "exportKey")) goto out;
+    if (key) *key = std::string(ephemeral_key.begin(), ephemeral_key.end());
+
+    ret = true;
+out:
+    zeroize_vector(ephemeral_key);
+    zeroize_vector(storageKey.blob.value());
+    return ret;
 }
 
 bool Keymaster::deleteKey(const std::string& key) {
-    auto keyBlob = km::support::blob2hidlVec(key);
-    auto error = mDevice->deleteKey(keyBlob);
-    if (!error.isOk()) {
-        LOG(ERROR) << "delete_key failed: " << error.description();
-        return false;
-    }
-    if (error != km::ErrorCode::OK) {
-        LOG(ERROR) << "delete_key failed, code " << int32_t(km::ErrorCode(error));
-        return false;
-    }
-    return true;
-}
-
-bool Keymaster::upgradeKey(const std::string& oldKey, const km::AuthorizationSet& inParams,
-                           std::string* newKey) {
-    auto oldKeyBlob = km::support::blob2hidlVec(oldKey);
-    km::ErrorCode km_error;
-    auto hidlCb = [&](km::ErrorCode ret, const hidl_vec<uint8_t>& upgradedKeyBlob) {
-        km_error = ret;
-        if (km_error != km::ErrorCode::OK) return;
-        if (newKey)
-            newKey->assign(reinterpret_cast<const char*>(&upgradedKeyBlob[0]),
-                           upgradedKeyBlob.size());
+    ks2::KeyDescriptor keyDesc = {
+            .domain = ks2::Domain::BLOB,
+            .alias = std::nullopt,
+            .nspace = VOLD_NAMESPACE,
     };
-    auto error = mDevice->upgradeKey(oldKeyBlob, inParams.hidl_data(), hidlCb);
-    if (!error.isOk()) {
-        LOG(ERROR) << "upgrade_key failed: " << error.description();
-        return false;
-    }
-    if (km_error != km::ErrorCode::OK) {
-        LOG(ERROR) << "upgrade_key failed, code " << int32_t(km_error);
-        return false;
-    }
-    return true;
+    keyDesc.blob =
+            std::optional<std::vector<uint8_t>>(std::vector<uint8_t>(key.begin(), key.end()));
+
+    auto rc = securityLevel->deleteKey(keyDesc);
+    return !logKeystore2ExceptionIfPresent(rc, "deleteKey");
 }
 
-KeymasterOperation Keymaster::begin(km::KeyPurpose purpose, const std::string& key,
-                                    const km::AuthorizationSet& inParams,
+KeymasterOperation Keymaster::begin(const std::string& key, const km::AuthorizationSet& inParams,
                                     km::AuthorizationSet* outParams) {
-    auto keyBlob = km::support::blob2hidlVec(key);
-    uint64_t mOpHandle;
-    km::ErrorCode km_error;
-
-    auto hidlCb = [&](km::ErrorCode ret, const hidl_vec<km::KeyParameter>& _outParams,
-                      uint64_t operationHandle) {
-        km_error = ret;
-        if (km_error != km::ErrorCode::OK) return;
-        if (outParams) *outParams = _outParams;
-        mOpHandle = operationHandle;
+    ks2::KeyDescriptor keyDesc = {
+            .domain = ks2::Domain::BLOB,
+            .alias = std::nullopt,
+            .nspace = VOLD_NAMESPACE,
     };
+    keyDesc.blob =
+            std::optional<std::vector<uint8_t>>(std::vector<uint8_t>(key.begin(), key.end()));
 
-    auto error =
-            mDevice->begin(purpose, keyBlob, inParams.hidl_data(), km::HardwareAuthToken(), hidlCb);
-    if (!error.isOk()) {
-        LOG(ERROR) << "begin failed: " << error.description();
-        return KeymasterOperation(km::ErrorCode::UNKNOWN_ERROR);
+    ks2::CreateOperationResponse cor;
+    auto rc = securityLevel->createOperation(keyDesc, inParams.vector_data(), true, &cor);
+    if (logKeystore2ExceptionIfPresent(rc, "createOperation")) {
+        if (rc.getExceptionCode() == EX_SERVICE_SPECIFIC)
+            return KeymasterOperation((km::ErrorCode)rc.getServiceSpecificError());
+        else
+            return KeymasterOperation();
     }
-    if (km_error != km::ErrorCode::OK) {
-        LOG(ERROR) << "begin failed, code " << int32_t(km_error);
-        return KeymasterOperation(km_error);
+
+    if (!cor.iOperation) {
+        LOG(ERROR) << "keystore2 createOperation didn't return an operation";
+        return KeymasterOperation();
     }
-    return KeymasterOperation(mDevice.get(), mOpHandle);
+
+    if (outParams && cor.parameters) *outParams = cor.parameters->keyParameter;
+
+    return KeymasterOperation(cor.iOperation, cor.upgradedBlob);
 }
 
 bool Keymaster::isSecure() {
-    return mDevice->halVersion().securityLevel != km::SecurityLevel::SOFTWARE;
+    return true;
 }
 
 void Keymaster::earlyBootEnded() {
-    auto devices = KmDevice::enumerateAvailableDevices();
-    for (auto& dev : devices) {
-        auto error = dev->earlyBootEnded();
-        if (!error.isOk()) {
-            LOG(ERROR) << "earlyBootEnded call failed: " << error.description() << " for "
-                       << dev->halVersion().keymasterName;
-        }
-        km::V4_1_ErrorCode km_error = error;
-        if (km_error != km::V4_1_ErrorCode::OK && km_error != km::V4_1_ErrorCode::UNIMPLEMENTED) {
-            LOG(ERROR) << "Error reporting early boot ending to keymaster: "
-                       << static_cast<int32_t>(km_error) << " for "
-                       << dev->halVersion().keymasterName;
-        }
+    ::ndk::SpAIBinder binder(AServiceManager_getService(maintenance_service_name));
+    auto maint_service = ks2_maint::IKeystoreMaintenance::fromBinder(binder);
+
+    if (!maint_service) {
+        LOG(ERROR) << "Unable to connect to keystore2 maintenance service for earlyBootEnded";
+        return;
     }
+
+    auto rc = maint_service->earlyBootEnded();
+    logKeystore2ExceptionIfPresent(rc, "earlyBootEnded");
 }
 
 }  // namespace vold
 }  // namespace android
 
-using namespace ::android::vold;
-
+// TODO: This always returns true right now since we hardcode the security level.
+// If it's alright to hardcode it, we should remove this function and simplify the callers.
 int keymaster_compatibility_cryptfs_scrypt() {
-    Keymaster dev;
+    android::vold::Keymaster dev;
     if (!dev) {
         LOG(ERROR) << "Failed to initiate keymaster session";
         return -1;
     }
     return dev.isSecure();
 }
-
-static bool write_string_to_buf(const std::string& towrite, uint8_t* buffer, uint32_t buffer_size,
-                                uint32_t* out_size) {
-    if (!buffer || !out_size) {
-        LOG(ERROR) << "Missing target pointers";
-        return false;
-    }
-    *out_size = towrite.size();
-    if (buffer_size < towrite.size()) {
-        LOG(ERROR) << "Buffer too small " << buffer_size << " < " << towrite.size();
-        return false;
-    }
-    memset(buffer, '\0', buffer_size);
-    std::copy(towrite.begin(), towrite.end(), buffer);
-    return true;
-}
-
-static km::AuthorizationSet keyParams(uint32_t rsa_key_size, uint64_t rsa_exponent,
-                                      uint32_t ratelimit) {
-    return km::AuthorizationSetBuilder()
-        .RsaSigningKey(rsa_key_size, rsa_exponent)
-        .NoDigestOrPadding()
-        .Authorization(km::TAG_BLOB_USAGE_REQUIREMENTS, km::KeyBlobUsageRequirements::STANDALONE)
-        .Authorization(km::TAG_NO_AUTH_REQUIRED)
-        .Authorization(km::TAG_MIN_SECONDS_BETWEEN_OPS, ratelimit);
-}
-
-int keymaster_create_key_for_cryptfs_scrypt(uint32_t rsa_key_size, uint64_t rsa_exponent,
-                                            uint32_t ratelimit, uint8_t* key_buffer,
-                                            uint32_t key_buffer_size, uint32_t* key_out_size) {
-    if (key_out_size) {
-        *key_out_size = 0;
-    }
-    Keymaster dev;
-    if (!dev) {
-        LOG(ERROR) << "Failed to initiate keymaster session";
-        return -1;
-    }
-    std::string key;
-    if (!dev.generateKey(keyParams(rsa_key_size, rsa_exponent, ratelimit), &key)) return -1;
-    if (!write_string_to_buf(key, key_buffer, key_buffer_size, key_out_size)) return -1;
-    return 0;
-}
-
-int keymaster_upgrade_key_for_cryptfs_scrypt(uint32_t rsa_key_size, uint64_t rsa_exponent,
-                                             uint32_t ratelimit, const uint8_t* key_blob,
-                                             size_t key_blob_size, uint8_t* key_buffer,
-                                             uint32_t key_buffer_size, uint32_t* key_out_size) {
-    if (key_out_size) {
-        *key_out_size = 0;
-    }
-    Keymaster dev;
-    if (!dev) {
-        LOG(ERROR) << "Failed to initiate keymaster session";
-        return -1;
-    }
-    std::string old_key(reinterpret_cast<const char*>(key_blob), key_blob_size);
-    std::string new_key;
-    if (!dev.upgradeKey(old_key, keyParams(rsa_key_size, rsa_exponent, ratelimit), &new_key))
-        return -1;
-    if (!write_string_to_buf(new_key, key_buffer, key_buffer_size, key_out_size)) return -1;
-    return 0;
-}
-
-KeymasterSignResult keymaster_sign_object_for_cryptfs_scrypt(
-    const uint8_t* key_blob, size_t key_blob_size, uint32_t ratelimit, const uint8_t* object,
-    const size_t object_size, uint8_t** signature_buffer, size_t* signature_buffer_size) {
-    Keymaster dev;
-    if (!dev) {
-        LOG(ERROR) << "Failed to initiate keymaster session";
-        return KeymasterSignResult::error;
-    }
-    if (!key_blob || !object || !signature_buffer || !signature_buffer_size) {
-        LOG(ERROR) << __FILE__ << ":" << __LINE__ << ":Invalid argument";
-        return KeymasterSignResult::error;
-    }
-
-    km::AuthorizationSet outParams;
-    std::string key(reinterpret_cast<const char*>(key_blob), key_blob_size);
-    std::string input(reinterpret_cast<const char*>(object), object_size);
-    std::string output;
-    KeymasterOperation op;
-
-    auto paramBuilder = km::AuthorizationSetBuilder().NoDigestOrPadding();
-    while (true) {
-        op = dev.begin(km::KeyPurpose::SIGN, key, paramBuilder, &outParams);
-        if (op.errorCode() == km::ErrorCode::KEY_RATE_LIMIT_EXCEEDED) {
-            sleep(ratelimit);
-            continue;
-        } else
-            break;
-    }
-
-    if (op.errorCode() == km::ErrorCode::KEY_REQUIRES_UPGRADE) {
-        LOG(ERROR) << "Keymaster key requires upgrade";
-        return KeymasterSignResult::upgrade;
-    }
-
-    if (op.errorCode() != km::ErrorCode::OK) {
-        LOG(ERROR) << "Error starting keymaster signature transaction: " << int32_t(op.errorCode());
-        return KeymasterSignResult::error;
-    }
-
-    if (!op.updateCompletely(input, &output)) {
-        LOG(ERROR) << "Error sending data to keymaster signature transaction: "
-                   << uint32_t(op.errorCode());
-        return KeymasterSignResult::error;
-    }
-
-    if (!op.finish(&output)) {
-        LOG(ERROR) << "Error finalizing keymaster signature transaction: "
-                   << int32_t(op.errorCode());
-        return KeymasterSignResult::error;
-    }
-
-    *signature_buffer = reinterpret_cast<uint8_t*>(malloc(output.size()));
-    if (*signature_buffer == nullptr) {
-        LOG(ERROR) << "Error allocation buffer for keymaster signature";
-        return KeymasterSignResult::error;
-    }
-    *signature_buffer_size = output.size();
-    std::copy(output.data(), output.data() + output.size(), *signature_buffer);
-    return KeymasterSignResult::ok;
-}
diff --git a/Keymaster.h b/Keymaster.h
index ccf018e..84b473e 100644
--- a/Keymaster.h
+++ b/Keymaster.h
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-
+// TODO: Maybe "Keymaster" should be replaced with Keystore2 everywhere?
 #ifndef ANDROID_VOLD_KEYMASTER_H
 #define ANDROID_VOLD_KEYMASTER_H
 
@@ -24,33 +24,25 @@
 #include <utility>
 
 #include <android-base/macros.h>
-#include <keymasterV4_1/Keymaster.h>
-#include <keymasterV4_1/authorization_set.h>
+#include <keymint_support/authorization_set.h>
+#include <keymint_support/keymint_tags.h>
+
+#include <aidl/android/hardware/security/keymint/ErrorCode.h>
+#include <aidl/android/system/keystore2/IKeystoreService.h>
+#include <android/binder_manager.h>
 
 namespace android {
 namespace vold {
 
-namespace km {
+namespace ks2 = ::aidl::android::system::keystore2;
+namespace km = ::aidl::android::hardware::security::keymint;
 
-using namespace ::android::hardware::keymaster::V4_1;
-
-// Surprisingly -- to me, at least -- this is totally fine.  You can re-define symbols that were
-// brought in via a using directive (the "using namespace") above.  In general this seems like a
-// dangerous thing to rely on, but in this case its implications are simple and straightforward:
-// km::ErrorCode refers to the 4.0 ErrorCode, though we pull everything else from 4.1.
-using ErrorCode = ::android::hardware::keymaster::V4_0::ErrorCode;
-using V4_1_ErrorCode = ::android::hardware::keymaster::V4_1::ErrorCode;
-
-}  // namespace km
-
-using KmDevice = km::support::Keymaster;
-
-// C++ wrappers to the Keymaster hidl interface.
+// C++ wrappers to the Keystore2 AIDL interface.
 // This is tailored to the needs of KeyStorage, but could be extended to be
 // a more general interface.
 
-// Wrapper for a Keymaster operation handle representing an
-// ongoing Keymaster operation.  Aborts the operation
+// Wrapper for a Keystore2 operation handle representing an
+// ongoing Keystore2 operation.  Aborts the operation
 // in the destructor if it is unfinished. Methods log failures
 // to LOG(ERROR).
 class KeymasterOperation {
@@ -58,8 +50,9 @@
     ~KeymasterOperation();
     // Is this instance valid? This is false if creation fails, and becomes
     // false on finish or if an update fails.
-    explicit operator bool() const { return mError == km::ErrorCode::OK; }
-    km::ErrorCode errorCode() const { return mError; }
+    explicit operator bool() const { return (bool)ks2Operation; }
+    km::ErrorCode getErrorCode() const { return errorCode; }
+    std::optional<std::string> getUpgradedBlob() const { return upgradedBlob; }
     // Call "update" repeatedly until all of the input is consumed, and
     // concatenate the output. Return true on success.
     template <class TI, class TO>
@@ -75,102 +68,74 @@
     // Move constructor
     KeymasterOperation(KeymasterOperation&& rhs) { *this = std::move(rhs); }
     // Construct an object in an error state for error returns
-    KeymasterOperation() : mDevice{nullptr}, mOpHandle{0}, mError{km::ErrorCode::UNKNOWN_ERROR} {}
+    KeymasterOperation() { errorCode = km::ErrorCode::UNKNOWN_ERROR; }
     // Move Assignment
     KeymasterOperation& operator=(KeymasterOperation&& rhs) {
-        mDevice = rhs.mDevice;
-        rhs.mDevice = nullptr;
+        ks2Operation = rhs.ks2Operation;
+        rhs.ks2Operation = nullptr;
 
-        mOpHandle = rhs.mOpHandle;
-        rhs.mOpHandle = 0;
+        upgradedBlob = rhs.upgradedBlob;
+        rhs.upgradedBlob = std::nullopt;
 
-        mError = rhs.mError;
-        rhs.mError = km::ErrorCode::UNKNOWN_ERROR;
+        errorCode = rhs.errorCode;
+        rhs.errorCode = km::ErrorCode::UNKNOWN_ERROR;
 
         return *this;
     }
 
   private:
-    KeymasterOperation(KmDevice* d, uint64_t h)
-        : mDevice{d}, mOpHandle{h}, mError{km::ErrorCode::OK} {}
-    KeymasterOperation(km::ErrorCode error) : mDevice{nullptr}, mOpHandle{0}, mError{error} {}
+    KeymasterOperation(std::shared_ptr<ks2::IKeystoreOperation> ks2Op,
+                       std::optional<std::vector<uint8_t>> blob)
+        : ks2Operation{ks2Op}, errorCode{km::ErrorCode::OK} {
+        if (blob)
+            upgradedBlob = std::optional(std::string(blob->begin(), blob->end()));
+        else
+            upgradedBlob = std::nullopt;
+    }
+
+    KeymasterOperation(km::ErrorCode errCode) : errorCode{errCode} {}
 
     bool updateCompletely(const char* input, size_t inputLen,
                           const std::function<void(const char*, size_t)> consumer);
 
-    KmDevice* mDevice;
-    uint64_t mOpHandle;
-    km::ErrorCode mError;
+    std::shared_ptr<ks2::IKeystoreOperation> ks2Operation;
+    std::optional<std::string> upgradedBlob;
+    km::ErrorCode errorCode;
     DISALLOW_COPY_AND_ASSIGN(KeymasterOperation);
     friend class Keymaster;
 };
 
-// Wrapper for a Keymaster device for methods that start a KeymasterOperation or are not
-// part of one.
+// Wrapper for keystore2 methods that vold uses.
 class Keymaster {
   public:
     Keymaster();
-    // false if we failed to open the keymaster device.
-    explicit operator bool() { return mDevice.get() != nullptr; }
-    // Generate a key in the keymaster from the given params.
+    // false if we failed to get a keystore2 security level.
+    explicit operator bool() { return (bool)securityLevel; }
+    // Generate a key using keystore2 from the given params.
     bool generateKey(const km::AuthorizationSet& inParams, std::string* key);
-    // Exports a keymaster key with STORAGE_KEY tag wrapped with a per-boot ephemeral key
+    // Exports a keystore2 key with STORAGE_KEY tag wrapped with a per-boot ephemeral key
     bool exportKey(const KeyBuffer& kmKey, std::string* key);
-    // If the keymaster supports it, permanently delete a key.
+    // If supported, permanently delete a key from the keymint device it belongs to.
     bool deleteKey(const std::string& key);
-    // Replace stored key blob in response to KM_ERROR_KEY_REQUIRES_UPGRADE.
-    bool upgradeKey(const std::string& oldKey, const km::AuthorizationSet& inParams,
-                    std::string* newKey);
     // Begin a new cryptographic operation, collecting output parameters if pointer is non-null
-    KeymasterOperation begin(km::KeyPurpose purpose, const std::string& key,
-                             const km::AuthorizationSet& inParams,
+    // If the key was upgraded as a result of a call to this method, the returned KeymasterOperation
+    // also stores the upgraded key blob.
+    KeymasterOperation begin(const std::string& key, const km::AuthorizationSet& inParams,
                              km::AuthorizationSet* outParams);
     bool isSecure();
 
-    // Tell all Keymaster instances that early boot has ended and early boot-only keys can no longer
+    // Tell all Keymint devices that early boot has ended and early boot-only keys can no longer
     // be created or used.
     static void earlyBootEnded();
 
   private:
-    sp<KmDevice> mDevice;
+    std::shared_ptr<ks2::IKeystoreSecurityLevel> securityLevel;
     DISALLOW_COPY_AND_ASSIGN(Keymaster);
-    static bool hmacKeyGenerated;
 };
 
 }  // namespace vold
 }  // namespace android
 
-// FIXME no longer needed now cryptfs is in C++.
-
-/*
- * The following functions provide C bindings to keymaster services
- * needed by cryptfs scrypt. The compatibility check checks whether
- * the keymaster implementation is considered secure, i.e., TEE backed.
- * The create_key function generates an RSA key for signing.
- * The sign_object function signes an object with the given keymaster
- * key.
- */
-
-/* Return values for keymaster_sign_object_for_cryptfs_scrypt */
-
-enum class KeymasterSignResult {
-    ok = 0,
-    error = -1,
-    upgrade = -2,
-};
-
 int keymaster_compatibility_cryptfs_scrypt();
-int keymaster_create_key_for_cryptfs_scrypt(uint32_t rsa_key_size, uint64_t rsa_exponent,
-                                            uint32_t ratelimit, uint8_t* key_buffer,
-                                            uint32_t key_buffer_size, uint32_t* key_out_size);
-
-int keymaster_upgrade_key_for_cryptfs_scrypt(uint32_t rsa_key_size, uint64_t rsa_exponent,
-                                             uint32_t ratelimit, const uint8_t* key_blob,
-                                             size_t key_blob_size, uint8_t* key_buffer,
-                                             uint32_t key_buffer_size, uint32_t* key_out_size);
-
-KeymasterSignResult keymaster_sign_object_for_cryptfs_scrypt(
-    const uint8_t* key_blob, size_t key_blob_size, uint32_t ratelimit, const uint8_t* object,
-    const size_t object_size, uint8_t** signature_buffer, size_t* signature_buffer_size);
 
 #endif
diff --git a/cryptfs.cpp b/cryptfs.cpp
index 6203003..deba6da 100644
--- a/cryptfs.cpp
+++ b/cryptfs.cpp
@@ -29,6 +29,7 @@
 #include "VoldUtil.h"
 #include "VolumeManager.h"
 
+#include <android-base/logging.h>
 #include <android-base/parseint.h>
 #include <android-base/properties.h>
 #include <android-base/stringprintf.h>
@@ -80,6 +81,7 @@
 using android::vold::CryptoType;
 using android::vold::KeyBuffer;
 using android::vold::KeyGeneration;
+using namespace android::vold;
 using namespace android::dm;
 using namespace std::chrono_literals;
 
@@ -331,6 +333,45 @@
     return keymaster_compatibility_cryptfs_scrypt();
 }
 
+static bool write_string_to_buf(const std::string& towrite, uint8_t* buffer, uint32_t buffer_size,
+                                uint32_t* out_size) {
+    if (!buffer || !out_size) {
+        LOG(ERROR) << "Missing target pointers";
+        return false;
+    }
+    *out_size = towrite.size();
+    if (buffer_size < towrite.size()) {
+        LOG(ERROR) << "Buffer too small " << buffer_size << " < " << towrite.size();
+        return false;
+    }
+    memset(buffer, '\0', buffer_size);
+    std::copy(towrite.begin(), towrite.end(), buffer);
+    return true;
+}
+
+static int keymaster_create_key_for_cryptfs_scrypt(uint32_t rsa_key_size, uint64_t rsa_exponent,
+                                                   uint32_t ratelimit, uint8_t* key_buffer,
+                                                   uint32_t key_buffer_size,
+                                                   uint32_t* key_out_size) {
+    if (key_out_size) {
+        *key_out_size = 0;
+    }
+    Keymaster dev;
+    if (!dev) {
+        LOG(ERROR) << "Failed to initiate keymaster session";
+        return -1;
+    }
+    auto keyParams = km::AuthorizationSetBuilder()
+                             .RsaSigningKey(rsa_key_size, rsa_exponent)
+                             .NoDigestOrPadding()
+                             .Authorization(km::TAG_NO_AUTH_REQUIRED)
+                             .Authorization(km::TAG_MIN_SECONDS_BETWEEN_OPS, ratelimit);
+    std::string key;
+    if (!dev.generateKey(keyParams, &key)) return -1;
+    if (!write_string_to_buf(key, key_buffer, key_buffer_size, key_out_size)) return -1;
+    return 0;
+}
+
 /* Create a new keymaster key and store it in this footer */
 static int keymaster_create_key(struct crypt_mnt_ftr* ftr) {
     if (ftr->keymaster_blob_size) {
@@ -352,6 +393,79 @@
     return 0;
 }
 
+static int keymaster_sign_object_for_cryptfs_scrypt(struct crypt_mnt_ftr* ftr, uint32_t ratelimit,
+                                                    const uint8_t* object, const size_t object_size,
+                                                    uint8_t** signature_buffer,
+                                                    size_t* signature_buffer_size) {
+    if (!object || !signature_buffer || !signature_buffer_size) {
+        LOG(ERROR) << __FILE__ << ":" << __LINE__ << ":Invalid argument";
+        return -1;
+    }
+
+    Keymaster dev;
+    if (!dev) {
+        LOG(ERROR) << "Failed to initiate keymaster session";
+        return -1;
+    }
+
+    km::AuthorizationSet outParams;
+    std::string key(reinterpret_cast<const char*>(ftr->keymaster_blob), ftr->keymaster_blob_size);
+    std::string input(reinterpret_cast<const char*>(object), object_size);
+    std::string output;
+    KeymasterOperation op;
+
+    auto paramBuilder = km::AuthorizationSetBuilder().NoDigestOrPadding().Authorization(
+            km::TAG_PURPOSE, km::KeyPurpose::SIGN);
+    while (true) {
+        op = dev.begin(key, paramBuilder, &outParams);
+        if (op.getErrorCode() == km::ErrorCode::KEY_RATE_LIMIT_EXCEEDED) {
+            sleep(ratelimit);
+            continue;
+        } else
+            break;
+    }
+
+    if (!op) {
+        LOG(ERROR) << "Error starting keymaster signature transaction: "
+                   << int32_t(op.getErrorCode());
+        return -1;
+    }
+
+    if (op.getUpgradedBlob()) {
+        write_string_to_buf(*op.getUpgradedBlob(), ftr->keymaster_blob, KEYMASTER_BLOB_SIZE,
+                            &ftr->keymaster_blob_size);
+
+        SLOGD("Upgrading key");
+        if (put_crypt_ftr_and_key(ftr) != 0) {
+            SLOGE("Failed to write upgraded key to disk");
+            return -1;
+        }
+        SLOGD("Key upgraded successfully");
+    }
+
+    if (!op.updateCompletely(input, &output)) {
+        LOG(ERROR) << "Error sending data to keymaster signature transaction: "
+                   << int32_t(op.getErrorCode());
+        return -1;
+    }
+
+    if (!op.finish(&output)) {
+        LOG(ERROR) << "Error finalizing keymaster signature transaction: "
+                   << int32_t(op.getErrorCode());
+        return -1;
+    }
+
+    *signature_buffer = reinterpret_cast<uint8_t*>(malloc(output.size()));
+    if (*signature_buffer == nullptr) {
+        LOG(ERROR) << "Error allocation buffer for keymaster signature";
+        return -1;
+    }
+    *signature_buffer_size = output.size();
+    std::copy(output.data(), output.data() + output.size(), *signature_buffer);
+
+    return 0;
+}
+
 /* This signs the given object using the keymaster key. */
 static int keymaster_sign_object(struct crypt_mnt_ftr* ftr, const unsigned char* object,
                                  const size_t object_size, unsigned char** signature,
@@ -389,31 +503,8 @@
             SLOGE("Unknown KDF type %d", ftr->kdf_type);
             return -1;
     }
-    for (;;) {
-        auto result = keymaster_sign_object_for_cryptfs_scrypt(
-            ftr->keymaster_blob, ftr->keymaster_blob_size, KEYMASTER_CRYPTFS_RATE_LIMIT, to_sign,
-            to_sign_size, signature, signature_size);
-        switch (result) {
-            case KeymasterSignResult::ok:
-                return 0;
-            case KeymasterSignResult::upgrade:
-                break;
-            default:
-                return -1;
-        }
-        SLOGD("Upgrading key");
-        if (keymaster_upgrade_key_for_cryptfs_scrypt(
-                RSA_KEY_SIZE, RSA_EXPONENT, KEYMASTER_CRYPTFS_RATE_LIMIT, ftr->keymaster_blob,
-                ftr->keymaster_blob_size, ftr->keymaster_blob, KEYMASTER_BLOB_SIZE,
-                &ftr->keymaster_blob_size) != 0) {
-            SLOGE("Failed to upgrade key");
-            return -1;
-        }
-        if (put_crypt_ftr_and_key(ftr) != 0) {
-            SLOGE("Failed to write upgraded key to disk");
-        }
-        SLOGD("Key upgraded successfully");
-    }
+    return keymaster_sign_object_for_cryptfs_scrypt(ftr, KEYMASTER_CRYPTFS_RATE_LIMIT, to_sign,
+                                                    to_sign_size, signature, signature_size);
 }
 
 /* Store password when userdata is successfully decrypted and mounted.