diff --git a/identity/Android.bp b/identity/Android.bp
index ed8ff2f..d66f4ec 100644
--- a/identity/Android.bp
+++ b/identity/Android.bp
@@ -39,13 +39,16 @@
     shared_libs: [
         "libbase",
         "libbinder",
-        "libkeystore_aidl",
+        "libbinder_ndk",
+        "android.hardware.keymaster@4.0",
         "libcredstore_aidl",
         "libutils",
         "libhidlbase",
         "android.hardware.identity-support-lib",
         "libkeymaster4support",
         "libkeystore-attestation-application-id",
+        "android.hardware.security.keymint-V1-ndk_platform",
+        "android.security.authorization-ndk_platform",
     ],
     static_libs: [
         "android.hardware.identity-V3-cpp",
diff --git a/identity/Credential.cpp b/identity/Credential.cpp
index 4a2bae1..2e6b9c1 100644
--- a/identity/Credential.cpp
+++ b/identity/Credential.cpp
@@ -17,13 +17,11 @@
 #define LOG_TAG "Credential"
 
 #include <android-base/logging.h>
-
+#include <android/binder_manager.h>
 #include <android/hardware/identity/support/IdentityCredentialSupport.h>
 
 #include <android/security/identity/ICredentialStore.h>
 
-#include <android/security/keystore/BnCredstoreTokenCallback.h>
-#include <android/security/keystore/IKeystoreService.h>
 #include <binder/IPCThreadState.h>
 #include <binder/IServiceManager.h>
 #include <keymasterV4_0/keymaster_utils.h>
@@ -33,6 +31,11 @@
 #include <future>
 #include <tuple>
 
+#include <aidl/android/hardware/security/keymint/HardwareAuthToken.h>
+#include <aidl/android/hardware/security/secureclock/TimeStampToken.h>
+#include <aidl/android/security/authorization/AuthorizationTokens.h>
+#include <aidl/android/security/authorization/IKeystoreAuthorization.h>
+
 #include "Credential.h"
 #include "CredentialData.h"
 #include "Util.h"
@@ -46,8 +49,6 @@
 using std::promise;
 using std::tuple;
 
-using android::security::keystore::IKeystoreService;
-
 using ::android::hardware::identity::IWritableIdentityCredential;
 
 using ::android::hardware::identity::support::ecKeyPairGetPkcs12;
@@ -55,11 +56,17 @@
 using ::android::hardware::identity::support::ecKeyPairGetPublicKey;
 using ::android::hardware::identity::support::sha256;
 
+using android::hardware::keymaster::SecurityLevel;
 using android::hardware::keymaster::V4_0::HardwareAuthToken;
 using android::hardware::keymaster::V4_0::VerificationToken;
 using AidlHardwareAuthToken = android::hardware::keymaster::HardwareAuthToken;
 using AidlVerificationToken = android::hardware::keymaster::VerificationToken;
 
+using KeyMintAuthToken = ::aidl::android::hardware::security::keymint::HardwareAuthToken;
+using ::aidl::android::hardware::security::secureclock::TimeStampToken;
+using ::aidl::android::security::authorization::AuthorizationTokens;
+using ::aidl::android::security::authorization::IKeystoreAuthorization;
+
 Credential::Credential(CipherSuite cipherSuite, const std::string& dataPath,
                        const std::string& credentialName, uid_t callingUid,
                        HardwareInformation hwInfo, sp<IIdentityCredentialStore> halStoreBinder,
@@ -117,73 +124,94 @@
                                                 "Error loading data for credential");
     }
 
-    selectedAuthKey_ = data->selectAuthKey(allowUsingExhaustedKeys, allowUsingExpiredKeys);
-    if (selectedAuthKey_ == nullptr) {
+    // We just check if a key is available, we actually don't store it since we
+    // don't keep CredentialData around between binder calls.
+    const AuthKeyData* authKey =
+        data->selectAuthKey(allowUsingExhaustedKeys, allowUsingExpiredKeys);
+    if (authKey == nullptr) {
         return Status::fromServiceSpecificError(
             ICredentialStore::ERROR_NO_AUTHENTICATION_KEY_AVAILABLE,
             "No suitable authentication key available");
     }
 
-    int64_t challenge;
-    Status status = halBinder_->createAuthChallenge(&challenge);
-    if (!status.isOk()) {
-        return halStatusToGenericError(status);
-    }
-    if (challenge == 0) {
+    if (!ensureChallenge()) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
-                                                "Returned challenge is 0 (bug in HAL or TA)");
+                                                "Error getting challenge (bug in HAL or TA)");
     }
-
-    selectedChallenge_ = challenge;
-    *_aidl_return = challenge;
+    *_aidl_return = selectedChallenge_;
     return Status::ok();
 }
 
-class CredstoreTokenCallback : public android::security::keystore::BnCredstoreTokenCallback,
-                               public promise<tuple<bool, vector<uint8_t>, vector<uint8_t>>> {
-  public:
-    CredstoreTokenCallback() {}
-    virtual Status onFinished(bool success, const vector<uint8_t>& authToken,
-                              const vector<uint8_t>& verificationToken) override {
-        this->set_value({success, authToken, verificationToken});
-        return Status::ok();
+bool Credential::ensureChallenge() {
+    if (selectedChallenge_ != 0) {
+        return true;
     }
-};
+
+    int64_t challenge;
+    Status status = halBinder_->createAuthChallenge(&challenge);
+    if (!status.isOk()) {
+        LOG(ERROR) << "Error getting challenge: " << status.exceptionMessage();
+        return false;
+    }
+    if (challenge == 0) {
+        LOG(ERROR) << "Returned challenge is 0 (bug in HAL or TA)";
+        return false;
+    }
+
+    selectedChallenge_ = challenge;
+    return true;
+}
 
 // Returns false if an error occurred communicating with keystore.
 //
-bool getTokensFromKeystore(uint64_t challenge, uint64_t secureUserId,
-                           unsigned int authTokenMaxAgeMillis, vector<uint8_t>& authToken,
-                           vector<uint8_t>& verificationToken) {
-    sp<IServiceManager> sm = defaultServiceManager();
-    sp<IBinder> binder = sm->getService(String16("android.security.keystore"));
-    sp<IKeystoreService> keystore = interface_cast<IKeystoreService>(binder);
-    if (keystore == nullptr) {
-        return false;
-    }
+bool getTokensFromKeystore2(uint64_t challenge, uint64_t secureUserId,
+                            unsigned int authTokenMaxAgeMillis,
+                            AidlHardwareAuthToken& aidlAuthToken,
+                            AidlVerificationToken& aidlVerificationToken) {
+    // try to connect to IKeystoreAuthorization AIDL service first.
+    AIBinder* authzAIBinder = AServiceManager_checkService("android.security.authorization");
+    ::ndk::SpAIBinder authzBinder(authzAIBinder);
+    auto authzService = IKeystoreAuthorization::fromBinder(authzBinder);
+    if (authzService) {
+        AuthorizationTokens authzTokens;
+        auto result = authzService->getAuthTokensForCredStore(challenge, secureUserId,
+                                                              authTokenMaxAgeMillis, &authzTokens);
+        // Convert KeyMint auth token to KeyMaster authtoken, only if tokens are
+        // returned
+        if (result.isOk()) {
+            KeyMintAuthToken keymintAuthToken = authzTokens.authToken;
+            aidlAuthToken.challenge = keymintAuthToken.challenge;
+            aidlAuthToken.userId = keymintAuthToken.userId;
+            aidlAuthToken.authenticatorId = keymintAuthToken.authenticatorId;
+            aidlAuthToken.authenticatorType =
+                ::android::hardware::keymaster::HardwareAuthenticatorType(
+                    int32_t(keymintAuthToken.authenticatorType));
+            aidlAuthToken.timestamp.milliSeconds = keymintAuthToken.timestamp.milliSeconds;
+            aidlAuthToken.mac = keymintAuthToken.mac;
 
-    sp<CredstoreTokenCallback> callback = new CredstoreTokenCallback();
-    auto future = callback->get_future();
-
-    Status status =
-        keystore->getTokensForCredstore(challenge, secureUserId, authTokenMaxAgeMillis, callback);
-    if (!status.isOk()) {
+            // Convert timestamp token to KeyMaster verification token
+            TimeStampToken timestampToken = authzTokens.timestampToken;
+            aidlVerificationToken.challenge = timestampToken.challenge;
+            aidlVerificationToken.timestamp.milliSeconds = timestampToken.timestamp.milliSeconds;
+            // Legacy verification tokens were always minted by TEE.
+            aidlVerificationToken.securityLevel = SecurityLevel::TRUSTED_ENVIRONMENT;
+            aidlVerificationToken.mac = timestampToken.mac;
+        } else {
+            if (result.getServiceSpecificError() == 0) {
+                // Here we differentiate the errors occurred during communication
+                // from the service specific errors.
+                LOG(ERROR) << "Error getting tokens from keystore2: " << result.getDescription();
+                return false;
+            } else {
+                // Log the reason for not receiving auth tokens from keystore2.
+                LOG(INFO) << "Auth tokens were not received due to: " << result.getDescription();
+            }
+        }
+        return true;
+    } else {
+        LOG(ERROR) << "Error connecting to IKeystoreAuthorization service";
         return false;
     }
-
-    auto fstatus = future.wait_for(std::chrono::milliseconds(5000));
-    if (fstatus != std::future_status::ready) {
-        LOG(ERROR) << "Waited 5 seconds from tokens for credstore, aborting";
-        return false;
-    }
-    auto [success, returnedAuthToken, returnedVerificationToken] = future.get();
-    if (!success) {
-        LOG(ERROR) << "Error getting tokens from credstore";
-        return false;
-    }
-    authToken = returnedAuthToken;
-    verificationToken = returnedVerificationToken;
-    return true;
 }
 
 Status Credential::getEntries(const vector<uint8_t>& requestMessage,
@@ -279,13 +307,6 @@
         }
     }
 
-    // If requesting a challenge-based authToken the idea is that authentication
-    // happens as part of the transaction. As such, authTokenMaxAgeMillis should
-    // be nearly zero. We'll use 10 seconds for this.
-    if (userAuthNeeded && selectedChallenge_ != 0) {
-        authTokenMaxAgeMillis = 10 * 1000;
-    }
-
     // Reset tokens and only get them if they're actually needed, e.g. if user authentication
     // is needed in any of the access control profiles for data items being requested.
     //
@@ -303,63 +324,58 @@
     aidlVerificationToken.securityLevel = ::android::hardware::keymaster::SecurityLevel::SOFTWARE;
     aidlVerificationToken.mac.clear();
     if (userAuthNeeded) {
-        vector<uint8_t> authTokenBytes;
-        vector<uint8_t> verificationTokenBytes;
-        if (!getTokensFromKeystore(selectedChallenge_, data->getSecureUserId(),
-                                   authTokenMaxAgeMillis, authTokenBytes, verificationTokenBytes)) {
-            LOG(ERROR) << "Error getting tokens from keystore";
+        // If user authentication is needed, always get a challenge from the
+        // HAL/TA since it'll need it to check the returned VerificationToken
+        // for freshness.
+        if (!ensureChallenge()) {
             return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
-                                                    "Error getting tokens from keystore");
+                                                    "Error getting challenge (bug in HAL or TA)");
         }
 
-        // It's entirely possible getTokensFromKeystore() succeeded but didn't
-        // return any tokens (in which case the returned byte-vectors are
-        // empty). For example, this can happen if no auth token is available
-        // which satifies e.g. |authTokenMaxAgeMillis|.
+        // Note: if all selected profiles require auth-on-every-presentation
+        // then authTokenMaxAgeMillis will be 0 (because timeoutMillis for each
+        // profile is 0). Which means that keystore will only return an
+        // AuthToken if its challenge matches what we pass, regardless of its
+        // age. This is intended b/c the HAL/TA will check not care about
+        // the age in this case, it only cares that the challenge matches.
         //
-        if (authTokenBytes.size() > 0) {
-            HardwareAuthToken authToken =
-                android::hardware::keymaster::V4_0::support::hidlVec2AuthToken(authTokenBytes);
-            // Convert from HIDL to AIDL...
-            aidlAuthToken.challenge = int64_t(authToken.challenge);
-            aidlAuthToken.userId = int64_t(authToken.userId);
-            aidlAuthToken.authenticatorId = int64_t(authToken.authenticatorId);
-            aidlAuthToken.authenticatorType =
-                ::android::hardware::keymaster::HardwareAuthenticatorType(
-                    int32_t(authToken.authenticatorType));
-            aidlAuthToken.timestamp.milliSeconds = int64_t(authToken.timestamp);
-            aidlAuthToken.mac = authToken.mac;
-        }
+        // Otherwise, if one or more of the profiles is auth-with-a-timeout then
+        // authTokenMaxAgeMillis will be set to the largest of those
+        // timeouts. We'll get an AuthToken which satisfies this deadline if it
+        // exists. This authToken _may_ have the requested challenge but it's
+        // not a guarantee and it's also not required.
+        //
 
-        if (verificationTokenBytes.size() > 0) {
-            optional<VerificationToken> token =
-                android::hardware::keymaster::V4_0::support::deserializeVerificationToken(
-                    verificationTokenBytes);
-            if (!token) {
-                LOG(ERROR) << "Error deserializing verification token";
-                return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
-                                                        "Error deserializing verification token");
-            }
-            aidlVerificationToken.challenge = token->challenge;
-            aidlVerificationToken.timestamp.milliSeconds = token->timestamp;
-            aidlVerificationToken.securityLevel =
-                ::android::hardware::keymaster::SecurityLevel(token->securityLevel);
-            aidlVerificationToken.mac = token->mac;
+        if (!getTokensFromKeystore2(selectedChallenge_, data->getSecureUserId(),
+                                    authTokenMaxAgeMillis, aidlAuthToken, aidlVerificationToken)) {
+            LOG(ERROR) << "Error getting tokens from keystore2";
+            return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                    "Error getting tokens from keystore2");
         }
     }
 
     // Note that the selectAuthKey() method is only called if a CryptoObject is involved at
     // the Java layer. So we could end up with no previously selected auth key and we may
     // need one.
-    const AuthKeyData* authKey = selectedAuthKey_;
-    if (sessionTranscript.size() > 0) {
-        if (authKey == nullptr) {
-            authKey = data->selectAuthKey(allowUsingExhaustedKeys, allowUsingExpiredKeys);
-            if (authKey == nullptr) {
-                return Status::fromServiceSpecificError(
-                    ICredentialStore::ERROR_NO_AUTHENTICATION_KEY_AVAILABLE,
-                    "No suitable authentication key available");
-            }
+    //
+    const AuthKeyData* authKey =
+        data->selectAuthKey(allowUsingExhaustedKeys, allowUsingExpiredKeys);
+    if (authKey == nullptr) {
+        // If no authKey is available, consider it an error only when a
+        // SessionTranscript was provided.
+        //
+        // We allow no SessionTranscript to be provided because it makes
+        // the API simpler to deal with insofar it can be used without having
+        // to generate any authentication keys.
+        //
+        // In this "no SessionTranscript is provided" mode we don't return
+        // DeviceNameSpaces nor a MAC over DeviceAuthentication so we don't
+        // need a device key.
+        //
+        if (sessionTranscript.size() > 0) {
+            return Status::fromServiceSpecificError(
+                ICredentialStore::ERROR_NO_AUTHENTICATION_KEY_AVAILABLE,
+                "No suitable authentication key available and one is needed");
         }
     }
     vector<uint8_t> signingKeyBlob;
@@ -750,31 +766,36 @@
     //
     // It is because of this we need to set the CredentialKey certificate chain,
     // keyCount, and maxUsesPerKey below.
-    sp<WritableCredential> writableCredential =
-        new WritableCredential(dataPath_, credentialName_, docType.value(), true, hwInfo_,
-                               halWritableCredential, halApiVersion_);
+    sp<WritableCredential> writableCredential = new WritableCredential(
+        dataPath_, credentialName_, docType.value(), true, hwInfo_, halWritableCredential);
 
     writableCredential->setAttestationCertificate(data->getAttestationCertificate());
     auto [keyCount, maxUsesPerKey] = data->getAvailableAuthenticationKeys();
     writableCredential->setAvailableAuthenticationKeys(keyCount, maxUsesPerKey);
 
-    // Because its data has changed, we need to reconnect to the HAL when the
-    // credential has been updated... otherwise the remote object will have
-    // stale data for future calls (e.g. getAuthKeysNeedingCertification().
+    // Because its data has changed, we need to replace the binder for the
+    // IIdentityCredential when the credential has been updated... otherwise the
+    // remote object will have stale data for future calls, for example
+    // getAuthKeysNeedingCertification().
     //
-    // The joys and pitfalls of mutable objects...
+    // The way this is implemented is that setCredentialToReloadWhenUpdated()
+    // instructs the WritableCredential to call writableCredentialPersonalized()
+    // on |this|.
     //
-    writableCredential->setCredentialUpdatedCallback([this] {
-        Status status = this->ensureOrReplaceHalBinder();
-        if (!status.isOk()) {
-            LOG(ERROR) << "Error loading credential";
-        }
-    });
+    //
+    writableCredential->setCredentialToReloadWhenUpdated(this);
 
     *_aidl_return = writableCredential;
     return Status::ok();
 }
 
+void Credential::writableCredentialPersonalized() {
+    Status status = ensureOrReplaceHalBinder();
+    if (!status.isOk()) {
+        LOG(ERROR) << "Error reloading credential";
+    }
+}
+
 }  // namespace identity
 }  // namespace security
 }  // namespace android
diff --git a/identity/Credential.h b/identity/Credential.h
index 7f08515..a76f3cc 100644
--- a/identity/Credential.h
+++ b/identity/Credential.h
@@ -50,6 +50,7 @@
     ~Credential();
 
     Status ensureOrReplaceHalBinder();
+    void writableCredentialPersonalized();
 
     // ICredential overrides
     Status createEphemeralKeyPair(vector<uint8_t>* _aidl_return) override;
@@ -94,12 +95,13 @@
     HardwareInformation hwInfo_;
     sp<IIdentityCredentialStore> halStoreBinder_;
 
-    const AuthKeyData* selectedAuthKey_ = nullptr;
     uint64_t selectedChallenge_ = 0;
 
     sp<IIdentityCredential> halBinder_;
     int halApiVersion_;
 
+    bool ensureChallenge();
+
     ssize_t
     calcExpectedDeviceNameSpacesSize(const vector<uint8_t>& requestMessage,
                                      const vector<RequestNamespaceParcel>& requestNamespaces,
diff --git a/identity/CredentialData.h b/identity/CredentialData.h
index b037997..24b55d3 100644
--- a/identity/CredentialData.h
+++ b/identity/CredentialData.h
@@ -55,7 +55,7 @@
 
     vector<uint8_t> certificate;
     vector<uint8_t> keyBlob;
-    int64_t expirationDateMillisSinceEpoch;
+    int64_t expirationDateMillisSinceEpoch = 0;
     vector<uint8_t> staticAuthenticationData;
     vector<uint8_t> pendingCertificate;
     vector<uint8_t> pendingKeyBlob;
diff --git a/identity/CredentialStore.cpp b/identity/CredentialStore.cpp
index f77294e..509e022 100644
--- a/identity/CredentialStore.cpp
+++ b/identity/CredentialStore.cpp
@@ -90,7 +90,7 @@
     }
 
     sp<IWritableCredential> writableCredential = new WritableCredential(
-        dataPath_, credentialName, docType, false, hwInfo_, halWritableCredential, halApiVersion_);
+        dataPath_, credentialName, docType, false, hwInfo_, halWritableCredential);
     *_aidl_return = writableCredential;
     return Status::ok();
 }
diff --git a/identity/WritableCredential.cpp b/identity/WritableCredential.cpp
index d0688b8..a300e51 100644
--- a/identity/WritableCredential.cpp
+++ b/identity/WritableCredential.cpp
@@ -41,15 +41,14 @@
 WritableCredential::WritableCredential(const string& dataPath, const string& credentialName,
                                        const string& docType, bool isUpdate,
                                        HardwareInformation hwInfo,
-                                       sp<IWritableIdentityCredential> halBinder, int halApiVersion)
+                                       sp<IWritableIdentityCredential> halBinder)
     : dataPath_(dataPath), credentialName_(credentialName), docType_(docType), isUpdate_(isUpdate),
-      hwInfo_(std::move(hwInfo)), halBinder_(halBinder), halApiVersion_(halApiVersion) {}
+      hwInfo_(std::move(hwInfo)), halBinder_(halBinder) {}
 
 WritableCredential::~WritableCredential() {}
 
-void WritableCredential::setCredentialUpdatedCallback(
-    std::function<void()>&& onCredentialUpdatedCallback) {
-    onCredentialUpdatedCallback_ = onCredentialUpdatedCallback;
+void WritableCredential::setCredentialToReloadWhenUpdated(sp<Credential> credential) {
+    credentialToReloadWhenUpdated_ = credential;
 }
 
 Status WritableCredential::ensureAttestationCertificateExists(const vector<uint8_t>& challenge) {
@@ -268,7 +267,10 @@
                                                 "Error saving credential data to disk");
     }
 
-    onCredentialUpdatedCallback_();
+    if (credentialToReloadWhenUpdated_) {
+        credentialToReloadWhenUpdated_->writableCredentialPersonalized();
+        credentialToReloadWhenUpdated_.clear();
+    }
 
     *_aidl_return = proofOfProvisioningSignature;
     return Status::ok();
diff --git a/identity/WritableCredential.h b/identity/WritableCredential.h
index 6ff31ae..838b956 100644
--- a/identity/WritableCredential.h
+++ b/identity/WritableCredential.h
@@ -24,6 +24,8 @@
 
 #include <android/hardware/identity/IIdentityCredentialStore.h>
 
+#include "Credential.h"
+
 namespace android {
 namespace security {
 namespace identity {
@@ -38,13 +40,15 @@
   public:
     WritableCredential(const string& dataPath, const string& credentialName, const string& docType,
                        bool isUpdate, HardwareInformation hwInfo,
-                       sp<IWritableIdentityCredential> halBinder, int halApiVersion);
+                       sp<IWritableIdentityCredential> halBinder);
     ~WritableCredential();
 
     // Used when updating a credential
     void setAttestationCertificate(const vector<uint8_t>& attestationCertificate);
     void setAvailableAuthenticationKeys(int keyCount, int maxUsesPerKey);
-    void setCredentialUpdatedCallback(std::function<void()>&& onCredentialUpdatedCallback);
+
+    // Used by Credential::update()
+    void setCredentialToReloadWhenUpdated(sp<Credential> credential);
 
     // IWritableCredential overrides
     Status getCredentialKeyCertificateChain(const vector<uint8_t>& challenge,
@@ -61,13 +65,12 @@
     bool isUpdate_;
     HardwareInformation hwInfo_;
     sp<IWritableIdentityCredential> halBinder_;
-    int halApiVersion_;
 
     vector<uint8_t> attestationCertificate_;
     int keyCount_ = 0;
     int maxUsesPerKey_ = 1;
 
-    std::function<void()> onCredentialUpdatedCallback_ = []() {};
+    sp<Credential> credentialToReloadWhenUpdated_;
 
     ssize_t calcExpectedProofOfProvisioningSize(
         const vector<AccessControlProfileParcel>& accessControlProfiles,
diff --git a/identity/main.cpp b/identity/main.cpp
index 8f4968d..08f2219 100644
--- a/identity/main.cpp
+++ b/identity/main.cpp
@@ -40,7 +40,7 @@
 using ::android::security::identity::CredentialStoreFactory;
 
 int main(int argc, char* argv[]) {
-    InitLogging(argv, StderrLogger);
+    InitLogging(argv);
 
     CHECK(argc == 2) << "A directory must be specified";
     string data_dir = string(argv[1]);
@@ -53,10 +53,5 @@
     CHECK(ret == ::android::OK) << "Couldn't register binder service";
     LOG(ERROR) << "Registered binder service";
 
-    // This is needed for binder callbacks from keystore on a ICredstoreTokenCallback binder.
-    android::ProcessState::self()->startThreadPool();
-
-    IPCThreadState::self()->joinThreadPool();
-
     return 0;
 }
diff --git a/keystore/auth_token_table.cpp b/keystore/auth_token_table.cpp
index 5e6d572..971f9ef 100644
--- a/keystore/auth_token_table.cpp
+++ b/keystore/auth_token_table.cpp
@@ -178,33 +178,39 @@
                                               int64_t authTokenMaxAgeMillis) {
     std::vector<uint64_t> sids = {secureUserId};
     HardwareAuthenticatorType auth_type = HardwareAuthenticatorType::ANY;
-
     time_t now = clock_function_();
+    int64_t nowMillis = now * 1000;
 
-    // challenge-based - the authToken has to contain the given challenge.
-    if (challenge != 0) {
-        auto matching_op = find_if(
-            entries_, [&](Entry& e) { return e.token().challenge == challenge && !e.completed(); });
-        if (matching_op == entries_.end()) {
-            return {AUTH_TOKEN_NOT_FOUND, {}};
-        }
-
-        if (!matching_op->SatisfiesAuth(sids, auth_type)) {
-            return {AUTH_TOKEN_WRONG_SID, {}};
-        }
-
-        if (authTokenMaxAgeMillis > 0) {
-            if (static_cast<int64_t>(matching_op->time_received()) + authTokenMaxAgeMillis <
-                static_cast<int64_t>(now)) {
-                return {AUTH_TOKEN_EXPIRED, {}};
-            }
-        }
-
-        return {OK, matching_op->token()};
+    // It's an error to call this without a non-zero challenge.
+    if (challenge == 0) {
+        return {OP_HANDLE_REQUIRED, {}};
     }
 
-    // Otherwise, no challenge - any authToken younger than the specified maximum
-    // age will do.
+    // First see if we can find a token which matches the given challenge. If we
+    // can, return the newest one. We specifically don't care about its age.
+    //
+    Entry* newest_match_for_challenge = nullptr;
+    for (auto& entry : entries_) {
+        if (entry.token().challenge == challenge && !entry.completed() &&
+            entry.SatisfiesAuth(sids, auth_type)) {
+            if (newest_match_for_challenge == nullptr ||
+                entry.is_newer_than(newest_match_for_challenge)) {
+                newest_match_for_challenge = &entry;
+            }
+        }
+    }
+    if (newest_match_for_challenge != nullptr) {
+        newest_match_for_challenge->UpdateLastUse(now);
+        return {OK, newest_match_for_challenge->token()};
+    }
+
+    // If that didn't work, we'll take the most recent token within the specified
+    // deadline, if any. Of course if the deadline is zero it doesn't make sense
+    // to look at all.
+    if (authTokenMaxAgeMillis == 0) {
+        return {AUTH_TOKEN_NOT_FOUND, {}};
+    }
+
     Entry* newest_match = nullptr;
     for (auto& entry : entries_) {
         if (entry.SatisfiesAuth(sids, auth_type) && entry.is_newer_than(newest_match)) {
@@ -216,11 +222,9 @@
         return {AUTH_TOKEN_NOT_FOUND, {}};
     }
 
-    if (authTokenMaxAgeMillis > 0) {
-        if (static_cast<int64_t>(newest_match->time_received()) + authTokenMaxAgeMillis <
-            static_cast<int64_t>(now)) {
-            return {AUTH_TOKEN_EXPIRED, {}};
-        }
+    int64_t tokenAgeMillis = nowMillis - newest_match->time_received() * 1000;
+    if (tokenAgeMillis >= authTokenMaxAgeMillis) {
+        return {AUTH_TOKEN_EXPIRED, {}};
     }
 
     newest_match->UpdateLastUse(now);
diff --git a/keystore/binder/android/security/keystore/IKeystoreService.aidl b/keystore/binder/android/security/keystore/IKeystoreService.aidl
index e0879dd..3b9a1b4 100644
--- a/keystore/binder/android/security/keystore/IKeystoreService.aidl
+++ b/keystore/binder/android/security/keystore/IKeystoreService.aidl
@@ -87,7 +87,20 @@
     int onKeyguardVisibilityChanged(in boolean isShowing, in int userId);
     int listUidsOfAuthBoundKeys(out @utf8InCpp List<String> uids);
 
-    // Called by credstore (and only credstore).
+    // This method looks through auth-tokens cached by keystore which match
+    // the passed-in |secureUserId|.
+    //
+    // If one or more of these tokens has a |challenge| field which matches
+    // the passed-in |challenge| parameter, the most recent is returned. In
+    // this case the |authTokenMaxAgeMillis| parameter is not used.
+    //
+    // Otherwise, the most recent auth-token of these tokens which is younger
+    // than |authTokenMaxAgeMillis| is returned.
+    //
+    // The passed in |challenge| parameter must always be non-zero.
+    //
+    // This method is called by credstore (and only credstore).
+    //
     void getTokensForCredstore(in long challenge, in long secureUserId, in int authTokenMaxAgeMillis,
                                in ICredstoreTokenCallback cb);
 }
diff --git a/keystore2/Android.bp b/keystore2/Android.bp
index d480244..9875d64 100644
--- a/keystore2/Android.bp
+++ b/keystore2/Android.bp
@@ -29,11 +29,12 @@
     rustlibs: [
         "android.hardware.security.keymint-V1-rust",
         "android.hardware.security.secureclock-V1-rust",
+        "android.os.permissions_aidl-rust",
         "android.security.apc-rust",
         "android.security.authorization-rust",
         "android.security.compat-rust",
+        "android.security.maintenance-rust",
         "android.security.remoteprovisioning-rust",
-        "android.security.usermanager-rust",
         "android.system.keystore2-V1-rust",
         "libanyhow",
         "libbinder_rs",
diff --git a/keystore2/aidl/Android.bp b/keystore2/aidl/Android.bp
index f30ad83..69ba0b4 100644
--- a/keystore2/aidl/Android.bp
+++ b/keystore2/aidl/Android.bp
@@ -116,8 +116,8 @@
 }
 
 aidl_interface {
-    name: "android.security.usermanager",
-    srcs: [ "android/security/usermanager/*.aidl" ],
+    name: "android.security.maintenance",
+    srcs: [ "android/security/maintenance/*.aidl" ],
     imports: [
         "android.system.keystore2",
     ],
diff --git a/keystore2/aidl/android/security/authorization/AuthorizationTokens.aidl b/keystore2/aidl/android/security/authorization/AuthorizationTokens.aidl
new file mode 100644
index 0000000..9061998
--- /dev/null
+++ b/keystore2/aidl/android/security/authorization/AuthorizationTokens.aidl
@@ -0,0 +1,33 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package android.security.authorization;
+
+import android.hardware.security.keymint.HardwareAuthToken;
+import android.hardware.security.secureclock.TimeStampToken;
+
+/**
+ * This parcelable is returned by `IKeystoreAuthorization::getAuthTokensForCredStore`.
+ * @hide
+ */
+parcelable AuthorizationTokens {
+    /**
+     * HardwareAuthToken provided by an authenticator.
+     */
+    HardwareAuthToken authToken;
+    /**
+     * TimeStampToken provided by a SecureClock.
+     */
+    TimeStampToken timestampToken;
+}
\ No newline at end of file
diff --git a/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl b/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
index df64401..6dc172e 100644
--- a/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
+++ b/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
@@ -16,6 +16,7 @@
 
 import android.hardware.security.keymint.HardwareAuthToken;
 import android.security.authorization.LockScreenEvent;
+import android.security.authorization.AuthorizationTokens;
 
 // TODO: mark the interface with @SensitiveData when the annotation is ready (b/176110256).
 
@@ -45,6 +46,8 @@
      * ## Error conditions:
      * `ResponseCode::PERMISSION_DENIED` - if the callers do not have the 'Unlock' permission.
      * `ResponseCode::SYSTEM_ERROR` - if failed to perform lock/unlock operations due to various
+     * `ResponseCode::VALUE_CORRUPTED` - if the super key can not be decrypted.
+     * `ResponseCode::KEY_NOT_FOUND` - if the super key is not found.
      *
      * @lockScreenEvent - Indicates what happened.
      *                    * LockScreenEvent.UNLOCK if the screen was unlocked.
@@ -56,4 +59,37 @@
      */
     void onLockScreenEvent(in LockScreenEvent lockScreenEvent, in int userId,
                            in @nullable byte[] password);
+
+    /**
+     * Allows Credstore to retrieve a HardwareAuthToken and a TimestampToken.
+     * Identity Credential Trusted App can run either in the TEE or in other secure Hardware.
+     * So, credstore always need to retrieve a TimestampToken along with a HardwareAuthToken.
+     *
+     * The passed in |challenge| parameter must always be non-zero.
+     *
+     * The returned TimestampToken will always have its |challenge| field set to
+     * the |challenge| parameter.
+     *
+     * This method looks through auth-tokens cached by keystore which match
+     * the passed-in |secureUserId|.
+     * The most recent matching auth token which has a |challenge| field which matches
+     * the passed-in |challenge| parameter is returned.
+     * In this case the |authTokenMaxAgeMillis| parameter is not used.
+     *
+     * Otherwise, the most recent matching auth token which is younger
+     * than |authTokenMaxAgeMillis| is returned.
+     *
+     * This method is called by credstore (and only credstore).
+     *
+     * The caller requires 'get_auth_token' permission.
+     *
+     * ## Error conditions:
+     * `ResponseCode::PERMISSION_DENIED` - if the caller does not have the 'get_auth_token'
+     *                                     permission.
+     * `ResponseCode::SYSTEM_ERROR` - if failed to obtain an authtoken from the database.
+     * `ResponseCode::NO_AUTH_TOKEN_FOUND` - a matching auth token is not found.
+     * `ResponseCode::INVALID_ARGUMENT` - if the passed-in |challenge| parameter is zero.
+     */
+    AuthorizationTokens getAuthTokensForCredStore(in long challenge, in long secureUserId,
+     in long authTokenMaxAgeMillis);
 }
diff --git a/keystore2/aidl/android/security/authorization/ResponseCode.aidl b/keystore2/aidl/android/security/authorization/ResponseCode.aidl
new file mode 100644
index 0000000..94f1120
--- /dev/null
+++ b/keystore2/aidl/android/security/authorization/ResponseCode.aidl
@@ -0,0 +1,56 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package android.security.authorization;
+
+/**
+ * Used as exception codes by IKeystoreAuthorization.
+ */
+@Backing(type="int")
+enum ResponseCode {
+    /**
+     * A matching auth token is not found.
+     */
+    NO_AUTH_TOKEN_FOUND = 1,
+    /**
+     * The matching auth token is expired.
+     */
+    AUTH_TOKEN_EXPIRED = 2,
+    /**
+     * Same as in keystore2/ResponseCode.aidl.
+     * Any unexpected Error such as IO or communication errors.
+     */
+    SYSTEM_ERROR = 4,
+    /**
+     * Same as in keystore2/ResponseCode.aidl.
+     * Indicates that the caller does not have the permissions for the attempted request.
+     */
+    PERMISSION_DENIED = 6,
+    /**
+     * Same as in keystore2/ResponseCode.aidl.
+     * Indicates that the requested key does not exist.
+     */
+    KEY_NOT_FOUND = 7,
+    /**
+     * Same as in keystore2/ResponseCode.aidl.
+     * Indicates that a value being processed is corrupted.
+     */
+    VALUE_CORRUPTED = 8,
+    /**
+     * Same as in keystore2/ResponseCode.aidl.
+     * Indicates that an invalid argument was passed to an API call.
+     */
+    INVALID_ARGUMENT = 20,
+
+ }
\ No newline at end of file
diff --git a/keystore2/aidl/android/security/usermanager/IKeystoreUserManager.aidl b/keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl
similarity index 81%
rename from keystore2/aidl/android/security/usermanager/IKeystoreUserManager.aidl
rename to keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl
index 83edb1a..3115e92 100644
--- a/keystore2/aidl/android/security/usermanager/IKeystoreUserManager.aidl
+++ b/keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl
@@ -12,18 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package android.security.usermanager;
+package android.security.maintenance;
 
 import android.system.keystore2.Domain;
+import android.security.maintenance.UserState;
 
 // TODO: mark the interface with @SensitiveData when the annotation is ready (b/176110256).
 
 /**
- * IKeystoreUserManager interface exposes the methods for adding/removing users and changing the
+ * IKeystoreMaintenance interface exposes the methods for adding/removing users and changing the
  * user's password.
  * @hide
  */
-interface IKeystoreUserManager {
+interface IKeystoreMaintenance {
 
     /**
      * Allows LockSettingsService to inform keystore about adding a new user.
@@ -75,4 +76,17 @@
      * @hide
      */
      void clearNamespace(Domain domain, long nspace);
+
+    /**
+     * Allows querying user state, given user id.
+     * Callers require 'GetState' permission.
+     * ## Error conditions:
+     * `ResponseCode::PERMISSION_DENIED` - if the callers do not have the 'GetState'
+     *                                     permission.
+     * `ResponseCode::SYSTEM_ERROR` - if an error occurred when querying the user state.
+     *
+     * @param userId - Android user id
+     * @hide
+     */
+    UserState getState(in int userId);
 }
diff --git a/keystore2/aidl/android/security/maintenance/UserState.aidl b/keystore2/aidl/android/security/maintenance/UserState.aidl
new file mode 100644
index 0000000..b6fe278
--- /dev/null
+++ b/keystore2/aidl/android/security/maintenance/UserState.aidl
@@ -0,0 +1,22 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package android.security.maintenance;
+
+@Backing(type="int")
+enum UserState {
+    UNINITIALIZED = 0,
+    LSKF_UNLOCKED = 1,
+    LSKF_LOCKED = 2,
+}
\ No newline at end of file
diff --git a/keystore2/aidl/android/security/remoteprovisioning/IRemoteProvisioning.aidl b/keystore2/aidl/android/security/remoteprovisioning/IRemoteProvisioning.aidl
index 0d4c30f..5c2d0b1 100644
--- a/keystore2/aidl/android/security/remoteprovisioning/IRemoteProvisioning.aidl
+++ b/keystore2/aidl/android/security/remoteprovisioning/IRemoteProvisioning.aidl
@@ -133,4 +133,13 @@
      * @return The array of security levels.
      */
      SecurityLevel[] getSecurityLevels();
+
+    /**
+     * This method deletes all remotely provisioned attestation keys in the database, regardless
+     * of what state in their life cycle they are in. This is primarily useful to facilitate
+     * testing.
+     *
+     * @return Number of keys deleted
+     */
+    long deleteAllKeys();
 }
diff --git a/keystore2/keystore2.rc b/keystore2/keystore2.rc
index 2d1f05a..82bf3b8 100644
--- a/keystore2/keystore2.rc
+++ b/keystore2/keystore2.rc
@@ -6,14 +6,8 @@
 #
 # See system/core/init/README.md for information on the init.rc language.
 
-# Start Keystore 2 conditionally
-# TODO b/171563717 Remove when Keystore 2 migration is complete.
-on property:persist.android.security.keystore2.enable=true
-    enable keystore2
-
 service keystore2 /system/bin/keystore2 /data/misc/keystore
     class early_hal
     user keystore
     group keystore readproc log
     writepid /dev/cpuset/foreground/tasks
-    disabled
diff --git a/keystore2/selinux/src/lib.rs b/keystore2/selinux/src/lib.rs
index 2b5091d..cc707e7 100644
--- a/keystore2/selinux/src/lib.rs
+++ b/keystore2/selinux/src/lib.rs
@@ -455,7 +455,6 @@
 
         check_keystore_perm!(add_auth);
         check_keystore_perm!(clear_ns);
-        check_keystore_perm!(get_state);
         check_keystore_perm!(lock);
         check_keystore_perm!(reset);
         check_keystore_perm!(unlock);
diff --git a/keystore2/src/apc.rs b/keystore2/src/apc.rs
index 767014e..f8259ea 100644
--- a/keystore2/src/apc.rs
+++ b/keystore2/src/apc.rs
@@ -31,7 +31,7 @@
     ExceptionCode, Interface, Result as BinderResult, SpIBinder, Status as BinderStatus, Strong,
 };
 use anyhow::{Context, Result};
-use binder::{IBinder, ThreadState};
+use binder::{IBinderInternal, ThreadState};
 use keystore2_apc_compat::ApcHal;
 use keystore2_selinux as selinux;
 use std::time::{Duration, Instant};
diff --git a/keystore2/src/async_task.rs b/keystore2/src/async_task.rs
index 2b36f1f..20a7458 100644
--- a/keystore2/src/async_task.rs
+++ b/keystore2/src/async_task.rs
@@ -89,8 +89,10 @@
 struct AsyncTaskState {
     state: State,
     thread: Option<thread::JoinHandle<()>>,
+    timeout: Duration,
     hi_prio_req: VecDeque<Box<dyn FnOnce(&mut Shelf) + Send>>,
     lo_prio_req: VecDeque<Box<dyn FnOnce(&mut Shelf) + Send>>,
+    idle_fns: Vec<Arc<dyn Fn(&mut Shelf) + Send + Sync>>,
     /// The store allows tasks to store state across invocations. It is passed to each invocation
     /// of each task. Tasks need to cooperate on the ids they use for storing state.
     shelf: Option<Shelf>,
@@ -107,25 +109,32 @@
 
 impl Default for AsyncTask {
     fn default() -> Self {
+        Self::new(Duration::from_secs(30))
+    }
+}
+
+impl AsyncTask {
+    /// Construct an [`AsyncTask`] with a specific timeout value.
+    pub fn new(timeout: Duration) -> Self {
         Self {
             state: Arc::new((
                 Condvar::new(),
                 Mutex::new(AsyncTaskState {
                     state: State::Exiting,
                     thread: None,
+                    timeout,
                     hi_prio_req: VecDeque::new(),
                     lo_prio_req: VecDeque::new(),
+                    idle_fns: Vec::new(),
                     shelf: None,
                 }),
             )),
         }
     }
-}
 
-impl AsyncTask {
-    /// Adds a job to the high priority queue. High priority jobs are completed before
-    /// low priority jobs and can also overtake low priority jobs. But they cannot
-    /// preempt them.
+    /// Adds a one-off job to the high priority queue. High priority jobs are
+    /// completed before low priority jobs and can also overtake low priority
+    /// jobs. But they cannot preempt them.
     pub fn queue_hi<F>(&self, f: F)
     where
         F: for<'r> FnOnce(&'r mut Shelf) + Send + 'static,
@@ -133,10 +142,10 @@
         self.queue(f, true)
     }
 
-    /// Adds a job to the low priority queue. Low priority jobs are completed after
-    /// high priority. And they are not executed as long as high priority jobs are
-    /// present. Jobs always run to completion and are never preempted by high
-    /// priority jobs.
+    /// Adds a one-off job to the low priority queue. Low priority jobs are
+    /// completed after high priority. And they are not executed as long as high
+    /// priority jobs are present. Jobs always run to completion and are never
+    /// preempted by high priority jobs.
     pub fn queue_lo<F>(&self, f: F)
     where
         F: FnOnce(&mut Shelf) + Send + 'static,
@@ -144,6 +153,17 @@
         self.queue(f, false)
     }
 
+    /// Adds an idle callback. This will be invoked whenever the worker becomes
+    /// idle (all high and low priority jobs have been performed).
+    pub fn add_idle<F>(&self, f: F)
+    where
+        F: Fn(&mut Shelf) + Send + Sync + 'static,
+    {
+        let (ref _condvar, ref state) = *self.state;
+        let mut state = state.lock().unwrap();
+        state.idle_fns.push(Arc::new(f));
+    }
+
     fn queue<F>(&self, f: F, hi_prio: bool)
     where
         F: for<'r> FnOnce(&'r mut Shelf) + Send + 'static,
@@ -169,39 +189,66 @@
         }
 
         let cloned_state = self.state.clone();
+        let timeout_period = state.timeout;
 
         state.thread = Some(thread::spawn(move || {
             let (ref condvar, ref state) = *cloned_state;
+
+            enum Action {
+                QueuedFn(Box<dyn FnOnce(&mut Shelf) + Send>),
+                IdleFns(Vec<Arc<dyn Fn(&mut Shelf) + Send + Sync>>),
+            };
+            let mut done_idle = false;
+
             // When the worker starts, it takes the shelf and puts it on the stack.
             let mut shelf = state.lock().unwrap().shelf.take().unwrap_or_default();
             loop {
-                if let Some(f) = {
-                    let (mut state, timeout) = condvar
-                        .wait_timeout_while(
-                            state.lock().unwrap(),
-                            Duration::from_secs(30),
-                            |state| state.hi_prio_req.is_empty() && state.lo_prio_req.is_empty(),
-                        )
-                        .unwrap();
-                    match (
-                        state.hi_prio_req.pop_front(),
-                        state.lo_prio_req.is_empty(),
-                        timeout.timed_out(),
-                    ) {
-                        (Some(f), _, _) => Some(f),
-                        (None, false, _) => state.lo_prio_req.pop_front(),
-                        (None, true, true) => {
-                            // When the worker exits it puts the shelf back into the shared
-                            // state for the next worker to use. So state is preserved not
-                            // only across invocations but also across worker thread shut down.
-                            state.shelf = Some(shelf);
-                            state.state = State::Exiting;
-                            break;
+                if let Some(action) = {
+                    let state = state.lock().unwrap();
+                    if !done_idle && state.hi_prio_req.is_empty() && state.lo_prio_req.is_empty() {
+                        // No jobs queued so invoke the idle callbacks.
+                        Some(Action::IdleFns(state.idle_fns.clone()))
+                    } else {
+                        // Wait for either a queued job to arrive or a timeout.
+                        let (mut state, timeout) = condvar
+                            .wait_timeout_while(state, timeout_period, |state| {
+                                state.hi_prio_req.is_empty() && state.lo_prio_req.is_empty()
+                            })
+                            .unwrap();
+                        match (
+                            state.hi_prio_req.pop_front(),
+                            state.lo_prio_req.is_empty(),
+                            timeout.timed_out(),
+                        ) {
+                            (Some(f), _, _) => Some(Action::QueuedFn(f)),
+                            (None, false, _) => {
+                                state.lo_prio_req.pop_front().map(|f| Action::QueuedFn(f))
+                            }
+                            (None, true, true) => {
+                                // When the worker exits it puts the shelf back into the shared
+                                // state for the next worker to use. So state is preserved not
+                                // only across invocations but also across worker thread shut down.
+                                state.shelf = Some(shelf);
+                                state.state = State::Exiting;
+                                break;
+                            }
+                            (None, true, false) => None,
                         }
-                        (None, true, false) => None,
                     }
                 } {
-                    f(&mut shelf)
+                    // Now that the lock has been dropped, perform the action.
+                    match action {
+                        Action::QueuedFn(f) => {
+                            f(&mut shelf);
+                            done_idle = false;
+                        }
+                        Action::IdleFns(idle_fns) => {
+                            for idle_fn in idle_fns {
+                                idle_fn(&mut shelf);
+                            }
+                            done_idle = true;
+                        }
+                    }
                 }
             }
         }));
@@ -212,7 +259,11 @@
 #[cfg(test)]
 mod tests {
     use super::{AsyncTask, Shelf};
-    use std::sync::mpsc::channel;
+    use std::sync::{
+        mpsc::{channel, sync_channel, RecvTimeoutError},
+        Arc,
+    };
+    use std::time::Duration;
 
     #[test]
     fn test_shelf() {
@@ -306,6 +357,21 @@
     }
 
     #[test]
+    fn test_async_task_chain() {
+        let at = Arc::new(AsyncTask::default());
+        let (sender, receiver) = channel();
+        // Queue up a job that will queue up another job. This confirms
+        // that the job is not invoked with any internal AsyncTask locks held.
+        let at_clone = at.clone();
+        at.queue_hi(move |_shelf| {
+            at_clone.queue_lo(move |_shelf| {
+                sender.send(()).unwrap();
+            });
+        });
+        receiver.recv().unwrap();
+    }
+
+    #[test]
     #[should_panic]
     fn test_async_task_panic() {
         let at = AsyncTask::default();
@@ -319,4 +385,147 @@
         });
         done_receiver.recv().unwrap();
     }
+
+    #[test]
+    fn test_async_task_idle() {
+        let at = AsyncTask::new(Duration::from_secs(3));
+        // Need a SyncSender as it is Send+Sync.
+        let (idle_done_sender, idle_done_receiver) = sync_channel::<()>(3);
+        at.add_idle(move |_shelf| {
+            idle_done_sender.send(()).unwrap();
+        });
+
+        // Queue up some high-priority and low-priority jobs that take time.
+        for _i in 0..3 {
+            at.queue_lo(|_shelf| {
+                std::thread::sleep(Duration::from_millis(500));
+            });
+            at.queue_hi(|_shelf| {
+                std::thread::sleep(Duration::from_millis(500));
+            });
+        }
+        // Final low-priority job.
+        let (done_sender, done_receiver) = channel();
+        at.queue_lo(move |_shelf| {
+            done_sender.send(()).unwrap();
+        });
+
+        // Nothing happens until the last job completes.
+        assert_eq!(
+            idle_done_receiver.recv_timeout(Duration::from_secs(1)),
+            Err(RecvTimeoutError::Timeout)
+        );
+        done_receiver.recv().unwrap();
+        idle_done_receiver.recv_timeout(Duration::from_millis(1)).unwrap();
+
+        // Idle callback not executed again even if we wait for a while.
+        assert_eq!(
+            idle_done_receiver.recv_timeout(Duration::from_secs(3)),
+            Err(RecvTimeoutError::Timeout)
+        );
+
+        // However, if more work is done then there's another chance to go idle.
+        let (done_sender, done_receiver) = channel();
+        at.queue_hi(move |_shelf| {
+            std::thread::sleep(Duration::from_millis(500));
+            done_sender.send(()).unwrap();
+        });
+        // Idle callback not immediately executed, because the high priority
+        // job is taking a while.
+        assert_eq!(
+            idle_done_receiver.recv_timeout(Duration::from_millis(1)),
+            Err(RecvTimeoutError::Timeout)
+        );
+        done_receiver.recv().unwrap();
+        idle_done_receiver.recv_timeout(Duration::from_millis(1)).unwrap();
+    }
+
+    #[test]
+    fn test_async_task_multiple_idle() {
+        let at = AsyncTask::new(Duration::from_secs(3));
+        let (idle_sender, idle_receiver) = sync_channel::<i32>(5);
+        // Queue a high priority job to start things off
+        at.queue_hi(|_shelf| {
+            std::thread::sleep(Duration::from_millis(500));
+        });
+
+        // Multiple idle callbacks.
+        for i in 0..3 {
+            let idle_sender = idle_sender.clone();
+            at.add_idle(move |_shelf| {
+                idle_sender.send(i).unwrap();
+            });
+        }
+
+        // Nothing happens immediately.
+        assert_eq!(
+            idle_receiver.recv_timeout(Duration::from_millis(1)),
+            Err(RecvTimeoutError::Timeout)
+        );
+        // Wait for a moment and the idle jobs should have run.
+        std::thread::sleep(Duration::from_secs(1));
+
+        let mut results = Vec::new();
+        while let Ok(i) = idle_receiver.recv_timeout(Duration::from_millis(1)) {
+            results.push(i);
+        }
+        assert_eq!(results, [0, 1, 2]);
+    }
+
+    #[test]
+    fn test_async_task_idle_queues_job() {
+        let at = Arc::new(AsyncTask::new(Duration::from_secs(1)));
+        let at_clone = at.clone();
+        let (idle_sender, idle_receiver) = sync_channel::<i32>(100);
+        // Add an idle callback that queues a low-priority job.
+        at.add_idle(move |shelf| {
+            at_clone.queue_lo(|_shelf| {
+                // Slow things down so the channel doesn't fill up.
+                std::thread::sleep(Duration::from_millis(50));
+            });
+            let i = shelf.get_mut::<i32>();
+            idle_sender.send(*i).unwrap();
+            *i += 1;
+        });
+
+        // Nothing happens immediately.
+        assert_eq!(
+            idle_receiver.recv_timeout(Duration::from_millis(1500)),
+            Err(RecvTimeoutError::Timeout)
+        );
+
+        // Once we queue a normal job, things start.
+        at.queue_hi(|_shelf| {});
+        assert_eq!(0, idle_receiver.recv_timeout(Duration::from_millis(200)).unwrap());
+
+        // The idle callback queues a job, and completion of that job
+        // means the task is going idle again...so the idle callback will
+        // be called repeatedly.
+        assert_eq!(1, idle_receiver.recv_timeout(Duration::from_millis(100)).unwrap());
+        assert_eq!(2, idle_receiver.recv_timeout(Duration::from_millis(100)).unwrap());
+        assert_eq!(3, idle_receiver.recv_timeout(Duration::from_millis(100)).unwrap());
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_async_task_idle_panic() {
+        let at = AsyncTask::new(Duration::from_secs(1));
+        let (idle_sender, idle_receiver) = sync_channel::<()>(3);
+        // Add an idle callback that panics.
+        at.add_idle(move |_shelf| {
+            idle_sender.send(()).unwrap();
+            panic!("Panic from idle callback");
+        });
+        // Queue a job to trigger idleness and ensuing panic.
+        at.queue_hi(|_shelf| {});
+        idle_receiver.recv().unwrap();
+
+        // Queue another job afterwards to ensure that the async thread gets joined
+        // and the panic detected.
+        let (done_sender, done_receiver) = channel();
+        at.queue_hi(move |_shelf| {
+            done_sender.send(()).unwrap();
+        });
+        done_receiver.recv().unwrap();
+    }
 }
diff --git a/keystore2/src/authorization.rs b/keystore2/src/authorization.rs
index 02b19c4..5abb426 100644
--- a/keystore2/src/authorization.rs
+++ b/keystore2/src/authorization.rs
@@ -15,7 +15,6 @@
 //! This module implements IKeystoreAuthorization AIDL interface.
 
 use crate::error::Error as KeystoreError;
-use crate::error::map_or_log_err;
 use crate::globals::{ENFORCEMENTS, SUPER_KEY, DB, LEGACY_MIGRATOR};
 use crate::permission::KeystorePerm;
 use crate::super_key::UserState;
@@ -23,14 +22,87 @@
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     HardwareAuthToken::HardwareAuthToken,
 };
-use android_security_authorization::binder::{Interface, Result as BinderResult, Strong};
-use android_security_authorization::aidl::android::security::authorization::IKeystoreAuthorization::{
-        BnKeystoreAuthorization, IKeystoreAuthorization,
+use android_security_authorization::binder::{ExceptionCode, Interface, Result as BinderResult,
+     Strong, Status as BinderStatus};
+use android_security_authorization::aidl::android::security::authorization::{
+    IKeystoreAuthorization::BnKeystoreAuthorization, IKeystoreAuthorization::IKeystoreAuthorization,
+    LockScreenEvent::LockScreenEvent, AuthorizationTokens::AuthorizationTokens,
+    ResponseCode::ResponseCode,
 };
-use android_security_authorization:: aidl::android::security::authorization::LockScreenEvent::LockScreenEvent;
-use android_system_keystore2::aidl::android::system::keystore2::ResponseCode::ResponseCode;
+use android_system_keystore2::aidl::android::system::keystore2::{
+    ResponseCode::ResponseCode as KsResponseCode };
 use anyhow::{Context, Result};
-use binder::IBinder;
+use binder::IBinderInternal;
+use keystore2_selinux as selinux;
+
+/// This is the Authorization error type, it wraps binder exceptions and the
+/// Authorization ResponseCode
+#[derive(Debug, thiserror::Error, PartialEq)]
+pub enum Error {
+    /// Wraps an IKeystoreAuthorization response code as defined by
+    /// android.security.authorization AIDL interface specification.
+    #[error("Error::Rc({0:?})")]
+    Rc(ResponseCode),
+    /// Wraps a Binder exception code other than a service specific exception.
+    #[error("Binder exception code {0:?}, {1:?}")]
+    Binder(ExceptionCode, i32),
+}
+
+/// This function should be used by authorization service calls to translate error conditions
+/// into service specific exceptions.
+///
+/// All error conditions get logged by this function.
+///
+/// `Error::Rc(x)` variants get mapped onto a service specific error code of `x`.
+/// Certain response codes may be returned from keystore/ResponseCode.aidl by the keystore2 modules,
+/// which are then converted to the corresponding response codes of android.security.authorization
+/// AIDL interface specification.
+///
+/// `selinux::Error::perm()` is mapped on `ResponseCode::PERMISSION_DENIED`.
+///
+/// All non `Error` error conditions get mapped onto ResponseCode::SYSTEM_ERROR`.
+///
+/// `handle_ok` will be called if `result` is `Ok(value)` where `value` will be passed
+/// as argument to `handle_ok`. `handle_ok` must generate a `BinderResult<T>`, but it
+/// typically returns Ok(value).
+pub fn map_or_log_err<T, U, F>(result: Result<U>, handle_ok: F) -> BinderResult<T>
+where
+    F: FnOnce(U) -> BinderResult<T>,
+{
+    result.map_or_else(
+        |e| {
+            log::error!("{:#?}", e);
+            let root_cause = e.root_cause();
+            if let Some(KeystoreError::Rc(ks_rcode)) = root_cause.downcast_ref::<KeystoreError>() {
+                let rc = match *ks_rcode {
+                    // Although currently keystore2/ResponseCode.aidl and
+                    // authorization/ResponseCode.aidl share the same integer values for the
+                    // common response codes, this may deviate in the future, hence the
+                    // conversion here.
+                    KsResponseCode::SYSTEM_ERROR => ResponseCode::SYSTEM_ERROR.0,
+                    KsResponseCode::KEY_NOT_FOUND => ResponseCode::KEY_NOT_FOUND.0,
+                    KsResponseCode::VALUE_CORRUPTED => ResponseCode::VALUE_CORRUPTED.0,
+                    KsResponseCode::INVALID_ARGUMENT => ResponseCode::INVALID_ARGUMENT.0,
+                    // If the code paths of IKeystoreAuthorization aidl's methods happen to return
+                    // other error codes from KsResponseCode in the future, they should be converted
+                    // as well.
+                    _ => ResponseCode::SYSTEM_ERROR.0,
+                };
+                return Err(BinderStatus::new_service_specific_error(rc, None));
+            }
+            let rc = match root_cause.downcast_ref::<Error>() {
+                Some(Error::Rc(rcode)) => rcode.0,
+                Some(Error::Binder(_, _)) => ResponseCode::SYSTEM_ERROR.0,
+                None => match root_cause.downcast_ref::<selinux::Error>() {
+                    Some(selinux::Error::PermissionDenied) => ResponseCode::PERMISSION_DENIED.0,
+                    _ => ResponseCode::SYSTEM_ERROR.0,
+                },
+            };
+            Err(BinderStatus::new_service_specific_error(rc, None))
+        },
+        handle_ok,
+    )
+}
 
 /// This struct is defined to implement the aforementioned AIDL interface.
 /// As of now, it is an empty struct.
@@ -99,11 +171,33 @@
             }
             _ => {
                 // Any other combination is not supported.
-                Err(KeystoreError::Rc(ResponseCode::INVALID_ARGUMENT))
+                Err(Error::Rc(ResponseCode::INVALID_ARGUMENT))
                     .context("In on_lock_screen_event: Unknown event.")
             }
         }
     }
+
+    fn get_auth_tokens_for_credstore(
+        &self,
+        challenge: i64,
+        secure_user_id: i64,
+        auth_token_max_age_millis: i64,
+    ) -> Result<AuthorizationTokens> {
+        // Check permission. Function should return if this failed. Therefore having '?' at the end
+        // is very important.
+        check_keystore_permission(KeystorePerm::get_auth_token())
+            .context("In get_auth_tokens_for_credstore.")?;
+
+        // if the challenge is zero, return error
+        if challenge == 0 {
+            return Err(Error::Rc(ResponseCode::INVALID_ARGUMENT))
+                .context("In get_auth_tokens_for_credstore. Challenge can not be zero.");
+        }
+        // Obtain the auth token and the timestamp token from the enforcement module.
+        let (auth_token, ts_token) =
+            ENFORCEMENTS.get_auth_tokens(challenge, secure_user_id, auth_token_max_age_millis)?;
+        Ok(AuthorizationTokens { authToken: auth_token, timestampToken: ts_token })
+    }
 }
 
 impl Interface for AuthorizationManager {}
@@ -121,4 +215,20 @@
     ) -> BinderResult<()> {
         map_or_log_err(self.on_lock_screen_event(lock_screen_event, user_id, password), Ok)
     }
+
+    fn getAuthTokensForCredStore(
+        &self,
+        challenge: i64,
+        secure_user_id: i64,
+        auth_token_max_age_millis: i64,
+    ) -> binder::public_api::Result<AuthorizationTokens> {
+        map_or_log_err(
+            self.get_auth_tokens_for_credstore(
+                challenge,
+                secure_user_id,
+                auth_token_max_age_millis,
+            ),
+            Ok,
+        )
+    }
 }
diff --git a/keystore2/src/crypto/crypto.hpp b/keystore2/src/crypto/crypto.hpp
index 1b8971f..6686c8c 100644
--- a/keystore2/src/crypto/crypto.hpp
+++ b/keystore2/src/crypto/crypto.hpp
@@ -67,7 +67,7 @@
 // cert_len, extract the subject, DER-encode it and write the result to
 // subject_buf, which has subject_buf_len capacity.
 //
-// Because the length of the issuer is unknown, and becaue we'd like to (a) be
+// Because the length of the subject is unknown, and because we'd like to (a) be
 // able to handle subjects of any size and (b) avoid parsing the certificate
 // twice most of the time, once to discover the length and once to parse it, the
 // return value is overloaded.
diff --git a/keystore2/src/crypto/lib.rs b/keystore2/src/crypto/lib.rs
index f23778c..bd5906c 100644
--- a/keystore2/src/crypto/lib.rs
+++ b/keystore2/src/crypto/lib.rs
@@ -58,10 +58,15 @@
 
 /// Generate a salt.
 pub fn generate_salt() -> Result<Vec<u8>, Error> {
-    // Safety: salt has the same length as the requested number of random bytes.
-    let mut salt = vec![0; SALT_LENGTH];
-    if unsafe { randomBytes(salt.as_mut_ptr(), SALT_LENGTH) } {
-        Ok(salt)
+    generate_random_data(SALT_LENGTH)
+}
+
+/// Generate random data of the given size.
+pub fn generate_random_data(size: usize) -> Result<Vec<u8>, Error> {
+    // Safety: data has the same length as the requested number of random bytes.
+    let mut data = vec![0; size];
+    if unsafe { randomBytes(data.as_mut_ptr(), size) } {
+        Ok(data)
     } else {
         Err(Error::RandomNumberGenerationFailed)
     }
@@ -354,11 +359,13 @@
     Ok(OwnedECPoint(result))
 }
 
-/// Uses BoringSSL to extract the DER-encoded issuer subject from a
-/// DER-encoded X.509 certificate.
-pub fn parse_issuer_subject_from_certificate(cert_buf: &[u8]) -> Result<Vec<u8>, Error> {
+/// Uses BoringSSL to extract the DER-encoded subject from a DER-encoded X.509 certificate.
+pub fn parse_subject_from_certificate(cert_buf: &[u8]) -> Result<Vec<u8>, Error> {
     // Try with a 200-byte output buffer, should be enough in all but bizarre cases.
     let mut retval = vec![0; 200];
+
+    // Safety: extractSubjectFromCertificate reads at most cert_buf.len() bytes from cert_buf and
+    // writes at most retval.len() bytes to retval.
     let mut size = unsafe {
         extractSubjectFromCertificate(
             cert_buf.as_ptr(),
@@ -374,12 +381,11 @@
 
     if size < 0 {
         // Our buffer wasn't big enough.  Make one that is just the right size and try again.
-        let negated_size = usize::try_from(-size);
-        retval = match negated_size.ok() {
-            None => return Err(Error::ExtractSubjectFailed),
-            Some(size) => vec![0; size],
-        };
+        let negated_size = usize::try_from(-size).map_err(|_e| Error::ExtractSubjectFailed)?;
+        retval = vec![0; negated_size];
 
+        // Safety: extractSubjectFromCertificate reads at most cert_buf.len() bytes from cert_buf
+        // and writes at most retval.len() bytes to retval.
         size = unsafe {
             extractSubjectFromCertificate(
                 cert_buf.as_ptr(),
@@ -395,14 +401,8 @@
     }
 
     // Reduce buffer size to the amount written.
-    let safe_size = usize::try_from(size);
-    retval.resize(
-        match safe_size.ok() {
-            None => return Err(Error::ExtractSubjectFailed),
-            Some(size) => size,
-        },
-        0,
-    );
+    let safe_size = usize::try_from(size).map_err(|_e| Error::ExtractSubjectFailed)?;
+    retval.truncate(safe_size);
 
     Ok(retval)
 }
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 40860be..22ab02e 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -581,11 +581,14 @@
 
 /// This type represents a certificate chain with a private key corresponding to the leaf
 /// certificate. TODO(jbires): This will be used in a follow-on CL, for now it's used in the tests.
-#[allow(dead_code)]
 pub struct CertificateChain {
-    private_key: ZVec,
-    batch_cert: ZVec,
-    cert_chain: ZVec,
+    /// A KM key blob
+    pub private_key: ZVec,
+    /// A batch cert for private_key
+    pub batch_cert: Vec<u8>,
+    /// A full certificate chain from root signing authority to private_key, including batch_cert
+    /// for convenience.
+    pub cert_chain: Vec<u8>,
 }
 
 /// This type represents a Keystore 2.0 key entry.
@@ -730,11 +733,21 @@
         Self(get_current_time_in_seconds())
     }
 
+    /// Constructs a new MonotonicRawTime from a given number of seconds.
+    pub fn from_secs(val: i64) -> Self {
+        Self(val)
+    }
+
     /// Returns the integer value of MonotonicRawTime as i64
     pub fn seconds(&self) -> i64 {
         self.0
     }
 
+    /// Returns the value of MonotonicRawTime in milli seconds as i64
+    pub fn milli_seconds(&self) -> i64 {
+        self.0 * 1000
+    }
+
     /// Like i64::checked_sub.
     pub fn checked_sub(&self, other: &Self) -> Option<Self> {
         self.0.checked_sub(other.0).map(Self)
@@ -787,6 +800,11 @@
     pub fn time_received(&self) -> MonotonicRawTime {
         self.time_received
     }
+
+    /// Returns the challenge value of the auth token.
+    pub fn challenge(&self) -> i64 {
+        self.auth_token.challenge
+    }
 }
 
 /// Shared in-memory databases get destroyed as soon as the last connection to them gets closed.
@@ -1776,6 +1794,33 @@
         .context("In delete_expired_attestation_keys: ")
     }
 
+    /// Deletes all remotely provisioned attestation keys in the system, regardless of the state
+    /// they are in. This is useful primarily as a testing mechanism.
+    pub fn delete_all_attestation_keys(&mut self) -> Result<i64> {
+        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+            let mut stmt = tx
+                .prepare(
+                    "SELECT id FROM persistent.keyentry
+                    WHERE key_type IS ?;",
+                )
+                .context("Failed to prepare statement")?;
+            let keys_to_delete = stmt
+                .query_map(params![KeyType::Attestation], |row| Ok(row.get(0)?))?
+                .collect::<rusqlite::Result<Vec<i64>>>()
+                .context("Failed to execute statement")?;
+            let num_deleted = keys_to_delete
+                .iter()
+                .map(|id| Self::mark_unreferenced(&tx, *id))
+                .collect::<Result<Vec<bool>>>()
+                .context("Failed to execute mark_unreferenced on a keyid")?
+                .into_iter()
+                .filter(|result| *result)
+                .count() as i64;
+            Ok(num_deleted).do_gc(num_deleted != 0)
+        })
+        .context("In delete_all_attestation_keys: ")
+    }
+
     /// Counts the number of keys that will expire by the provided epoch date and the number of
     /// keys not currently assigned to a domain.
     pub fn get_attestation_pool_status(
@@ -1914,8 +1959,8 @@
             }
             Ok(Some(CertificateChain {
                 private_key: ZVec::try_from(km_blob)?,
-                batch_cert: ZVec::try_from(batch_cert_blob)?,
-                cert_chain: ZVec::try_from(cert_chain_blob)?,
+                batch_cert: batch_cert_blob,
+                cert_chain: cert_chain_blob,
             }))
             .no_gc()
         })
@@ -3212,8 +3257,8 @@
         assert_eq!(true, chain.is_some());
         let cert_chain = chain.unwrap();
         assert_eq!(cert_chain.private_key.to_vec(), loaded_values.priv_key);
-        assert_eq!(cert_chain.batch_cert.to_vec(), loaded_values.batch_cert);
-        assert_eq!(cert_chain.cert_chain.to_vec(), loaded_values.cert_chain);
+        assert_eq!(cert_chain.batch_cert, loaded_values.batch_cert);
+        assert_eq!(cert_chain.cert_chain, loaded_values.cert_chain);
         Ok(())
     }
 
@@ -3306,8 +3351,8 @@
             db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace, &KEYSTORE_UUID)?;
         assert!(cert_chain.is_some());
         let value = cert_chain.unwrap();
-        assert_eq!(entry_values.batch_cert, value.batch_cert.to_vec());
-        assert_eq!(entry_values.cert_chain, value.cert_chain.to_vec());
+        assert_eq!(entry_values.batch_cert, value.batch_cert);
+        assert_eq!(entry_values.cert_chain, value.cert_chain);
         assert_eq!(entry_values.priv_key, value.private_key.to_vec());
 
         cert_chain = db.retrieve_attestation_key_and_cert_chain(
@@ -3338,6 +3383,23 @@
     }
 
     #[test]
+    fn test_delete_all_attestation_keys() -> Result<()> {
+        let mut db = new_test_db()?;
+        load_attestation_key_pool(&mut db, 45 /* expiration */, 1 /* namespace */, 0x02)?;
+        load_attestation_key_pool(&mut db, 80 /* expiration */, 2 /* namespace */, 0x03)?;
+        db.create_key_entry(&Domain::APP, &42, &KEYSTORE_UUID)?;
+        let result = db.delete_all_attestation_keys()?;
+
+        // Give the garbage collector half a second to catch up.
+        std::thread::sleep(Duration::from_millis(500));
+
+        // Attestation keys should be deleted, and the regular key should remain.
+        assert_eq!(result, 2);
+
+        Ok(())
+    }
+
+    #[test]
     fn test_rebind_alias() -> Result<()> {
         fn extractor(
             ke: &KeyEntryRow,
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index aa852f4..06648f1 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -14,6 +14,7 @@
 
 //! This is the Keystore 2.0 Enforcements module.
 // TODO: more description to follow.
+use crate::authorization::Error as AuthzError;
 use crate::database::{AuthTokenEntry, MonotonicRawTime};
 use crate::error::{map_binder_status, Error, ErrorCode};
 use crate::globals::{get_timestamp_service, ASYNC_TASK, DB, ENFORCEMENTS};
@@ -26,6 +27,7 @@
 use android_hardware_security_secureclock::aidl::android::hardware::security::secureclock::{
     ISecureClock::ISecureClock, TimeStampToken::TimeStampToken,
 };
+use android_security_authorization::aidl::android::security::authorization::ResponseCode::ResponseCode as AuthzResponseCode;
 use android_system_keystore2::aidl::android::system::keystore2::{
     IKeystoreSecurityLevel::KEY_FLAG_AUTH_BOUND_WITHOUT_CRYPTOGRAPHIC_LSKF_BINDING,
     OperationChallenge::OperationChallenge,
@@ -218,7 +220,7 @@
     if let Err(e) = sender.send(get_timestamp_token(challenge)) {
         log::info!(
             concat!(
-                "In timestamp_token_request: Operation hung up ",
+                "In timestamp_token_request: Receiver hung up ",
                 "before timestamp token could be delivered. {:?}"
             ),
             e
@@ -767,6 +769,66 @@
         auth_bound && !skip_lskf_binding
     }
 
+    /// Finds a matching auth token along with a timestamp token.
+    /// This method looks through auth-tokens cached by keystore which satisfy the given
+    /// authentication information (i.e. |secureUserId|).
+    /// The most recent matching auth token which has a |challenge| field which matches
+    /// the passed-in |challenge| parameter is returned.
+    /// In this case the |authTokenMaxAgeMillis| parameter is not used.
+    ///
+    /// Otherwise, the most recent matching auth token which is younger than |authTokenMaxAgeMillis|
+    /// is returned.
+    pub fn get_auth_tokens(
+        &self,
+        challenge: i64,
+        secure_user_id: i64,
+        auth_token_max_age_millis: i64,
+    ) -> Result<(HardwareAuthToken, TimeStampToken)> {
+        let auth_type = HardwareAuthenticatorType::ANY;
+        let sids: Vec<i64> = vec![secure_user_id];
+        // Filter the matching auth tokens by challenge
+        let result = Self::find_auth_token(|hat: &AuthTokenEntry| {
+            (challenge == hat.challenge()) && hat.satisfies(&sids, auth_type)
+        })
+        .context(
+            "In get_auth_tokens: Failed to get a matching auth token filtered by challenge.",
+        )?;
+
+        let auth_token = if let Some((auth_token_entry, _)) = result {
+            auth_token_entry.take_auth_token()
+        } else {
+            // Filter the matching auth tokens by age.
+            if auth_token_max_age_millis != 0 {
+                let now_in_millis = MonotonicRawTime::now().milli_seconds();
+                let result = Self::find_auth_token(|auth_token_entry: &AuthTokenEntry| {
+                    let token_valid = now_in_millis
+                        .checked_sub(auth_token_entry.time_received().milli_seconds())
+                        .map_or(false, |token_age_in_millis| {
+                            token_age_in_millis > auth_token_max_age_millis
+                        });
+                    token_valid && auth_token_entry.satisfies(&sids, auth_type)
+                })
+                .context(
+                    "In get_auth_tokens: Failed to get a matching auth token filtered by age.",
+                )?;
+
+                if let Some((auth_token_entry, _)) = result {
+                    auth_token_entry.take_auth_token()
+                } else {
+                    return Err(AuthzError::Rc(AuthzResponseCode::NO_AUTH_TOKEN_FOUND))
+                        .context("In get_auth_tokens: No auth token found.");
+                }
+            } else {
+                return Err(AuthzError::Rc(AuthzResponseCode::NO_AUTH_TOKEN_FOUND))
+                    .context("In get_auth_tokens: Passed-in auth token max age is zero.");
+            }
+        };
+        // Wait and obtain the timestamp token from secure clock service.
+        let tst = get_timestamp_token(challenge)
+            .context("In get_auth_tokens. Error in getting timestamp token.")?;
+        Ok((auth_token, tst))
+    }
+
     /// Watch the `keystore.boot_level` system property, and keep self.boot_level up to date.
     /// Blocks waiting for system property changes, so must be run in its own thread.
     pub fn watch_boot_level(&self) -> Result<()> {
diff --git a/keystore2/src/entropy.rs b/keystore2/src/entropy.rs
new file mode 100644
index 0000000..de38187
--- /dev/null
+++ b/keystore2/src/entropy.rs
@@ -0,0 +1,98 @@
+// Copyright 2021, The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! This module holds functionality for retrieving and distributing entropy.
+
+use anyhow::{Context, Result};
+use log::error;
+use std::time::{Duration, Instant};
+
+static ENTROPY_SIZE: usize = 64;
+static MIN_FEED_INTERVAL_SECS: u64 = 30;
+
+#[derive(Default)]
+struct FeederInfo {
+    last_feed: Option<Instant>,
+}
+
+/// Register the entropy feeder as an idle callback.
+pub fn register_feeder() {
+    crate::globals::ASYNC_TASK.add_idle(|shelf| {
+        let mut info = shelf.get_mut::<FeederInfo>();
+        let now = Instant::now();
+        let feed_needed = match info.last_feed {
+            None => true,
+            Some(last) => now.duration_since(last) > Duration::from_secs(MIN_FEED_INTERVAL_SECS),
+        };
+        if feed_needed {
+            info.last_feed = Some(now);
+            feed_devices();
+        }
+    });
+}
+
+fn get_entropy(size: usize) -> Result<Vec<u8>> {
+    keystore2_crypto::generate_random_data(size).context("Retrieving entropy for KeyMint device")
+}
+
+/// Feed entropy to all known KeyMint devices.
+pub fn feed_devices() {
+    let km_devs = crate::globals::get_keymint_devices();
+    if km_devs.is_empty() {
+        return;
+    }
+    let data = match get_entropy(km_devs.len() * ENTROPY_SIZE) {
+        Ok(data) => data,
+        Err(e) => {
+            error!(
+                "Failed to retrieve {}*{} bytes of entropy: {:?}",
+                km_devs.len(),
+                ENTROPY_SIZE,
+                e
+            );
+            return;
+        }
+    };
+    for (i, km_dev) in km_devs.iter().enumerate() {
+        let offset = i * ENTROPY_SIZE;
+        let sub_data = &data[offset..(offset + ENTROPY_SIZE)];
+        if let Err(e) = km_dev.addRngEntropy(sub_data) {
+            error!("Failed to feed entropy to KeyMint device: {:?}", e);
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::HashSet;
+
+    #[test]
+    fn test_entropy_size() {
+        for size in &[0, 1, 4, 8, 256, 4096] {
+            let data = get_entropy(*size).expect("failed to get entropy");
+            assert_eq!(data.len(), *size);
+        }
+    }
+    #[test]
+    fn test_entropy_uniqueness() {
+        let count = 10;
+        let mut seen = HashSet::new();
+        for _i in 0..count {
+            let data = get_entropy(16).expect("failed to get entropy");
+            seen.insert(data);
+        }
+        assert_eq!(seen.len(), count);
+    }
+}
diff --git a/keystore2/src/error.rs b/keystore2/src/error.rs
index d67f5f4..388487c 100644
--- a/keystore2/src/error.rs
+++ b/keystore2/src/error.rs
@@ -171,9 +171,31 @@
 where
     F: FnOnce(U) -> BinderResult<T>,
 {
-    result.map_or_else(
+    map_err_with(
+        result,
         |e| {
             log::error!("{:?}", e);
+            e
+        },
+        handle_ok,
+    )
+}
+
+/// This function behaves similar to map_or_log_error, but it does not log the errors, instead
+/// it calls map_err on the error before mapping it to a binder result allowing callers to
+/// log or transform the error before mapping it.
+pub fn map_err_with<T, U, F1, F2>(
+    result: anyhow::Result<U>,
+    map_err: F1,
+    handle_ok: F2,
+) -> BinderResult<T>
+where
+    F1: FnOnce(anyhow::Error) -> anyhow::Error,
+    F2: FnOnce(U) -> BinderResult<T>,
+{
+    result.map_or_else(
+        |e| {
+            let e = map_err(e);
             let root_cause = e.root_cause();
             let rc = match root_cause.downcast_ref::<Error>() {
                 Some(Error::Rc(rcode)) => rcode.0,
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index 04bfbc9..54f7dc7 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -35,6 +35,7 @@
 use android_hardware_security_keymint::binder::{StatusCode, Strong};
 use android_security_compat::aidl::android::security::compat::IKeystoreCompatService::IKeystoreCompatService;
 use anyhow::{Context, Result};
+use binder::FromIBinder;
 use keystore2_vintf::get_aidl_instances;
 use lazy_static::lazy_static;
 use std::sync::{Arc, Mutex};
@@ -117,6 +118,10 @@
             .map(|(dev, hw_info)| ((*dev).clone(), (*hw_info).clone(), *uuid))
     }
 
+    fn devices<T: FromIBinder + ?Sized>(&self) -> Vec<Strong<T>> {
+        self.devices_by_uuid.values().filter_map(|(asp, _)| asp.get_interface::<T>().ok()).collect()
+    }
+
     /// The requested security level and the security level of the actual implementation may
     /// differ. So we map the requested security level to the uuid of the implementation
     /// so that there cannot be any confusion as to which KeyMint instance is requested.
@@ -256,6 +261,11 @@
     }
 }
 
+/// Return all known keymint devices.
+pub fn get_keymint_devices() -> Vec<Strong<dyn IKeyMintDevice>> {
+    KEY_MINT_DEVICES.lock().unwrap().devices()
+}
+
 static TIME_STAMP_SERVICE_NAME: &str = "android.hardware.security.secureclock.ISecureClock";
 
 /// Make a new connection to a secure clock service.
diff --git a/keystore2/src/keystore2_main.rs b/keystore2/src/keystore2_main.rs
index 7739f5e..5d99449 100644
--- a/keystore2/src/keystore2_main.rs
+++ b/keystore2/src/keystore2_main.rs
@@ -16,10 +16,11 @@
 
 use keystore2::apc::ApcManager;
 use keystore2::authorization::AuthorizationManager;
+use keystore2::entropy;
 use keystore2::globals::ENFORCEMENTS;
 use keystore2::remote_provisioning::RemoteProvisioningService;
 use keystore2::service::KeystoreService;
-use keystore2::user_manager::UserManager;
+use keystore2::user_manager::Maintenance;
 use log::{error, info};
 use std::{panic, path::Path, sync::mpsc::channel};
 use vpnprofilestore::VpnProfileStore;
@@ -28,7 +29,7 @@
 static APC_SERVICE_NAME: &str = "android.security.apc";
 static AUTHORIZATION_SERVICE_NAME: &str = "android.security.authorization";
 static REMOTE_PROVISIONING_SERVICE_NAME: &str = "android.security.remoteprovisioning";
-static USER_MANAGER_SERVICE_NAME: &str = "android.security.usermanager";
+static USER_MANAGER_SERVICE_NAME: &str = "android.security.maintenance";
 static VPNPROFILESTORE_SERVICE_NAME: &str = "android.security.vpnprofilestore";
 
 /// Keystore 2.0 takes one argument which is a path indicating its designated working directory.
@@ -74,6 +75,8 @@
             .unwrap_or_else(|e| error!("watch_boot_level failed: {}", e));
     });
 
+    entropy::register_feeder();
+
     info!("Starting thread pool now.");
     binder::ProcessState::start_thread_pool();
 
@@ -100,10 +103,10 @@
             panic!("Failed to register service {} because of {:?}.", AUTHORIZATION_SERVICE_NAME, e);
         });
 
-    let usermanager_service = UserManager::new_native_binder().unwrap_or_else(|e| {
+    let maintenance_service = Maintenance::new_native_binder().unwrap_or_else(|e| {
         panic!("Failed to create service {} because of {:?}.", USER_MANAGER_SERVICE_NAME, e);
     });
-    binder::add_service(USER_MANAGER_SERVICE_NAME, usermanager_service.as_binder()).unwrap_or_else(
+    binder::add_service(USER_MANAGER_SERVICE_NAME, maintenance_service.as_binder()).unwrap_or_else(
         |e| {
             panic!("Failed to register service {} because of {:?}.", USER_MANAGER_SERVICE_NAME, e);
         },
diff --git a/keystore2/src/km_compat/km_compat.cpp b/keystore2/src/km_compat/km_compat.cpp
index 3d43918..26b099a 100644
--- a/keystore2/src/km_compat/km_compat.cpp
+++ b/keystore2/src/km_compat/km_compat.cpp
@@ -216,6 +216,8 @@
  */
 bool isNewAndKeystoreEnforceable(const KMV1::KeyParameter& param) {
     switch (param.tag) {
+    case KMV1::Tag::MAX_BOOT_LEVEL:
+        return true;
     case KMV1::Tag::USAGE_COUNT_LIMIT:
         return true;
     default:
@@ -511,15 +513,23 @@
 
 ScopedAStatus
 KeyMintDevice::importWrappedKey(const std::vector<uint8_t>& in_inWrappedKeyData,
-                                const std::vector<uint8_t>& in_inWrappingKeyBlob,  //
+                                const std::vector<uint8_t>& in_inPrefixedWrappingKeyBlob,
                                 const std::vector<uint8_t>& in_inMaskingKey,
                                 const std::vector<KeyParameter>& in_inUnwrappingParams,
                                 int64_t in_inPasswordSid, int64_t in_inBiometricSid,
                                 KeyCreationResult* out_creationResult) {
+    const std::vector<uint8_t>& wrappingKeyBlob =
+        prefixedKeyBlobRemovePrefix(in_inPrefixedWrappingKeyBlob);
+    if (prefixedKeyBlobIsSoftKeyMint(in_inPrefixedWrappingKeyBlob)) {
+        return softKeyMintDevice_->importWrappedKey(
+            in_inWrappedKeyData, wrappingKeyBlob, in_inMaskingKey, in_inUnwrappingParams,
+            in_inPasswordSid, in_inBiometricSid, out_creationResult);
+    }
+
     auto legacyUnwrappingParams = convertKeyParametersToLegacy(in_inUnwrappingParams);
     KMV1::ErrorCode errorCode;
     auto result = mDevice->importWrappedKey(
-        in_inWrappedKeyData, in_inWrappingKeyBlob, in_inMaskingKey, legacyUnwrappingParams,
+        in_inWrappedKeyData, wrappingKeyBlob, in_inMaskingKey, legacyUnwrappingParams,
         in_inPasswordSid, in_inBiometricSid,
         [&](V4_0_ErrorCode error, const hidl_vec<uint8_t>& keyBlob,
             const V4_0_KeyCharacteristics& keyCharacteristics) {
@@ -540,11 +550,12 @@
                                         std::vector<uint8_t>* _aidl_return) {
     auto legacyUpgradeParams = convertKeyParametersToLegacy(in_inUpgradeParams);
     V4_0_ErrorCode errorCode;
+
     auto result =
-        mDevice->upgradeKey(in_inKeyBlobToUpgrade, legacyUpgradeParams,
+        mDevice->upgradeKey(prefixedKeyBlobRemovePrefix(in_inKeyBlobToUpgrade), legacyUpgradeParams,
                             [&](V4_0_ErrorCode error, const hidl_vec<uint8_t>& upgradedKeyBlob) {
                                 errorCode = error;
-                                *_aidl_return = upgradedKeyBlob;
+                                *_aidl_return = keyBlobPrefix(upgradedKeyBlob, false);
                             });
     if (!result.isOk()) {
         LOG(ERROR) << __func__ << " transaction failed. " << result.description();
@@ -556,7 +567,7 @@
 ScopedAStatus KeyMintDevice::deleteKey(const std::vector<uint8_t>& prefixedKeyBlob) {
     const std::vector<uint8_t>& keyBlob = prefixedKeyBlobRemovePrefix(prefixedKeyBlob);
     if (prefixedKeyBlobIsSoftKeyMint(prefixedKeyBlob)) {
-        return softKeyMintDevice_->deleteKey(prefixedKeyBlob);
+        return softKeyMintDevice_->deleteKey(keyBlob);
     }
 
     auto result = mDevice->deleteKey(keyBlob);
@@ -658,8 +669,8 @@
 
     KMV1::ErrorCode errorCode;
     auto result = mDevice->update(
-        mOperationHandle, {V4_0::makeKeyParameter(V4_0::TAG_ASSOCIATED_DATA, input)}, input,
-        authToken, verificationToken,
+        mOperationHandle, {V4_0::makeKeyParameter(V4_0::TAG_ASSOCIATED_DATA, input)}, {}, authToken,
+        verificationToken,
         [&](V4_0_ErrorCode error, auto, auto, auto) { errorCode = convert(error); });
 
     if (!result.isOk()) {
@@ -723,7 +734,7 @@
 
     KMV1::ErrorCode errorCode;
     auto result = mDevice->finish(
-        mOperationHandle, {} /* inParams */, input, signature, authToken, verificationToken,
+        mOperationHandle, inParams, input, signature, authToken, verificationToken,
         [&](V4_0_ErrorCode error, auto /* outParams */, const hidl_vec<uint8_t>& output) {
             errorCode = convert(error);
             *out_output = output;
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 8fef6cf..0e51eff 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -20,6 +20,7 @@
 pub mod authorization;
 pub mod database;
 pub mod enforcements;
+pub mod entropy;
 pub mod error;
 pub mod globals;
 /// Internal Representation of Key Parameter and convenience functions.
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 1de24c7..0f5bcaf 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -126,7 +126,7 @@
 //! Either way, we have to revaluate the pruning scores.
 
 use crate::enforcements::AuthInfo;
-use crate::error::{map_km_error, map_or_log_err, Error, ErrorCode, ResponseCode};
+use crate::error::{map_err_with, map_km_error, map_or_log_err, Error, ErrorCode, ResponseCode};
 use crate::utils::Asp;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     IKeyMintOperation::IKeyMintOperation,
@@ -135,7 +135,7 @@
     IKeystoreOperation::BnKeystoreOperation, IKeystoreOperation::IKeystoreOperation,
 };
 use anyhow::{anyhow, Context, Result};
-use binder::IBinder;
+use binder::IBinderInternal;
 use std::{
     collections::HashMap,
     sync::{Arc, Mutex, MutexGuard, Weak},
@@ -729,7 +729,7 @@
 impl KeystoreOperation {
     /// Creates a new operation instance wrapped in a
     /// BnKeystoreOperation proxy object. It also
-    /// calls `IBinder::set_requesting_sid` on the new interface, because
+    /// calls `IBinderInternal::set_requesting_sid` on the new interface, because
     /// we need it for checking Keystore permissions.
     pub fn new_native_binder(
         operation: Arc<Operation>,
@@ -815,11 +815,21 @@
     }
 
     fn abort(&self) -> binder::public_api::Result<()> {
-        map_or_log_err(
+        map_err_with(
             self.with_locked_operation(
                 |op| op.abort(Outcome::Abort).context("In KeystoreOperation::abort"),
                 true,
             ),
+            |e| {
+                match e.root_cause().downcast_ref::<Error>() {
+                    // Calling abort on expired operations is something very common.
+                    // There is no reason to clutter the log with it. It is never the cause
+                    // for a true problem.
+                    Some(Error::Km(ErrorCode::INVALID_OPERATION_HANDLE)) => {}
+                    _ => log::error!("{:?}", e),
+                };
+                e
+            },
             Ok,
         )
     }
diff --git a/keystore2/src/permission.rs b/keystore2/src/permission.rs
index 576ac3f..7f63834 100644
--- a/keystore2/src/permission.rs
+++ b/keystore2/src/permission.rs
@@ -291,7 +291,7 @@
         AddAuth = 1,    selinux name: add_auth;
         /// Checked when an app is uninstalled or wiped.
         ClearNs = 2,    selinux name: clear_ns;
-        /// Checked when Keystore 2.0 gets locked.
+        /// Checked when the user state is queried from Keystore 2.0.
         GetState = 4,   selinux name: get_state;
         /// Checked when Keystore 2.0 is asked to list a namespace that the caller
         /// does not have the get_info permission for.
@@ -308,6 +308,8 @@
         ChangePassword = 0x100,    selinux name: change_password;
         /// Checked when a UID is cleared.
         ClearUID = 0x200,    selinux name: clear_uid;
+        /// Checked when Credstore calls IKeystoreAuthorization to obtain auth tokens.
+        GetAuthToken = 0x400,  selinux name: get_auth_token;
     }
 );
 
@@ -673,7 +675,7 @@
         let shell_ctx = Context::new("u:r:shell:s0")?;
         assert_perm_failed!(check_keystore_permission(&shell_ctx, KeystorePerm::add_auth()));
         assert_perm_failed!(check_keystore_permission(&shell_ctx, KeystorePerm::clear_ns()));
-        assert_perm_failed!(check_keystore_permission(&shell_ctx, KeystorePerm::get_state()));
+        assert!(check_keystore_permission(&shell_ctx, KeystorePerm::get_state()).is_ok());
         assert_perm_failed!(check_keystore_permission(&shell_ctx, KeystorePerm::list()));
         assert_perm_failed!(check_keystore_permission(&shell_ctx, KeystorePerm::lock()));
         assert_perm_failed!(check_keystore_permission(&shell_ctx, KeystorePerm::reset()));
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index d606b6a..cc97573 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -22,20 +22,184 @@
 use std::collections::HashMap;
 
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    IRemotelyProvisionedComponent::IRemotelyProvisionedComponent, MacedPublicKey::MacedPublicKey,
-    ProtectedData::ProtectedData, SecurityLevel::SecurityLevel,
+    Algorithm::Algorithm, AttestationKey::AttestationKey, Certificate::Certificate,
+    IRemotelyProvisionedComponent::IRemotelyProvisionedComponent, KeyParameter::KeyParameter,
+    KeyParameterValue::KeyParameterValue, MacedPublicKey::MacedPublicKey,
+    ProtectedData::ProtectedData, SecurityLevel::SecurityLevel, Tag::Tag,
 };
 use android_security_remoteprovisioning::aidl::android::security::remoteprovisioning::{
     AttestationPoolStatus::AttestationPoolStatus, IRemoteProvisioning::BnRemoteProvisioning,
     IRemoteProvisioning::IRemoteProvisioning,
 };
 use android_security_remoteprovisioning::binder::Strong;
+use android_system_keystore2::aidl::android::system::keystore2::{
+    Domain::Domain, KeyDescriptor::KeyDescriptor,
+};
 use anyhow::{Context, Result};
+use keystore2_crypto::parse_subject_from_certificate;
+use std::sync::atomic::{AtomicBool, Ordering};
 
-use crate::error::{self, map_or_log_err, map_rem_prov_error};
+use crate::database::{CertificateChain, KeystoreDB, Uuid};
+use crate::error::{self, map_or_log_err, map_rem_prov_error, Error};
 use crate::globals::{get_keymint_device, get_remotely_provisioned_component, DB};
 use crate::utils::Asp;
 
+/// Contains helper functions to check if remote provisioning is enabled on the system and, if so,
+/// to assign and retrieve attestation keys and certificate chains.
+#[derive(Default)]
+pub struct RemProvState {
+    security_level: SecurityLevel,
+    km_uuid: Uuid,
+    is_hal_present: AtomicBool,
+}
+
+impl RemProvState {
+    /// Creates a RemProvState struct.
+    pub fn new(security_level: SecurityLevel, km_uuid: Uuid) -> Self {
+        Self { security_level, km_uuid, is_hal_present: AtomicBool::new(true) }
+    }
+
+    /// Checks if remote provisioning is enabled and partially caches the result. On a hybrid system
+    /// remote provisioning can flip from being disabled to enabled depending on responses from the
+    /// server, so unfortunately caching the presence or absence of the HAL is not enough to fully
+    /// make decisions about the state of remote provisioning during runtime.
+    fn check_rem_prov_enabled(&self, db: &mut KeystoreDB) -> Result<bool> {
+        if !self.is_hal_present.load(Ordering::Relaxed)
+            || get_remotely_provisioned_component(&self.security_level).is_err()
+        {
+            self.is_hal_present.store(false, Ordering::Relaxed);
+            return Ok(false);
+        }
+        // To check if remote provisioning is enabled on a system that supports both remote
+        // provisioning and factory provisioned keys, we only need to check if there are any
+        // keys at all generated to indicate if the app has gotten the signal to begin filling
+        // the key pool from the server.
+        let pool_status = db
+            .get_attestation_pool_status(0 /* date */, &self.km_uuid)
+            .context("In check_rem_prov_enabled: failed to get attestation pool status.")?;
+        Ok(pool_status.total != 0)
+    }
+
+    /// Fetches a remote provisioning attestation key and certificate chain inside of the
+    /// returned `CertificateChain` struct if one exists for the given caller_uid. If one has not
+    /// been assigned, this function will assign it. If there are no signed attestation keys
+    /// available to be assigned, it will return the ResponseCode `OUT_OF_KEYS`
+    fn get_rem_prov_attest_key(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        db: &mut KeystoreDB,
+    ) -> Result<Option<CertificateChain>> {
+        match key.domain {
+            Domain::APP => {
+                // Attempt to get an Attestation Key once. If it fails, then the app doesn't
+                // have a valid chain assigned to it. The helper function will return None after
+                // attempting to assign a key. An error will be thrown if the pool is simply out
+                // of usable keys. Then another attempt to fetch the just-assigned key will be
+                // made. If this fails too, something is very wrong.
+                self.get_rem_prov_attest_key_helper(key, caller_uid, db)
+                    .context("In get_rem_prov_attest_key: Failed to get a key")?
+                    .map_or_else(
+                        || self.get_rem_prov_attest_key_helper(key, caller_uid, db),
+                        |v| Ok(Some(v)),
+                    )
+                    .context(concat!(
+                        "In get_rem_prov_attest_key: Failed to get a key after",
+                        "attempting to assign one."
+                    ))?
+                    .map_or_else(
+                        || {
+                            Err(Error::sys()).context(concat!(
+                                "In get_rem_prov_attest_key: Attempted to assign a ",
+                                "key and failed silently. Something is very wrong."
+                            ))
+                        },
+                        |cert_chain| Ok(Some(cert_chain)),
+                    )
+            }
+            _ => Ok(None),
+        }
+    }
+
+    /// Returns None if an AttestationKey fails to be assigned. Errors if no keys are available.
+    fn get_rem_prov_attest_key_helper(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        db: &mut KeystoreDB,
+    ) -> Result<Option<CertificateChain>> {
+        let cert_chain = db
+            .retrieve_attestation_key_and_cert_chain(key.domain, caller_uid as i64, &self.km_uuid)
+            .context("In get_rem_prov_attest_key_helper: Failed to retrieve a key + cert chain")?;
+        match cert_chain {
+            Some(cert_chain) => Ok(Some(cert_chain)),
+            // Either this app needs to be assigned a key, or the pool is empty. An error will
+            // be thrown if there is no key available to assign. This will indicate that the app
+            // should be nudged to provision more keys so keystore can retry.
+            None => {
+                db.assign_attestation_key(key.domain, caller_uid as i64, &self.km_uuid)
+                    .context("In get_rem_prov_attest_key_helper: Failed to assign a key")?;
+                Ok(None)
+            }
+        }
+    }
+
+    fn is_asymmetric_key(&self, params: &[KeyParameter]) -> bool {
+        params.iter().any(|kp| {
+            matches!(
+                kp,
+                KeyParameter {
+                    tag: Tag::ALGORITHM,
+                    value: KeyParameterValue::Algorithm(Algorithm::RSA)
+                } | KeyParameter {
+                    tag: Tag::ALGORITHM,
+                    value: KeyParameterValue::Algorithm(Algorithm::EC)
+                }
+            )
+        })
+    }
+
+    /// Checks to see (1) if the key in question should be attested to based on the algorithm and
+    /// (2) if remote provisioning is present and enabled on the system. If these conditions are
+    /// met, it makes an attempt to fetch the attestation key assigned to the `caller_uid`.
+    ///
+    /// It returns the ResponseCode `OUT_OF_KEYS` if there is not one key currently assigned to the
+    /// `caller_uid` and there are none available to assign.
+    pub fn get_remote_provisioning_key_and_certs(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        params: &[KeyParameter],
+        db: &mut KeystoreDB,
+    ) -> Result<(Option<AttestationKey>, Option<Certificate>)> {
+        if !self.is_asymmetric_key(params) || !self.check_rem_prov_enabled(db)? {
+            // There is no remote provisioning component for this security level on the
+            // device. Return None so the underlying KM instance knows to use its
+            // factory provisioned key instead. Alternatively, it's not an asymmetric key
+            // and therefore will not be attested.
+            Ok((None, None))
+        } else {
+            match self.get_rem_prov_attest_key(&key, caller_uid, db).context(concat!(
+                "In get_remote_provisioning_key_and_certs: Failed to get ",
+                "attestation key"
+            ))? {
+                Some(cert_chain) => Ok((
+                    Some(AttestationKey {
+                        keyBlob: cert_chain.private_key.to_vec(),
+                        attestKeyParams: vec![],
+                        issuerSubjectName: parse_subject_from_certificate(&cert_chain.batch_cert)
+                            .context(concat!(
+                            "In get_remote_provisioning_key_and_certs: Failed to ",
+                            "parse subject."
+                        ))?,
+                    }),
+                    Some(Certificate { encodedCertificate: cert_chain.cert_chain }),
+                )),
+                None => Ok((None, None)),
+            }
+        }
+    }
+}
 /// Implementation of the IRemoteProvisioning service.
 #[derive(Default)]
 pub struct RemoteProvisioningService {
@@ -126,7 +290,21 @@
             protected_data,
         ))
         .context("In generate_csr: Failed to generate csr")?;
-        Ok(mac)
+        let mut cose_mac_0 = Vec::<u8>::new();
+        // TODO(b/180392379): Replace this manual CBOR generation with the cbor-serde crate as well.
+        //                    This generates an array consisting of the mac and the public key Maps.
+        //                    Just generate the actual MacedPublicKeys structure when the crate is
+        //                    available.
+        cose_mac_0.push((0b100_00000 | (keys_to_sign.len() + 1)) as u8);
+        cose_mac_0.push(0b010_11000); //push mac
+        cose_mac_0.push(mac.len() as u8);
+        cose_mac_0.append(&mut mac);
+        for maced_public_key in keys_to_sign {
+            if maced_public_key.macedKey.len() > 83 + 8 {
+                cose_mac_0.extend_from_slice(&maced_public_key.macedKey[8..83 + 8]);
+            }
+        }
+        Ok(cose_mac_0)
     }
 
     /// Provisions a certificate chain for a key whose CSR was included in generate_csr. The
@@ -189,6 +367,15 @@
     pub fn get_security_levels(&self) -> Result<Vec<SecurityLevel>> {
         Ok(self.device_by_sec_level.keys().cloned().collect())
     }
+
+    /// Deletes all attestation keys generated by the IRemotelyProvisionedComponent from the device,
+    /// regardless of what state of the attestation key lifecycle they were in.
+    pub fn delete_all_keys(&self) -> Result<i64> {
+        DB.with::<_, Result<i64>>(|db| {
+            let mut db = db.borrow_mut();
+            Ok(db.delete_all_attestation_keys()?)
+        })
+    }
 }
 
 impl binder::Interface for RemoteProvisioningService {}
@@ -244,4 +431,8 @@
     fn getSecurityLevels(&self) -> binder::public_api::Result<Vec<SecurityLevel>> {
         map_or_log_err(self.get_security_levels(), Ok)
     }
+
+    fn deleteAllKeys(&self) -> binder::public_api::Result<i64> {
+        map_or_log_err(self.delete_all_keys(), Ok)
+    }
 }
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 8368c93..2ad8fe1 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -18,7 +18,7 @@
 
 use crate::globals::get_keymint_device;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    Algorithm::Algorithm, AttestationKey::AttestationKey,
+    Algorithm::Algorithm, AttestationKey::AttestationKey, Certificate::Certificate,
     HardwareAuthenticatorType::HardwareAuthenticatorType, IKeyMintDevice::IKeyMintDevice,
     KeyCreationResult::KeyCreationResult, KeyFormat::KeyFormat,
     KeyMintHardwareInfo::KeyMintHardwareInfo, KeyParameter::KeyParameter,
@@ -32,12 +32,16 @@
     KeyMetadata::KeyMetadata, KeyParameters::KeyParameters,
 };
 
-use crate::database::{CertificateInfo, KeyIdGuard};
+use crate::database::{CertificateInfo, KeyIdGuard, KeystoreDB};
 use crate::globals::{DB, ENFORCEMENTS, LEGACY_MIGRATOR, SUPER_KEY};
 use crate::key_parameter::KeyParameter as KsKeyParam;
 use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
+use crate::remote_provisioning::RemProvState;
 use crate::super_key::{KeyBlob, SuperKeyManager};
-use crate::utils::{check_key_permission, uid_to_android_user, Asp};
+use crate::utils::{
+    check_device_attestation_permissions, check_key_permission, is_device_id_attestation_tag,
+    uid_to_android_user, Asp,
+};
 use crate::{
     database::{
         BlobMetaData, BlobMetaEntry, DateTime, KeyEntry, KeyEntryLoadBits, KeyMetaData,
@@ -52,8 +56,8 @@
     utils::key_characteristics_to_internal,
 };
 use anyhow::{anyhow, Context, Result};
-use binder::{IBinder, Strong, ThreadState};
-use keystore2_crypto::parse_issuer_subject_from_certificate;
+use binder::{IBinderInternal, Strong, ThreadState};
+use keystore2_crypto::parse_subject_from_certificate;
 
 /// Implementation of the IKeystoreSecurityLevel Interface.
 pub struct KeystoreSecurityLevel {
@@ -63,6 +67,7 @@
     hw_info: KeyMintHardwareInfo,
     km_uuid: Uuid,
     operation_db: OperationDb,
+    rem_prov_state: RemProvState,
 }
 
 // Blob of 32 zeroes used as empty masking key.
@@ -75,7 +80,7 @@
 impl KeystoreSecurityLevel {
     /// Creates a new security level instance wrapped in a
     /// BnKeystoreSecurityLevel proxy object. It also
-    /// calls `IBinder::set_requesting_sid` on the new interface, because
+    /// calls `IBinderInternal::set_requesting_sid` on the new interface, because
     /// we need it for checking keystore permissions.
     pub fn new_native_binder(
         security_level: SecurityLevel,
@@ -88,6 +93,7 @@
             hw_info,
             km_uuid,
             operation_db: OperationDb::new(),
+            rem_prov_state: RemProvState::new(security_level, km_uuid),
         });
         result.as_binder().set_requesting_sid(true);
         Ok((result, km_uuid))
@@ -131,33 +137,34 @@
             SecurityLevel::SOFTWARE,
         ));
 
-        let (key_blob, mut blob_metadata) = DB
-            .with(|db| {
-                SUPER_KEY.handle_super_encryption_on_key_init(
-                    &mut db.borrow_mut(),
-                    &LEGACY_MIGRATOR,
-                    &(key.domain),
-                    &key_parameters,
-                    flags,
-                    user_id,
-                    &key_blob,
-                )
-            })
-            .context("In store_new_key. Failed to handle super encryption.")?;
-
         let creation_date = DateTime::now().context("Trying to make creation time.")?;
 
         let key = match key.domain {
-            Domain::BLOB => {
-                KeyDescriptor { domain: Domain::BLOB, blob: Some(key_blob), ..Default::default() }
-            }
+            Domain::BLOB => KeyDescriptor {
+                domain: Domain::BLOB,
+                blob: Some(key_blob.to_vec()),
+                ..Default::default()
+            },
             _ => DB
                 .with::<_, Result<KeyDescriptor>>(|db| {
+                    let mut db = db.borrow_mut();
+
+                    let (key_blob, mut blob_metadata) = SUPER_KEY
+                        .handle_super_encryption_on_key_init(
+                            &mut db,
+                            &LEGACY_MIGRATOR,
+                            &(key.domain),
+                            &key_parameters,
+                            flags,
+                            user_id,
+                            &key_blob,
+                        )
+                        .context("In store_new_key. Failed to handle super encryption.")?;
+
                     let mut key_metadata = KeyMetaData::new();
                     key_metadata.add(KeyMetaEntry::CreationDate(creation_date));
                     blob_metadata.add(BlobMetaEntry::KmUuid(self.km_uuid));
 
-                    let mut db = db.borrow_mut();
                     let key_id = db
                         .store_new_key(
                             &key,
@@ -274,10 +281,7 @@
                 purpose,
                 key_properties.as_ref(),
                 operation_parameters.as_ref(),
-                // TODO b/178222844 Replace this with the configuration returned by
-                //      KeyMintDevice::getHardwareInfo.
-                //      For now we assume that strongbox implementations need secure timestamps.
-                self.security_level == SecurityLevel::STRONGBOX,
+                self.hw_info.timestampTokenRequired,
             )
             .context("In create_operation.")?;
 
@@ -364,6 +368,15 @@
             ))?;
         }
 
+        // If the caller requests any device identifier attestation tag, check that they hold the
+        // correct Android permission.
+        if params.iter().any(|kp| is_device_id_attestation_tag(kp.tag)) {
+            check_device_attestation_permissions().context(concat!(
+                "In add_certificate_parameters: ",
+                "Caller does not have the permission to attest device identifiers."
+            ))?;
+        }
+
         // If we are generating/importing an asymmetric key, we need to make sure
         // that NOT_BEFORE and NOT_AFTER are present.
         match params.iter().find(|kp| kp.tag == Tag::ALGORITHM) {
@@ -413,34 +426,65 @@
 
         // generate_key requires the rebind permission.
         check_key_permission(KeyPerm::rebind(), &key, &None).context("In generate_key.")?;
-
-        let attest_key = match attest_key_descriptor {
-            None => None,
-            Some(key) => Some(
-                self.get_attest_key(key, caller_uid)
-                    .context("In generate_key: Trying to load attest key")?,
-            ),
+        let (attest_key, cert_chain) = match (key.domain, attest_key_descriptor) {
+            (Domain::BLOB, None) => (None, None),
+            _ => DB
+                .with::<_, Result<(Option<AttestationKey>, Option<Certificate>)>>(|db| {
+                    self.get_attest_key_and_cert_chain(
+                        &key,
+                        caller_uid,
+                        attest_key_descriptor,
+                        params,
+                        &mut db.borrow_mut(),
+                    )
+                })
+                .context("In generate_key: Trying to get an attestation key")?,
         };
-
         let params = Self::add_certificate_parameters(caller_uid, params, &key)
             .context("In generate_key: Trying to get aaid.")?;
 
         let km_dev: Strong<dyn IKeyMintDevice> = self.keymint.get_interface()?;
         map_km_error(km_dev.addRngEntropy(entropy))
             .context("In generate_key: Trying to add entropy.")?;
-        let creation_result = map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
+        let mut creation_result = map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
             .context("In generate_key: While generating Key")?;
-
+        // The certificate chain ultimately gets flattened into a big DER encoded byte array,
+        // so providing that blob upfront in a single certificate entry should be fine.
+        if let Some(cert) = cert_chain {
+            creation_result.certificateChain.push(cert);
+        }
         let user_id = uid_to_android_user(caller_uid);
         self.store_new_key(key, creation_result, user_id, Some(flags)).context("In generate_key.")
     }
 
+    fn get_attest_key_and_cert_chain(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        attest_key_descriptor: Option<&KeyDescriptor>,
+        params: &[KeyParameter],
+        db: &mut KeystoreDB,
+    ) -> Result<(Option<AttestationKey>, Option<Certificate>)> {
+        match attest_key_descriptor {
+            None => self
+                .rem_prov_state
+                .get_remote_provisioning_key_and_certs(&key, caller_uid, params, db),
+            Some(attest_key) => Ok((
+                Some(
+                    self.get_attest_key(&attest_key, caller_uid)
+                        .context("In generate_key: Trying to load attest key")?,
+                ),
+                None,
+            )),
+        }
+    }
+
     fn get_attest_key(&self, key: &KeyDescriptor, caller_uid: u32) -> Result<AttestationKey> {
         let (km_blob, cert) = self
             .load_attest_key_blob_and_cert(&key, caller_uid)
             .context("In get_attest_key: Failed to load blob and cert")?;
 
-        let issuer_subject: Vec<u8> = parse_issuer_subject_from_certificate(&cert)
+        let issuer_subject: Vec<u8> = parse_subject_from_certificate(&cert)
             .context("In get_attest_key: Failed to parse subject from certificate.")?;
 
         Ok(AttestationKey {
diff --git a/keystore2/src/service.rs b/keystore2/src/service.rs
index 3a4bf82..15b729d 100644
--- a/keystore2/src/service.rs
+++ b/keystore2/src/service.rs
@@ -43,7 +43,7 @@
     KeyDescriptor::KeyDescriptor, KeyEntryResponse::KeyEntryResponse, KeyMetadata::KeyMetadata,
 };
 use anyhow::{Context, Result};
-use binder::{IBinder, Strong, ThreadState};
+use binder::{IBinderInternal, Strong, ThreadState};
 use error::Error;
 use keystore2_selinux as selinux;
 
diff --git a/keystore2/src/user_manager.rs b/keystore2/src/user_manager.rs
index 3c393c5..123f3a1 100644
--- a/keystore2/src/user_manager.rs
+++ b/keystore2/src/user_manager.rs
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-//! This module implements IKeystoreUserManager AIDL interface.
+//! This module implements IKeystoreMaintenance AIDL interface.
 
 use crate::error::map_or_log_err;
 use crate::error::Error as KeystoreError;
@@ -20,23 +20,24 @@
 use crate::permission::KeystorePerm;
 use crate::super_key::UserState;
 use crate::utils::check_keystore_permission;
-use android_security_usermanager::aidl::android::security::usermanager::IKeystoreUserManager::{
-    BnKeystoreUserManager, IKeystoreUserManager,
+use android_security_maintenance::aidl::android::security::maintenance::{
+    IKeystoreMaintenance::{BnKeystoreMaintenance, IKeystoreMaintenance},
+    UserState::UserState as AidlUserState,
 };
-use android_security_usermanager::binder::{Interface, Result as BinderResult};
+use android_security_maintenance::binder::{Interface, Result as BinderResult};
 use android_system_keystore2::aidl::android::system::keystore2::Domain::Domain;
 use android_system_keystore2::aidl::android::system::keystore2::ResponseCode::ResponseCode;
 use anyhow::{Context, Result};
-use binder::{IBinder, Strong};
+use binder::{IBinderInternal, Strong};
 
 /// This struct is defined to implement the aforementioned AIDL interface.
 /// As of now, it is an empty struct.
-pub struct UserManager;
+pub struct Maintenance;
 
-impl UserManager {
+impl Maintenance {
     /// Create a new instance of Keystore User Manager service.
-    pub fn new_native_binder() -> Result<Strong<dyn IKeystoreUserManager>> {
-        let result = BnKeystoreUserManager::new_binder(Self);
+    pub fn new_native_binder() -> Result<Strong<dyn IKeystoreMaintenance>> {
+        let result = BnKeystoreMaintenance::new_binder(Self);
         result.as_binder().set_requesting_sid(true);
         Ok(result)
     }
@@ -97,11 +98,28 @@
         DB.with(|db| db.borrow_mut().unbind_keys_for_namespace(domain, nspace))
             .context("In clear_namespace: Trying to delete keys from db.")
     }
+
+    fn get_state(user_id: i32) -> Result<AidlUserState> {
+        // Check permission. Function should return if this failed. Therefore having '?' at the end
+        // is very important.
+        check_keystore_permission(KeystorePerm::get_state()).context("In get_state.")?;
+        let state = DB
+            .with(|db| {
+                UserState::get(&mut db.borrow_mut(), &LEGACY_MIGRATOR, &SUPER_KEY, user_id as u32)
+            })
+            .context("In get_state. Trying to get UserState.")?;
+
+        match state {
+            UserState::Uninitialized => Ok(AidlUserState::UNINITIALIZED),
+            UserState::LskfUnlocked(_) => Ok(AidlUserState::LSKF_UNLOCKED),
+            UserState::LskfLocked => Ok(AidlUserState::LSKF_LOCKED),
+        }
+    }
 }
 
-impl Interface for UserManager {}
+impl Interface for Maintenance {}
 
-impl IKeystoreUserManager for UserManager {
+impl IKeystoreMaintenance for Maintenance {
     fn onUserPasswordChanged(&self, user_id: i32, password: Option<&[u8]>) -> BinderResult<()> {
         map_or_log_err(Self::on_user_password_changed(user_id, password), Ok)
     }
@@ -117,4 +135,8 @@
     fn clearNamespace(&self, domain: Domain, nspace: i64) -> BinderResult<()> {
         map_or_log_err(Self::clear_namespace(domain, nspace), Ok)
     }
+
+    fn getState(&self, user_id: i32) -> binder::public_api::Result<AidlUserState> {
+        map_or_log_err(Self::get_state(user_id), Ok)
+    }
 }
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index 8e161b7..2748025 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -15,12 +15,13 @@
 //! This module implements utility functions used by the Keystore 2.0 service
 //! implementation.
 
-use crate::error::Error;
+use crate::error::{map_binder_status, Error, ErrorCode};
 use crate::permission;
 use crate::permission::{KeyPerm, KeyPermSet, KeystorePerm};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    KeyCharacteristics::KeyCharacteristics,
+    KeyCharacteristics::KeyCharacteristics, Tag::Tag,
 };
+use android_os_permissions_aidl::aidl::android::os::IPermissionController;
 use android_security_apc::aidl::android::security::apc::{
     IProtectedConfirmation::{FLAG_UI_OPTION_INVERTED, FLAG_UI_OPTION_MAGNIFIED},
     ResponseCode::ResponseCode as ApcResponseCode,
@@ -88,6 +89,34 @@
     })
 }
 
+/// This function checks whether a given tag corresponds to the access of device identifiers.
+pub fn is_device_id_attestation_tag(tag: Tag) -> bool {
+    matches!(tag, Tag::ATTESTATION_ID_IMEI | Tag::ATTESTATION_ID_MEID | Tag::ATTESTATION_ID_SERIAL)
+}
+
+/// This function checks whether the calling app has the Android permissions needed to attest device
+/// identifiers. It throws an error if the permissions cannot be verified, or if the caller doesn't
+/// have the right permissions, and returns silently otherwise.
+pub fn check_device_attestation_permissions() -> anyhow::Result<()> {
+    let permission_controller: binder::Strong<dyn IPermissionController::IPermissionController> =
+        binder::get_interface("permission")?;
+
+    let binder_result = permission_controller.checkPermission(
+        "android.permission.READ_PRIVILEGED_PHONE_STATE",
+        ThreadState::get_calling_pid(),
+        ThreadState::get_calling_uid() as i32,
+    );
+    let has_permissions = map_binder_status(binder_result)
+        .context("In check_device_attestation_permissions: checkPermission failed")?;
+    match has_permissions {
+        true => Ok(()),
+        false => Err(Error::Km(ErrorCode::CANNOT_ATTEST_IDS)).context(concat!(
+            "In check_device_attestation_permissions: ",
+            "caller does not have the permission to attest device IDs"
+        )),
+    }
+}
+
 /// Thread safe wrapper around SpIBinder. It is safe to have SpIBinder smart pointers to the
 /// same object in multiple threads, but cloning a SpIBinder is not thread safe.
 /// Keystore frequently hands out binder tokens to the security level interface. If this
@@ -193,3 +222,21 @@
 pub fn uid_to_android_user(uid: u32) -> u32 {
     uid / AID_USER_OFFSET
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use anyhow::Result;
+
+    #[test]
+    fn check_device_attestation_permissions_test() -> Result<()> {
+        check_device_attestation_permissions().or_else(|error| {
+            match error.root_cause().downcast_ref::<Error>() {
+                // Expected: the context for this test might not be allowed to attest device IDs.
+                Some(Error::Km(ErrorCode::CANNOT_ATTEST_IDS)) => Ok(()),
+                // Other errors are unexpected
+                _ => Err(error),
+            }
+        })
+    }
+}
diff --git a/ondevice-signing/Android.bp b/ondevice-signing/Android.bp
index 5db19b7..1c3706d 100644
--- a/ondevice-signing/Android.bp
+++ b/ondevice-signing/Android.bp
@@ -26,7 +26,6 @@
 tidy_errors = [
   "cert-err34-c",
   "google-default-arguments",
-  "google-explicit-constructor",
   "google-runtime-int",
   "google-runtime-member-string-references",
   "misc-move-const-arg",
@@ -87,17 +86,22 @@
     "CertUtils.cpp",
     "Keymaster.cpp",
     "KeymasterSigningKey.cpp",
+    "KeystoreKey.cpp",
     "VerityUtils.cpp",
   ],
 
   static_libs: [
     "libmini_keyctl_static", // TODO need static?
     "libc++fs",
+    "lib_odsign_proto",
   ],
 
   shared_libs: [
     "android.hardware.keymaster@4.1",
+    "android.system.keystore2-V1-cpp",
+    "android.hardware.security.keymint-V1-cpp",
     "libbase",
+    "libbinder",
     "libcrypto",
     "libcrypto_utils",
     "libfsverity",
@@ -106,6 +110,7 @@
     "libkeymaster4support", // For authorization_set
     "libkeymaster4_1support",
     "libkeyutils",
+    "libprotobuf-cpp-full",
     "libutils",
   ],
 }
diff --git a/ondevice-signing/CertUtils.cpp b/ondevice-signing/CertUtils.cpp
index 6b24391..b0b75a6 100644
--- a/ondevice-signing/CertUtils.cpp
+++ b/ondevice-signing/CertUtils.cpp
@@ -25,6 +25,9 @@
 
 #include <fcntl.h>
 #include <vector>
+
+#include "KeyConstants.h"
+
 const char kBasicConstraints[] = "CA:TRUE";
 const char kKeyUsage[] = "critical,keyCertSign,cRLSign,digitalSignature";
 const char kSubjectKeyIdentifier[] = "hash";
@@ -52,6 +55,33 @@
     return true;
 }
 
+Result<bssl::UniquePtr<RSA>> getRsa(const std::vector<uint8_t>& publicKey) {
+    bssl::UniquePtr<RSA> rsaPubkey(RSA_new());
+    rsaPubkey->n = BN_new();
+    rsaPubkey->e = BN_new();
+
+    BN_bin2bn(publicKey.data(), publicKey.size(), rsaPubkey->n);
+    BN_set_word(rsaPubkey->e, kRsaKeyExponent);
+
+    return rsaPubkey;
+}
+
+Result<void> verifySignature(const std::string& message, const std::string& signature,
+                             const std::vector<uint8_t>& publicKey) {
+    auto rsaKey = getRsa(publicKey);
+    uint8_t hashBuf[SHA256_DIGEST_LENGTH];
+    SHA256(const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(message.c_str())),
+           message.length(), hashBuf);
+
+    bool success = RSA_verify(NID_sha256, hashBuf, sizeof(hashBuf),
+                              (const uint8_t*)signature.c_str(), signature.length(), rsaKey->get());
+
+    if (!success) {
+        return Error() << "Failed to verify signature.";
+    }
+    return {};
+}
+
 Result<void> createSelfSignedCertificate(
     const std::vector<uint8_t>& publicKey,
     const std::function<Result<std::string>(const std::string&)>& signFunction,
@@ -66,8 +96,13 @@
     X509_gmtime_adj(X509_get_notBefore(x509.get()), 0);
     X509_gmtime_adj(X509_get_notAfter(x509.get()), kCertLifetimeSeconds);
 
-    auto pubKeyData = publicKey.data();
-    EVP_PKEY* public_key = d2i_PUBKEY(nullptr, &pubKeyData, publicKey.size());
+    // "publicKey" corresponds to the raw public key bytes - need to create
+    // a new RSA key with the correct exponent.
+    auto rsaPubkey = getRsa(publicKey);
+
+    EVP_PKEY* public_key = EVP_PKEY_new();
+    EVP_PKEY_assign_RSA(public_key, rsaPubkey->release());
+
     if (!X509_set_pubkey(x509.get(), public_key)) {
         return Error() << "Unable to set x509 public key";
     }
@@ -112,11 +147,14 @@
 
     x509->signature->flags &= ~(ASN1_STRING_FLAG_BITS_LEFT | 0x07);
     x509->signature->flags |= ASN1_STRING_FLAG_BITS_LEFT;
-    auto f = fopen(path.c_str(), "wb");
-    // TODO error checking
+    auto f = fopen(path.c_str(), "wbe");
+    if (f == nullptr) {
+        return Error() << "Failed to open " << path;
+    }
     i2d_X509_fp(f, x509.get());
     fclose(f);
 
+    EVP_PKEY_free(public_key);
     return {};
 }
 
@@ -142,17 +180,33 @@
     return pubKey;
 }
 
-Result<std::vector<uint8_t>> extractPublicKeyFromX509(const std::vector<uint8_t>& keyData) {
+Result<std::vector<uint8_t>>
+extractPublicKeyFromSubjectPublicKeyInfo(const std::vector<uint8_t>& keyData) {
     auto keyDataBytes = keyData.data();
     EVP_PKEY* public_key = d2i_PUBKEY(nullptr, &keyDataBytes, keyData.size());
 
     return extractPublicKey(public_key);
 }
 
+Result<std::vector<uint8_t>> extractPublicKeyFromX509(const std::vector<uint8_t>& keyData) {
+    auto keyDataBytes = keyData.data();
+    bssl::UniquePtr<X509> decoded_cert(d2i_X509(nullptr, &keyDataBytes, keyData.size()));
+    if (decoded_cert.get() == nullptr) {
+        return Error() << "Failed to decode X509 certificate.";
+    }
+    bssl::UniquePtr<EVP_PKEY> decoded_pkey(X509_get_pubkey(decoded_cert.get()));
+
+    return extractPublicKey(decoded_pkey.get());
+}
+
 Result<std::vector<uint8_t>> extractPublicKeyFromX509(const std::string& path) {
     X509* cert;
-    auto f = fopen(path.c_str(), "r");
+    auto f = fopen(path.c_str(), "re");
+    if (f == nullptr) {
+        return Error() << "Failed to open " << path;
+    }
     if (!d2i_X509_fp(f, &cert)) {
+        fclose(f);
         return Error() << "Unable to decode x509 cert at " << path;
     }
 
diff --git a/ondevice-signing/CertUtils.h b/ondevice-signing/CertUtils.h
index d9172d0..66dff04 100644
--- a/ondevice-signing/CertUtils.h
+++ b/ondevice-signing/CertUtils.h
@@ -25,5 +25,11 @@
 android::base::Result<std::vector<uint8_t>> createPkcs7(const std::vector<uint8_t>& signedData);
 
 android::base::Result<std::vector<uint8_t>>
-extractPublicKeyFromX509(const std::vector<uint8_t>& path);
+extractPublicKeyFromX509(const std::vector<uint8_t>& x509);
+android::base::Result<std::vector<uint8_t>>
+extractPublicKeyFromSubjectPublicKeyInfo(const std::vector<uint8_t>& subjectKeyInfo);
 android::base::Result<std::vector<uint8_t>> extractPublicKeyFromX509(const std::string& path);
+
+android::base::Result<void> verifySignature(const std::string& message,
+                                            const std::string& signature,
+                                            const std::vector<uint8_t>& publicKey);
diff --git a/ondevice-signing/KeyConstants.h b/ondevice-signing/KeyConstants.h
new file mode 100644
index 0000000..9e1a513
--- /dev/null
+++ b/ondevice-signing/KeyConstants.h
@@ -0,0 +1,18 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+static constexpr int kRsaKeySize = 2048;
+static constexpr int kRsaKeyExponent = 65537;
diff --git a/ondevice-signing/KeymasterSigningKey.cpp b/ondevice-signing/KeymasterSigningKey.cpp
index 2b748e4..dc3ef8a 100644
--- a/ondevice-signing/KeymasterSigningKey.cpp
+++ b/ondevice-signing/KeymasterSigningKey.cpp
@@ -33,30 +33,36 @@
 using android::base::Result;
 using android::base::unique_fd;
 
+const std::string kSigningKeyBlob = "/data/misc/odsign/key.blob";
+
 KeymasterSigningKey::KeymasterSigningKey() {}
 
-Result<KeymasterSigningKey> KeymasterSigningKey::loadFromBlobAndVerify(const std::string& path) {
-    KeymasterSigningKey signingKey;
+Result<std::unique_ptr<KeymasterSigningKey>>
+KeymasterSigningKey::loadFromBlobAndVerify(const std::string& path) {
+    auto signingKey = std::make_unique<KeymasterSigningKey>();
 
-    auto status = signingKey.initializeFromKeyblob(path);
+    auto status = signingKey->initializeFromKeyblob(path);
 
     if (!status.ok()) {
         return status.error();
     }
 
-    return std::move(signingKey);
+    return signingKey;
 }
 
-Result<KeymasterSigningKey> KeymasterSigningKey::createNewKey() {
-    KeymasterSigningKey signingKey;
+Result<void> KeymasterSigningKey::saveKeyblob(const std::string& path) const {
+    int flags = O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC;
 
-    auto status = signingKey.createSigningKey();
-
-    if (!status.ok()) {
-        return status.error();
+    unique_fd fd(TEMP_FAILURE_RETRY(open(path.c_str(), flags, 0600)));
+    if (fd == -1) {
+        return ErrnoError() << "Error creating key blob file " << path;
     }
 
-    return std::move(signingKey);
+    if (!android::base::WriteFully(fd, mVerifiedKeyBlob.data(), mVerifiedKeyBlob.size())) {
+        return ErrnoError() << "Error writing key blob file " << path;
+    } else {
+        return {};
+    }
 }
 
 Result<void> KeymasterSigningKey::createSigningKey() {
@@ -78,41 +84,45 @@
     return {};
 }
 
-Result<void> KeymasterSigningKey::saveKeyblob(const std::string& path) const {
-    int flags = O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC;
+Result<std::unique_ptr<KeymasterSigningKey>> KeymasterSigningKey::createAndPersistNewKey() {
+    auto signingKey = std::make_unique<KeymasterSigningKey>();
 
-    unique_fd fd(TEMP_FAILURE_RETRY(open(path.c_str(), flags, 0600)));
-    if (fd == -1) {
-        return ErrnoError() << "Error creating key blob file " << path;
+    auto status = signingKey->createSigningKey();
+
+    if (!status.ok()) {
+        return status.error();
     }
 
-    if (!android::base::WriteFully(fd, mVerifiedKeyBlob.data(), mVerifiedKeyBlob.size())) {
-        return ErrnoError() << "Error writing key blob file " << path;
-    } else {
-        return {};
+    status = signingKey->saveKeyblob(kSigningKeyBlob);
+    if (!status.ok()) {
+        return status.error();
     }
+
+    return signingKey;
+}
+
+Result<SigningKey*> KeymasterSigningKey::getInstance() {
+    auto key = loadFromBlobAndVerify(kSigningKeyBlob);
+
+    if (!key.ok()) {
+        key = createAndPersistNewKey();
+        if (!key.ok()) {
+            return key.error();
+        }
+    }
+
+    return key->release();
 }
 
 Result<std::vector<uint8_t>> KeymasterSigningKey::getPublicKey() const {
-    auto publicKeyX509 = mKeymaster->extractPublicKey(mVerifiedKeyBlob);
-    if (!publicKeyX509.ok()) {
-        return publicKeyX509.error();
-    }
-    return extractPublicKeyFromX509(publicKeyX509.value());
-}
-
-Result<void> KeymasterSigningKey::createX509Cert(const std::string& outPath) const {
     auto publicKey = mKeymaster->extractPublicKey(mVerifiedKeyBlob);
-
     if (!publicKey.ok()) {
         return publicKey.error();
     }
 
-    auto keymasterSignFunction = [&](const std::string& to_be_signed) {
-        return this->sign(to_be_signed);
-    };
-    createSelfSignedCertificate(*publicKey, keymasterSignFunction, outPath);
-    return {};
+    // Keymaster returns the public key not in a full X509 cert, but just the
+    // "SubjectPublicKeyInfo"
+    return extractPublicKeyFromSubjectPublicKeyInfo(publicKey.value());
 }
 
 Result<void> KeymasterSigningKey::initializeFromKeyblob(const std::string& path) {
diff --git a/ondevice-signing/KeymasterSigningKey.h b/ondevice-signing/KeymasterSigningKey.h
index 7631059..e66781f 100644
--- a/ondevice-signing/KeymasterSigningKey.h
+++ b/ondevice-signing/KeymasterSigningKey.h
@@ -23,30 +23,36 @@
 #include <utils/StrongPointer.h>
 
 #include "Keymaster.h"
+#include "SigningKey.h"
 
-class KeymasterSigningKey {
+class KeymasterSigningKey : public SigningKey {
     using KmDevice = ::android::hardware::keymaster::V4_1::IKeymasterDevice;
 
   public:
+    friend std::unique_ptr<KeymasterSigningKey> std::make_unique<KeymasterSigningKey>();
+    virtual ~KeymasterSigningKey(){};
+
     // Allow the key to be moved around
     KeymasterSigningKey& operator=(KeymasterSigningKey&& other) = default;
     KeymasterSigningKey(KeymasterSigningKey&& other) = default;
 
-    static android::base::Result<KeymasterSigningKey>
-    loadFromBlobAndVerify(const std::string& path);
-    static android::base::Result<KeymasterSigningKey> createNewKey();
+    static android::base::Result<SigningKey*> getInstance();
 
-    /* Sign a message with an initialized signing key */
-    android::base::Result<std::string> sign(const std::string& message) const;
-    android::base::Result<void> saveKeyblob(const std::string& path) const;
-    android::base::Result<std::vector<uint8_t>> getPublicKey() const;
-    android::base::Result<void> createX509Cert(const std::string& path) const;
+    virtual android::base::Result<std::string> sign(const std::string& message) const;
+    virtual android::base::Result<std::vector<uint8_t>> getPublicKey() const;
 
   private:
     KeymasterSigningKey();
 
+    static android::base::Result<std::unique_ptr<KeymasterSigningKey>> createAndPersistNewKey();
+    static android::base::Result<std::unique_ptr<KeymasterSigningKey>>
+    loadFromBlobAndVerify(const std::string& path);
+
     android::base::Result<void> createSigningKey();
     android::base::Result<void> initializeFromKeyblob(const std::string& path);
+    android::base::Result<void> saveKeyblob(const std::string& path) const;
+
+    static android::base::Result<KeymasterSigningKey> createNewKey();
 
     std::optional<Keymaster> mKeymaster;
     std::vector<uint8_t> mVerifiedKeyBlob;
diff --git a/ondevice-signing/KeystoreKey.cpp b/ondevice-signing/KeystoreKey.cpp
new file mode 100644
index 0000000..de7033f
--- /dev/null
+++ b/ondevice-signing/KeystoreKey.cpp
@@ -0,0 +1,243 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <string>
+
+#include <android-base/file.h>
+#include <android-base/logging.h>
+#include <binder/IServiceManager.h>
+
+#include <fcntl.h>
+#include <sys/stat.h>
+#include <sys/types.h>
+
+#include "CertUtils.h"
+#include "KeyConstants.h"
+#include "KeystoreKey.h"
+
+using android::defaultServiceManager;
+using android::IServiceManager;
+using android::sp;
+using android::String16;
+
+using android::hardware::security::keymint::Algorithm;
+using android::hardware::security::keymint::Digest;
+using android::hardware::security::keymint::KeyParameter;
+using android::hardware::security::keymint::KeyParameterValue;
+using android::hardware::security::keymint::KeyPurpose;
+using android::hardware::security::keymint::PaddingMode;
+using android::hardware::security::keymint::SecurityLevel;
+using android::hardware::security::keymint::Tag;
+
+using android::system::keystore2::CreateOperationResponse;
+using android::system::keystore2::Domain;
+using android::system::keystore2::KeyDescriptor;
+using android::system::keystore2::KeyEntryResponse;
+using android::system::keystore2::KeyMetadata;
+
+using android::base::Error;
+using android::base::Result;
+
+using android::base::unique_fd;
+
+// Keystore boot level that the odsign key uses
+static const int kOdsignBootLevel = 30;
+
+static KeyDescriptor getKeyDescriptor() {
+    // AIDL parcelable objects don't have constructor
+    static KeyDescriptor descriptor;
+    static std::once_flag flag;
+    std::call_once(flag, [&]() {
+        descriptor.domain = Domain::SELINUX;
+        descriptor.alias = String16("ondevice-signing");
+        descriptor.nspace = 101;  // odsign_key
+    });
+
+    return descriptor;
+}
+
+KeystoreKey::KeystoreKey() {}
+
+Result<KeyMetadata> KeystoreKey::createNewKey(const KeyDescriptor& descriptor) {
+    std::vector<KeyParameter> params;
+
+    KeyParameter algo;
+    algo.tag = Tag::ALGORITHM;
+    algo.value = KeyParameterValue::make<KeyParameterValue::algorithm>(Algorithm::RSA);
+    params.push_back(algo);
+
+    KeyParameter key_size;
+    key_size.tag = Tag::KEY_SIZE;
+    key_size.value = KeyParameterValue::make<KeyParameterValue::integer>(kRsaKeySize);
+    params.push_back(key_size);
+
+    KeyParameter digest;
+    digest.tag = Tag::DIGEST;
+    digest.value = KeyParameterValue::make<KeyParameterValue::digest>(Digest::SHA_2_256);
+    params.push_back(digest);
+
+    KeyParameter padding;
+    padding.tag = Tag::PADDING;
+    padding.value =
+        KeyParameterValue::make<KeyParameterValue::paddingMode>(PaddingMode::RSA_PKCS1_1_5_SIGN);
+    params.push_back(padding);
+
+    KeyParameter exponent;
+    exponent.tag = Tag::RSA_PUBLIC_EXPONENT;
+    exponent.value = KeyParameterValue::make<KeyParameterValue::longInteger>(kRsaKeyExponent);
+    params.push_back(exponent);
+
+    KeyParameter purpose;
+    purpose.tag = Tag::PURPOSE;
+    purpose.value = KeyParameterValue::make<KeyParameterValue::keyPurpose>(KeyPurpose::SIGN);
+    params.push_back(purpose);
+
+    KeyParameter auth;
+    auth.tag = Tag::NO_AUTH_REQUIRED;
+    auth.value = KeyParameterValue::make<KeyParameterValue::boolValue>(true);
+    params.push_back(auth);
+
+    KeyParameter boot_level;
+    boot_level.tag = Tag::MAX_BOOT_LEVEL;
+    boot_level.value = KeyParameterValue::make<KeyParameterValue::integer>(kOdsignBootLevel);
+    params.push_back(boot_level);
+
+    KeyMetadata metadata;
+    auto status = mSecurityLevel->generateKey(descriptor, {}, params, 0, {}, &metadata);
+    if (!status.isOk()) {
+        return Error() << "Failed to create new key";
+    }
+
+    return metadata;
+}
+
+bool KeystoreKey::initialize() {
+    sp<IServiceManager> sm = defaultServiceManager();
+    if (sm == nullptr) {
+        return false;
+    }
+    auto service = sm->getService(String16("android.system.keystore2"));
+    if (service == nullptr) {
+        return false;
+    }
+    mService = interface_cast<android::system::keystore2::IKeystoreService>(service);
+    if (mService == nullptr) {
+        return false;
+    }
+
+    auto status = mService->getSecurityLevel(SecurityLevel::STRONGBOX, &mSecurityLevel);
+    if (!status.isOk()) {
+        status = mService->getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT, &mSecurityLevel);
+        if (!status.isOk()) {
+            return false;
+        }
+    }
+
+    auto descriptor = getKeyDescriptor();
+    // See if we can fetch an existing key
+    KeyEntryResponse keyEntryResponse;
+    LOG(INFO) << "Trying to retrieve existing keystore key...";
+    status = mService->getKeyEntry(descriptor, &keyEntryResponse);
+    if (!status.isOk()) {
+        LOG(INFO) << "Existing keystore key not found, creating new key";
+        auto newKeyStatus = createNewKey(descriptor);
+        if (!newKeyStatus.ok()) {
+            LOG(ERROR) << "Failed to create new key";
+            return false;
+        }
+        mKeyMetadata = *newKeyStatus;
+    } else {
+        mKeyMetadata = keyEntryResponse.metadata;
+    }
+
+    LOG(ERROR) << "Initialized Keystore key.";
+    return true;
+}
+
+Result<SigningKey*> KeystoreKey::getInstance() {
+    static KeystoreKey keystoreKey;
+
+    if (!keystoreKey.initialize()) {
+        return Error() << "Failed to initialize keystore key.";
+    } else {
+        return &keystoreKey;
+    }
+}
+
+static std::vector<KeyParameter> getSignOpParameters() {
+    std::vector<KeyParameter> opParameters;
+
+    KeyParameter algo;
+    algo.tag = Tag::ALGORITHM;
+    algo.value = KeyParameterValue::make<KeyParameterValue::algorithm>(Algorithm::RSA);
+    opParameters.push_back(algo);
+
+    KeyParameter digest;
+    digest.tag = Tag::DIGEST;
+    digest.value = KeyParameterValue::make<KeyParameterValue::digest>(Digest::SHA_2_256);
+    opParameters.push_back(digest);
+
+    KeyParameter padding;
+    padding.tag = Tag::PADDING;
+    padding.value =
+        KeyParameterValue::make<KeyParameterValue::paddingMode>(PaddingMode::RSA_PKCS1_1_5_SIGN);
+    opParameters.push_back(padding);
+
+    KeyParameter purpose;
+    purpose.tag = Tag::PURPOSE;
+    purpose.value = KeyParameterValue::make<KeyParameterValue::keyPurpose>(KeyPurpose::SIGN);
+    opParameters.push_back(purpose);
+
+    return opParameters;
+}
+
+Result<std::string> KeystoreKey::sign(const std::string& message) const {
+    static auto opParameters = getSignOpParameters();
+
+    CreateOperationResponse opResponse;
+
+    auto status =
+        mSecurityLevel->createOperation(getKeyDescriptor(), opParameters, false, &opResponse);
+    if (!status.isOk()) {
+        return Error() << "Failed to create keystore signing operation: "
+                       << status.serviceSpecificErrorCode();
+    }
+    auto operation = opResponse.iOperation;
+
+    std::optional<std::vector<uint8_t>> out;
+    status = operation->update({message.begin(), message.end()}, &out);
+    if (!status.isOk()) {
+        return Error() << "Failed to call keystore update operation.";
+    }
+
+    std::optional<std::vector<uint8_t>> signature;
+    status = operation->finish({}, {}, &signature);
+    if (!status.isOk()) {
+        return Error() << "Failed to call keystore finish operation.";
+    }
+
+    if (!signature.has_value()) {
+        return Error() << "Didn't receive a signature from keystore finish operation.";
+    }
+
+    std::string result{signature.value().begin(), signature.value().end()};
+
+    return result;
+}
+
+Result<std::vector<uint8_t>> KeystoreKey::getPublicKey() const {
+    return extractPublicKeyFromX509(mKeyMetadata.certificate.value());
+}
diff --git a/ondevice-signing/KeystoreKey.h b/ondevice-signing/KeystoreKey.h
new file mode 100644
index 0000000..6b9cb57
--- /dev/null
+++ b/ondevice-signing/KeystoreKey.h
@@ -0,0 +1,52 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <optional>
+
+#include <android-base/macros.h>
+#include <android-base/result.h>
+#include <android-base/unique_fd.h>
+
+#include <utils/StrongPointer.h>
+
+#include <android/system/keystore2/IKeystoreService.h>
+
+#include "SigningKey.h"
+
+class KeystoreKey : public SigningKey {
+    using IKeystoreService = ::android::system::keystore2::IKeystoreService;
+    using IKeystoreSecurityLevel = ::android::system::keystore2::IKeystoreSecurityLevel;
+    using KeyDescriptor = ::android::system::keystore2::KeyDescriptor;
+    using KeyMetadata = ::android::system::keystore2::KeyMetadata;
+
+  public:
+    virtual ~KeystoreKey(){};
+    static android::base::Result<SigningKey*> getInstance();
+
+    virtual android::base::Result<std::string> sign(const std::string& message) const;
+    virtual android::base::Result<std::vector<uint8_t>> getPublicKey() const;
+
+  private:
+    KeystoreKey();
+    bool initialize();
+    android::base::Result<KeyMetadata> createNewKey(const KeyDescriptor& descriptor);
+
+    android::sp<IKeystoreService> mService;
+    android::sp<IKeystoreSecurityLevel> mSecurityLevel;
+    KeyMetadata mKeyMetadata;
+};
diff --git a/ondevice-signing/SigningKey.h b/ondevice-signing/SigningKey.h
new file mode 100644
index 0000000..89294fc
--- /dev/null
+++ b/ondevice-signing/SigningKey.h
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <android-base/macros.h>
+#include <android-base/result.h>
+
+class SigningKey {
+  public:
+    virtual ~SigningKey(){};
+    /* Sign a message with an initialized signing key */
+    virtual android::base::Result<std::string> sign(const std::string& message) const = 0;
+    /* Retrieve the associated public key */
+    virtual android::base::Result<std::vector<uint8_t>> getPublicKey() const = 0;
+};
diff --git a/ondevice-signing/VerityUtils.cpp b/ondevice-signing/VerityUtils.cpp
index b4a6a54..cab92e2 100644
--- a/ondevice-signing/VerityUtils.cpp
+++ b/ondevice-signing/VerityUtils.cpp
@@ -15,12 +15,15 @@
  */
 
 #include <filesystem>
+#include <map>
+#include <span>
 #include <string>
 
 #include <fcntl.h>
 #include <linux/fs.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <sys/wait.h>
 
 #include <android-base/logging.h>
 #include <android-base/unique_fd.h>
@@ -28,13 +31,17 @@
 #include <linux/fsverity.h>
 
 #include "CertUtils.h"
-#include "KeymasterSigningKey.h"
+#include "SigningKey.h"
+
+#define FS_VERITY_MAX_DIGEST_SIZE 64
 
 using android::base::ErrnoError;
 using android::base::Error;
 using android::base::Result;
 using android::base::unique_fd;
 
+static const char* kFsVerityInitPath = "/system/bin/fsverity_init";
+
 #if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
 #define cpu_to_le16(v) ((__force __le16)(uint16_t)(v))
 #define le16_to_cpu(v) ((__force uint16_t)(__le16)(v))
@@ -50,17 +57,31 @@
     __u8 digest[];
 };
 
+static std::string toHex(std::span<uint8_t> data) {
+    std::stringstream ss;
+    for (auto it = data.begin(); it != data.end(); ++it) {
+        ss << std::setfill('0') << std::setw(2) << std::hex << static_cast<unsigned>(*it);
+    }
+    return ss.str();
+}
+
 static int read_callback(void* file, void* buf, size_t count) {
     int* fd = (int*)file;
     if (TEMP_FAILURE_RETRY(read(*fd, buf, count)) < 0) return errno ? -errno : -EIO;
     return 0;
 }
 
-static Result<std::vector<uint8_t>> createDigest(const std::string& path) {
+Result<std::vector<uint8_t>> createDigest(const std::string& path) {
     struct stat filestat;
     unique_fd fd(TEMP_FAILURE_RETRY(open(path.c_str(), O_RDONLY | O_CLOEXEC)));
+    if (fd < 0) {
+        return ErrnoError() << "Failed to open " << path;
+    }
 
-    stat(path.c_str(), &filestat);
+    int ret = stat(path.c_str(), &filestat);
+    if (ret < 0) {
+        return ErrnoError() << "Failed to stat " << path;
+    }
     struct libfsverity_merkle_tree_params params = {
         .version = 1,
         .hash_algorithm = FS_VERITY_HASH_ALG_SHA256,
@@ -69,24 +90,45 @@
     };
 
     struct libfsverity_digest* digest;
-    libfsverity_compute_digest(&fd, &read_callback, &params, &digest);
-
-    return std::vector<uint8_t>(&digest->digest[0], &digest->digest[32]);
+    ret = libfsverity_compute_digest(&fd, &read_callback, &params, &digest);
+    if (ret < 0) {
+        return ErrnoError() << "Failed to compute fs-verity digest for " << path;
+    }
+    std::vector<uint8_t> digestVector(&digest->digest[0], &digest->digest[32]);
+    free(digest);
+    return digestVector;
 }
 
-static Result<std::vector<uint8_t>> signDigest(const KeymasterSigningKey& key,
+namespace {
+template <typename T> struct DeleteAsPODArray {
+    void operator()(T* x) {
+        if (x) {
+            x->~T();
+            delete[](uint8_t*) x;
+        }
+    }
+};
+}  // namespace
+
+template <typename T> using trailing_unique_ptr = std::unique_ptr<T, DeleteAsPODArray<T>>;
+
+template <typename T>
+static trailing_unique_ptr<T> makeUniqueWithTrailingData(size_t trailing_data_size) {
+    uint8_t* memory = new uint8_t[sizeof(T*) + trailing_data_size];
+    T* ptr = new (memory) T;
+    return trailing_unique_ptr<T>{ptr};
+}
+
+static Result<std::vector<uint8_t>> signDigest(const SigningKey& key,
                                                const std::vector<uint8_t>& digest) {
-    fsverity_signed_digest* d;
-    size_t signed_digest_size = sizeof(*d) + digest.size();
-    std::unique_ptr<uint8_t[]> digest_buffer{new uint8_t[signed_digest_size]};
-    d = (fsverity_signed_digest*)digest_buffer.get();
+    auto d = makeUniqueWithTrailingData<fsverity_signed_digest>(digest.size());
 
     memcpy(d->magic, "FSVerity", 8);
     d->digest_algorithm = cpu_to_le16(FS_VERITY_HASH_ALG_SHA256);
     d->digest_size = cpu_to_le16(digest.size());
     memcpy(d->digest, digest.data(), digest.size());
 
-    auto signed_digest = key.sign(std::string((char*)d, signed_digest_size));
+    auto signed_digest = key.sign(std::string((char*)d.get(), sizeof(*d) + digest.size()));
     if (!signed_digest.ok()) {
         return signed_digest.error();
     }
@@ -94,7 +136,7 @@
     return std::vector<uint8_t>(signed_digest->begin(), signed_digest->end());
 }
 
-Result<void> enableFsVerity(const std::string& path, const KeymasterSigningKey& key) {
+Result<std::string> enableFsVerity(const std::string& path, const SigningKey& key) {
     auto digest = createDigest(path);
     if (!digest.ok()) {
         return digest.error();
@@ -121,10 +163,13 @@
         return ErrnoError() << "Failed to call FS_IOC_ENABLE_VERITY on " << path;
     }
 
-    return {};
+    // Return the root hash as a hex string
+    return toHex(digest.value());
 }
 
-Result<void> addFilesToVerityRecursive(const std::string& path, const KeymasterSigningKey& key) {
+Result<std::map<std::string, std::string>> addFilesToVerityRecursive(const std::string& path,
+                                                                     const SigningKey& key) {
+    std::map<std::string, std::string> digests;
     std::error_code ec;
 
     auto it = std::filesystem::recursive_directory_iterator(path, ec);
@@ -137,14 +182,18 @@
             if (!result.ok()) {
                 return result.error();
             }
+            digests[it->path()] = *result;
         }
         ++it;
     }
+    if (ec) {
+        return Error() << "Failed to iterate " << path << ": " << ec;
+    }
 
-    return {};
+    return digests;
 }
 
-Result<bool> isFileInVerity(const std::string& path) {
+Result<std::string> isFileInVerity(const std::string& path) {
     unsigned int flags;
 
     unique_fd fd(TEMP_FAILURE_RETRY(open(path.c_str(), O_RDONLY | O_CLOEXEC)));
@@ -156,11 +205,21 @@
     if (ret < 0) {
         return ErrnoError() << "Failed to FS_IOC_GETFLAGS for " << path;
     }
+    if (!(flags & FS_VERITY_FL)) {
+        return Error() << "File is not in fs-verity: " << path;
+    }
 
-    return (flags & FS_VERITY_FL);
+    auto d = makeUniqueWithTrailingData<fsverity_digest>(FS_VERITY_MAX_DIGEST_SIZE);
+    d->digest_size = FS_VERITY_MAX_DIGEST_SIZE;
+    ret = ioctl(fd, FS_IOC_MEASURE_VERITY, d.get());
+    if (ret < 0) {
+        return ErrnoError() << "Failed to FS_IOC_MEASURE_VERITY for " << path;
+    }
+    return toHex({&d->digest[0], &d->digest[d->digest_size]});
 }
 
-Result<void> verifyAllFilesInVerity(const std::string& path) {
+Result<std::map<std::string, std::string>> verifyAllFilesInVerity(const std::string& path) {
+    std::map<std::string, std::string> digests;
     std::error_code ec;
 
     auto it = std::filesystem::recursive_directory_iterator(path, ec);
@@ -173,12 +232,50 @@
             if (!result.ok()) {
                 return result.error();
             }
-            if (!*result) {
-                return Error() << "File " << it->path() << " not in fs-verity";
-            }
+            digests[it->path()] = *result;
         }  // TODO reject other types besides dirs?
         ++it;
     }
+    if (ec) {
+        return Error() << "Failed to iterate " << path << ": " << ec;
+    }
+
+    return digests;
+}
+
+Result<void> addCertToFsVerityKeyring(const std::string& path) {
+    const char* const argv[] = {kFsVerityInitPath, "--load-extra-key", "fsv_ods"};
+
+    int fd = open(path.c_str(), O_RDONLY | O_CLOEXEC);
+    pid_t pid = fork();
+    if (pid == 0) {
+        dup2(fd, STDIN_FILENO);
+        close(fd);
+        int argc = arraysize(argv);
+        char* argv_child[argc + 1];
+        memcpy(argv_child, argv, argc * sizeof(char*));
+        argv_child[argc] = nullptr;
+        execvp(argv_child[0], const_cast<char**>(argv_child));
+        PLOG(ERROR) << "exec in ForkExecvp";
+        _exit(EXIT_FAILURE);
+    } else {
+        close(fd);
+    }
+    if (pid == -1) {
+        return ErrnoError() << "Failed to fork.";
+    }
+    int status;
+    if (waitpid(pid, &status, 0) == -1) {
+        return ErrnoError() << "waitpid() failed.";
+    }
+    if (!WIFEXITED(status)) {
+        return Error() << kFsVerityInitPath << ": abnormal process exit";
+    }
+    if (WEXITSTATUS(status)) {
+        if (status != 0) {
+            return Error() << kFsVerityInitPath << " exited with " << status;
+        }
+    }
 
     return {};
 }
diff --git a/ondevice-signing/VerityUtils.h b/ondevice-signing/VerityUtils.h
index 1eca5a6..84af319 100644
--- a/ondevice-signing/VerityUtils.h
+++ b/ondevice-signing/VerityUtils.h
@@ -18,8 +18,11 @@
 
 #include <android-base/result.h>
 
-#include "KeymasterSigningKey.h"
+#include "SigningKey.h"
 
-android::base::Result<void> verifyAllFilesInVerity(const std::string& path);
-android::base::Result<void> addFilesToVerityRecursive(const std::string& path,
-                                                      const KeymasterSigningKey& key);
+android::base::Result<void> addCertToFsVerityKeyring(const std::string& path);
+android::base::Result<std::vector<uint8_t>> createDigest(const std::string& path);
+android::base::Result<std::map<std::string, std::string>>
+verifyAllFilesInVerity(const std::string& path);
+android::base::Result<std::map<std::string, std::string>>
+addFilesToVerityRecursive(const std::string& path, const SigningKey& key);
diff --git a/ondevice-signing/odsign_main.cpp b/ondevice-signing/odsign_main.cpp
index 3baba68..58e49ec 100644
--- a/ondevice-signing/odsign_main.cpp
+++ b/ondevice-signing/odsign_main.cpp
@@ -16,83 +16,56 @@
 
 #include <fcntl.h>
 #include <filesystem>
+#include <fstream>
 #include <iomanip>
 #include <iostream>
+#include <iterator>
 #include <sys/stat.h>
 #include <sys/types.h>
-#include <sys/wait.h>
 #include <unistd.h>
 
 #include <android-base/file.h>
 #include <android-base/logging.h>
+#include <android-base/properties.h>
 #include <android-base/scopeguard.h>
 #include <logwrap/logwrap.h>
 
 #include "CertUtils.h"
 #include "KeymasterSigningKey.h"
+#include "KeystoreKey.h"
 #include "VerityUtils.h"
 
+#include "odsign_info.pb.h"
+
 using android::base::ErrnoError;
 using android::base::Error;
 using android::base::Result;
+using android::base::SetProperty;
+
+using OdsignInfo = ::odsign::proto::OdsignInfo;
 
 const std::string kSigningKeyBlob = "/data/misc/odsign/key.blob";
 const std::string kSigningKeyCert = "/data/misc/odsign/key.cert";
+const std::string kOdsignInfo = "/data/misc/odsign/odsign.info";
+const std::string kOdsignInfoSignature = "/data/misc/odsign/odsign.info.signature";
 
 const std::string kArtArtifactsDir = "/data/misc/apexdata/com.android.art/dalvik-cache";
 
 static const char* kOdrefreshPath = "/apex/com.android.art/bin/odrefresh";
 
-static const char* kFsVerityInitPath = "/system/bin/fsverity_init";
+static const char* kFsVerityProcPath = "/proc/sys/fs/verity";
 
 static const bool kForceCompilation = false;
+static const bool kUseKeystore = false;
 
-Result<void> addCertToFsVerityKeyring(const std::string& path) {
-    const char* const argv[] = {kFsVerityInitPath, "--load-extra-key", "fsv_ods"};
+static const char* kOdsignVerificationDoneProp = "odsign.verification.done";
+static const char* kOdsignKeyDoneProp = "odsign.key.done";
 
-    // NOLINTNEXTLINE(android-cloexec-open): Deliberately not O_CLOEXEC
-    int fd = open(path.c_str(), O_RDONLY);
-    pid_t pid = fork();
-    if (pid == 0) {
-        dup2(fd, STDIN_FILENO);
-        close(fd);
-        int argc = arraysize(argv);
-        char* argv_child[argc + 1];
-        memcpy(argv_child, argv, argc * sizeof(char*));
-        argv_child[argc] = nullptr;
-        execvp(argv_child[0], const_cast<char**>(argv_child));
-        PLOG(ERROR) << "exec in ForkExecvp";
-        _exit(EXIT_FAILURE);
-    } else {
-        close(fd);
-    }
-    if (pid == -1) {
-        return ErrnoError() << "Failed to fork.";
-    }
-    int status;
-    if (waitpid(pid, &status, 0) == -1) {
-        return ErrnoError() << "waitpid() failed.";
-    }
-    if (!WIFEXITED(status)) {
-        return Error() << kFsVerityInitPath << ": abnormal process exit";
-    }
-    if (WEXITSTATUS(status)) {
-        if (status != 0) {
-            return Error() << kFsVerityInitPath << " exited with " << status;
-        }
-    }
+static const char* kOdsignVerificationStatusProp = "odsign.verification.success";
+static const char* kOdsignVerificationStatusValid = "1";
+static const char* kOdsignVerificationStatusError = "0";
 
-    return {};
-}
-
-Result<KeymasterSigningKey> loadAndVerifyExistingKey() {
-    if (access(kSigningKeyBlob.c_str(), F_OK) < 0) {
-        return ErrnoError() << "Key blob not found: " << kSigningKeyBlob;
-    }
-    return KeymasterSigningKey::loadFromBlobAndVerify(kSigningKeyBlob);
-}
-
-Result<void> verifyExistingCert(const KeymasterSigningKey& key) {
+Result<void> verifyExistingCert(const SigningKey& key) {
     if (access(kSigningKeyCert.c_str(), F_OK) < 0) {
         return ErrnoError() << "Key certificate not found: " << kSigningKeyCert;
     }
@@ -114,19 +87,18 @@
     return {};
 }
 
-Result<KeymasterSigningKey> createAndPersistKey(const std::string& path) {
-    auto key = KeymasterSigningKey::createNewKey();
+Result<void> createX509Cert(const SigningKey& key, const std::string& outPath) {
+    auto publicKey = key.getPublicKey();
 
-    if (!key.ok()) {
-        return key.error();
+    if (!publicKey.ok()) {
+        return publicKey.error();
     }
 
-    auto result = key->saveKeyblob(path);
-    if (!result.ok()) {
-        return result.error();
-    }
-
-    return key;
+    auto keymasterSignFunction = [&](const std::string& to_be_signed) {
+        return key.sign(to_be_signed);
+    };
+    createSelfSignedCertificate(*publicKey, keymasterSignFunction, outPath);
+    return {};
 }
 
 bool compileArtifacts(bool force) {
@@ -143,81 +115,291 @@
            0;
 }
 
-int main(int /* argc */, char** /* argv */) {
-    auto removeArtifacts = []() {
-        std::error_code ec;
-        auto num_removed = std::filesystem::remove_all(kArtArtifactsDir, ec);
-        if (ec) {
-            // TODO can't remove artifacts, signal Zygote shouldn't use them
-            LOG(ERROR) << "Can't remove " << kArtArtifactsDir << ": " << ec.message();
-        } else {
+static std::string toHex(const std::vector<uint8_t>& digest) {
+    std::stringstream ss;
+    for (auto it = digest.begin(); it != digest.end(); ++it) {
+        ss << std::setfill('0') << std::setw(2) << std::hex << static_cast<unsigned>(*it);
+    }
+    return ss.str();
+}
+
+Result<std::map<std::string, std::string>> computeDigests(const std::string& path) {
+    std::error_code ec;
+    std::map<std::string, std::string> digests;
+
+    auto it = std::filesystem::recursive_directory_iterator(path, ec);
+    auto end = std::filesystem::recursive_directory_iterator();
+
+    while (!ec && it != end) {
+        if (it->is_regular_file()) {
+            auto digest = createDigest(it->path());
+            if (!digest.ok()) {
+                return Error() << "Failed to compute digest for " << it->path();
+            }
+            digests[it->path()] = toHex(*digest);
+        }
+        ++it;
+    }
+    if (ec) {
+        return Error() << "Failed to iterate " << path << ": " << ec;
+    }
+
+    return digests;
+}
+
+Result<void> verifyDigests(const std::map<std::string, std::string>& digests,
+                           const std::map<std::string, std::string>& trusted_digests) {
+    for (const auto& path_digest : digests) {
+        auto path = path_digest.first;
+        auto digest = path_digest.second;
+        if ((trusted_digests.count(path) == 0)) {
+            return Error() << "Couldn't find digest for " << path;
+        }
+        if (trusted_digests.at(path) != digest) {
+            return Error() << "Digest mismatch for " << path;
+        }
+    }
+
+    // All digests matched!
+    if (digests.size() > 0) {
+        LOG(INFO) << "All root hashes match.";
+    }
+    return {};
+}
+
+Result<void> verifyIntegrityFsVerity(const std::map<std::string, std::string>& trusted_digests) {
+    // Just verify that the files are in verity, and get their digests
+    auto result = verifyAllFilesInVerity(kArtArtifactsDir);
+    if (!result.ok()) {
+        return result.error();
+    }
+
+    return verifyDigests(*result, trusted_digests);
+}
+
+Result<void> verifyIntegrityNoFsVerity(const std::map<std::string, std::string>& trusted_digests) {
+    // On these devices, just compute the digests, and verify they match the ones we trust
+    auto result = computeDigests(kArtArtifactsDir);
+    if (!result.ok()) {
+        return result.error();
+    }
+
+    return verifyDigests(*result, trusted_digests);
+}
+
+Result<OdsignInfo> getOdsignInfo(const SigningKey& key) {
+    std::string persistedSignature;
+    OdsignInfo odsignInfo;
+
+    if (!android::base::ReadFileToString(kOdsignInfoSignature, &persistedSignature)) {
+        return ErrnoError() << "Failed to read " << kOdsignInfoSignature;
+    }
+
+    std::fstream odsign_info(kOdsignInfo, std::ios::in | std::ios::binary);
+    if (!odsign_info) {
+        return Error() << "Failed to open " << kOdsignInfo;
+    }
+    odsign_info.seekg(0);
+    // Verify the hash
+    std::string odsign_info_str((std::istreambuf_iterator<char>(odsign_info)),
+                                std::istreambuf_iterator<char>());
+
+    auto publicKey = key.getPublicKey();
+    auto signResult = verifySignature(odsign_info_str, persistedSignature, *publicKey);
+    if (!signResult.ok()) {
+        return Error() << kOdsignInfoSignature << " does not match.";
+    } else {
+        LOG(INFO) << kOdsignInfoSignature << " matches.";
+    }
+
+    odsign_info.seekg(0);
+    if (!odsignInfo.ParseFromIstream(&odsign_info)) {
+        return Error() << "Failed to parse " << kOdsignInfo;
+    }
+
+    LOG(INFO) << "Loaded " << kOdsignInfo;
+    return odsignInfo;
+}
+
+Result<void> persistDigests(const std::map<std::string, std::string>& digests,
+                            const SigningKey& key) {
+    OdsignInfo signInfo;
+    google::protobuf::Map<std::string, std::string> proto_hashes(digests.begin(), digests.end());
+    auto map = signInfo.mutable_file_hashes();
+    *map = proto_hashes;
+
+    std::fstream odsign_info(kOdsignInfo,
+                             std::ios::in | std::ios::out | std::ios::trunc | std::ios::binary);
+    if (!signInfo.SerializeToOstream(&odsign_info)) {
+        return Error() << "Failed to persist root hashes in " << kOdsignInfo;
+    }
+
+    // Sign the signatures with our key itself, and write that to storage
+    odsign_info.seekg(0, std::ios::beg);
+    std::string odsign_info_str((std::istreambuf_iterator<char>(odsign_info)),
+                                std::istreambuf_iterator<char>());
+    auto signResult = key.sign(odsign_info_str);
+    if (!signResult.ok()) {
+        return Error() << "Failed to sign " << kOdsignInfo;
+    }
+    android::base::WriteStringToFile(*signResult, kOdsignInfoSignature);
+    return {};
+}
+
+static int removeArtifacts() {
+    std::error_code ec;
+    auto num_removed = std::filesystem::remove_all(kArtArtifactsDir, ec);
+    if (ec) {
+        LOG(ERROR) << "Can't remove " << kArtArtifactsDir << ": " << ec.message();
+        return 0;
+    } else {
+        if (num_removed > 0) {
             LOG(INFO) << "Removed " << num_removed << " entries from " << kArtArtifactsDir;
         }
-    };
-    // Make sure we delete the artifacts in all early (error) exit paths
-    auto scope_guard = android::base::make_scope_guard(removeArtifacts);
+        return num_removed;
+    }
+}
 
-    auto key = loadAndVerifyExistingKey();
-    if (!key.ok()) {
-        LOG(WARNING) << key.error().message();
+static Result<void> verifyArtifacts(const SigningKey& key, bool supportsFsVerity) {
+    auto signInfo = getOdsignInfo(key);
+    // Tell init we're done with the key; this is a boot time optimization
+    // in particular for the no fs-verity case, where we need to do a
+    // costly verification. If the files haven't been tampered with, which
+    // should be the common path, the verification will succeed, and we won't
+    // need the key anymore. If it turns out the artifacts are invalid (eg not
+    // in fs-verity) or the hash doesn't match, we won't be able to generate
+    // new artifacts without the key, so in those cases, remove the artifacts,
+    // and use JIT zygote for the current boot. We should recover automatically
+    // by the next boot.
+    SetProperty(kOdsignKeyDoneProp, "1");
+    if (!signInfo.ok()) {
+        return Error() << signInfo.error().message();
+    }
+    std::map<std::string, std::string> trusted_digests(signInfo->file_hashes().begin(),
+                                                       signInfo->file_hashes().end());
+    Result<void> integrityStatus;
 
-        key = createAndPersistKey(kSigningKeyBlob);
-        if (!key.ok()) {
-            LOG(ERROR) << "Failed to create or persist new key: " << key.error().message();
-            return -1;
-        }
+    if (supportsFsVerity) {
+        integrityStatus = verifyIntegrityFsVerity(trusted_digests);
     } else {
-        LOG(INFO) << "Found and verified existing key: " << kSigningKeyBlob;
+        integrityStatus = verifyIntegrityNoFsVerity(trusted_digests);
+    }
+    if (!integrityStatus.ok()) {
+        return Error() << integrityStatus.error().message();
     }
 
-    auto existing_cert = verifyExistingCert(key.value());
-    if (!existing_cert.ok()) {
-        LOG(WARNING) << existing_cert.error().message();
+    return {};
+}
 
-        // Try to create a new cert
-        auto new_cert = key->createX509Cert(kSigningKeyCert);
-        if (!new_cert.ok()) {
-            LOG(ERROR) << "Failed to create X509 certificate: " << new_cert.error().message();
-            // TODO apparently the key become invalid - delete the blob / cert
-            return -1;
-        }
-    } else {
-        LOG(INFO) << "Found and verified existing public key certificate: " << kSigningKeyCert;
-    }
-    auto cert_add_result = addCertToFsVerityKeyring(kSigningKeyCert);
-    if (!cert_add_result.ok()) {
-        LOG(ERROR) << "Failed to add certificate to fs-verity keyring: "
-                   << cert_add_result.error().message();
-        return -1;
-    }
-
-    auto verityStatus = verifyAllFilesInVerity(kArtArtifactsDir);
-    if (!verityStatus.ok()) {
-        LOG(WARNING) << verityStatus.error().message() << ", removing " << kArtArtifactsDir;
+int main(int /* argc */, char** /* argv */) {
+    auto errorScopeGuard = []() {
+        // In case we hit any error, remove the artifacts and tell Zygote not to use anything
         removeArtifacts();
+        // Tell init we don't need to use our key anymore
+        SetProperty(kOdsignKeyDoneProp, "1");
+        // Tell init we're done with verification, and that it was an error
+        SetProperty(kOdsignVerificationDoneProp, "1");
+        SetProperty(kOdsignVerificationStatusProp, kOdsignVerificationStatusError);
+    };
+    auto scope_guard = android::base::make_scope_guard(errorScopeGuard);
+
+    SigningKey* key;
+    if (kUseKeystore) {
+        auto keystoreResult = KeystoreKey::getInstance();
+        if (!keystoreResult.ok()) {
+            LOG(ERROR) << "Could not create keystore key: " << keystoreResult.error().message();
+            return -1;
+        }
+        key = keystoreResult.value();
+    } else {
+        // TODO - keymaster will go away
+        auto keymasterResult = KeymasterSigningKey::getInstance();
+        if (!keymasterResult.ok()) {
+            LOG(ERROR) << "Failed to create keymaster key: " << keymasterResult.error().message();
+            return -1;
+        }
+        key = keymasterResult.value();
     }
 
+    bool supportsFsVerity = access(kFsVerityProcPath, F_OK) == 0;
+    if (!supportsFsVerity) {
+        LOG(INFO) << "Device doesn't support fsverity. Falling back to full verification.";
+    }
+
+    if (supportsFsVerity) {
+        auto existing_cert = verifyExistingCert(*key);
+        if (!existing_cert.ok()) {
+            LOG(WARNING) << existing_cert.error().message();
+
+            // Try to create a new cert
+            auto new_cert = createX509Cert(*key, kSigningKeyCert);
+            if (!new_cert.ok()) {
+                LOG(ERROR) << "Failed to create X509 certificate: " << new_cert.error().message();
+                // TODO apparently the key become invalid - delete the blob / cert
+                return -1;
+            }
+        } else {
+            LOG(INFO) << "Found and verified existing public key certificate: " << kSigningKeyCert;
+        }
+        auto cert_add_result = addCertToFsVerityKeyring(kSigningKeyCert);
+        if (!cert_add_result.ok()) {
+            LOG(ERROR) << "Failed to add certificate to fs-verity keyring: "
+                       << cert_add_result.error().message();
+            return -1;
+        }
+    }
+
+    // Ask ART whether it considers the artifacts valid
+    LOG(INFO) << "Asking odrefresh to verify artifacts (if present)...";
     bool artifactsValid = validateArtifacts();
+    LOG(INFO) << "odrefresh said they are " << (artifactsValid ? "VALID" : "INVALID");
+
+    // A post-condition of validating artifacts is that if the ones on /system
+    // are used, kArtArtifactsDir is removed. Conversely, if kArtArtifactsDir
+    // exists, those are artifacts that will be used, and we should verify them.
+    int err = access(kArtArtifactsDir.c_str(), F_OK);
+    // If we receive any error other than ENOENT, be suspicious
+    bool artifactsPresent = (err == 0) || (err < 0 && errno != ENOENT);
+    if (artifactsPresent) {
+        auto verificationResult = verifyArtifacts(*key, supportsFsVerity);
+        if (!verificationResult.ok()) {
+            LOG(ERROR) << verificationResult.error().message();
+            return -1;
+        }
+    }
 
     if (!artifactsValid || kForceCompilation) {
-        removeArtifacts();
-
         LOG(INFO) << "Starting compilation... ";
         bool ret = compileArtifacts(kForceCompilation);
         LOG(INFO) << "Compilation done, returned " << ret;
 
-        verityStatus = addFilesToVerityRecursive(kArtArtifactsDir, key.value());
+        Result<std::map<std::string, std::string>> digests;
+        if (supportsFsVerity) {
+            digests = addFilesToVerityRecursive(kArtArtifactsDir, *key);
+        } else {
+            // If we can't use verity, just compute the root hashes and store
+            // those, so we can reverify them at the next boot.
+            digests = computeDigests(kArtArtifactsDir);
+        }
+        if (!digests.ok()) {
+            LOG(ERROR) << digests.error().message();
+            return -1;
+        }
 
-        if (!verityStatus.ok()) {
-            LOG(ERROR) << "Failed to add " << verityStatus.error().message();
+        auto persistStatus = persistDigests(*digests, *key);
+        if (!persistStatus.ok()) {
+            LOG(ERROR) << persistStatus.error().message();
             return -1;
         }
     }
 
-    // TODO we want to make sure Zygote only picks up the artifacts if we deemed
-    // everything was ok here. We could use a sysprop, or some other mechanism?
     LOG(INFO) << "On-device signing done.";
 
     scope_guard.Disable();
+    // At this point, we're done with the key for sure
+    SetProperty(kOdsignKeyDoneProp, "1");
+    // And we did a successful verification
+    SetProperty(kOdsignVerificationDoneProp, "1");
+    SetProperty(kOdsignVerificationStatusProp, kOdsignVerificationStatusValid);
     return 0;
 }
diff --git a/ondevice-signing/proto/Android.bp b/ondevice-signing/proto/Android.bp
new file mode 100644
index 0000000..fd48f31
--- /dev/null
+++ b/ondevice-signing/proto/Android.bp
@@ -0,0 +1,29 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+cc_library_static {
+    name: "lib_odsign_proto",
+    host_supported: true,
+    proto: {
+        export_proto_headers: true,
+        type: "full",
+    },
+    srcs: ["odsign_info.proto"],
+}
diff --git a/ondevice-signing/proto/odsign_info.proto b/ondevice-signing/proto/odsign_info.proto
new file mode 100644
index 0000000..9d49c6c
--- /dev/null
+++ b/ondevice-signing/proto/odsign_info.proto
@@ -0,0 +1,24 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = "proto3";
+
+package odsign.proto;
+
+message OdsignInfo {
+  // Map of artifact files to their hashes
+  map<string, string> file_hashes = 1;
+}