Merge "Keystore 2.0: Allow by key id usage of granted keys."
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..99cf5ff
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+*.*~
diff --git a/OWNERS b/OWNERS
index 684cca3..fca66f8 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,5 +1,4 @@
 swillden@google.com
 cbrubaker@google.com
 jdanis@google.com
-kroot@google.com
-bdc@google.com
+kroot@google.com
\ No newline at end of file
diff --git a/identity/Android.bp b/identity/Android.bp
index c0f1635..dd61930 100644
--- a/identity/Android.bp
+++ b/identity/Android.bp
@@ -5,6 +5,7 @@
         // "-Werror",
         "-Wextra",
         "-Wunused",
+        "-Wno-deprecated-declarations",
     ],
     sanitize: {
         misc_undefined : ["integer"],
@@ -38,8 +39,8 @@
         "libkeystore-attestation-application-id",
     ],
     static_libs: [
-        "android.hardware.identity-cpp",
-        "android.hardware.keymaster-cpp",
+        "android.hardware.identity-unstable-cpp",
+        "android.hardware.keymaster-unstable-cpp",
         "libcppbor",
     ]
 }
diff --git a/identity/Credential.cpp b/identity/Credential.cpp
index 28ba752..4a2bae1 100644
--- a/identity/Credential.cpp
+++ b/identity/Credential.cpp
@@ -36,6 +36,7 @@
 #include "Credential.h"
 #include "CredentialData.h"
 #include "Util.h"
+#include "WritableCredential.h"
 
 namespace android {
 namespace security {
@@ -47,6 +48,8 @@
 
 using android::security::keystore::IKeystoreService;
 
+using ::android::hardware::identity::IWritableIdentityCredential;
+
 using ::android::hardware::identity::support::ecKeyPairGetPkcs12;
 using ::android::hardware::identity::support::ecKeyPairGetPrivateKey;
 using ::android::hardware::identity::support::ecKeyPairGetPublicKey;
@@ -58,25 +61,26 @@
 using AidlVerificationToken = android::hardware::keymaster::VerificationToken;
 
 Credential::Credential(CipherSuite cipherSuite, const std::string& dataPath,
-                       const std::string& credentialName)
-    : cipherSuite_(cipherSuite), dataPath_(dataPath), credentialName_(credentialName) {}
+                       const std::string& credentialName, uid_t callingUid,
+                       HardwareInformation hwInfo, sp<IIdentityCredentialStore> halStoreBinder,
+                       int halApiVersion)
+    : cipherSuite_(cipherSuite), dataPath_(dataPath), credentialName_(credentialName),
+      callingUid_(callingUid), hwInfo_(std::move(hwInfo)), halStoreBinder_(halStoreBinder),
+      halApiVersion_(halApiVersion) {}
 
 Credential::~Credential() {}
 
-Status Credential::loadCredential(sp<IIdentityCredentialStore> halStoreBinder) {
-    uid_t callingUid = android::IPCThreadState::self()->getCallingUid();
-    sp<CredentialData> data = new CredentialData(dataPath_, callingUid, credentialName_);
+Status Credential::ensureOrReplaceHalBinder() {
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
     if (!data->loadFromDisk()) {
         LOG(ERROR) << "Error loading data for credential";
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error loading data for credential");
     }
 
-    data_ = data;
-
     sp<IIdentityCredential> halBinder;
     Status status =
-        halStoreBinder->getCredential(cipherSuite_, data_->getCredentialData(), &halBinder);
+        halStoreBinder_->getCredential(cipherSuite_, data->getCredentialData(), &halBinder);
     if (!status.isOk() && status.exceptionCode() == binder::Status::EX_SERVICE_SPECIFIC) {
         int code = status.serviceSpecificErrorCode();
         if (code == IIdentityCredentialStore::STATUS_CIPHER_SUITE_NOT_SUPPORTED) {
@@ -87,21 +91,33 @@
         LOG(ERROR) << "Error getting HAL binder";
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC);
     }
-
     halBinder_ = halBinder;
 
     return Status::ok();
 }
 
 Status Credential::getCredentialKeyCertificateChain(std::vector<uint8_t>* _aidl_return) {
-    *_aidl_return = data_->getAttestationCertificate();
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+    *_aidl_return = data->getAttestationCertificate();
     return Status::ok();
 }
 
 // Returns operation handle
-Status Credential::selectAuthKey(bool allowUsingExhaustedKeys, int64_t* _aidl_return) {
+Status Credential::selectAuthKey(bool allowUsingExhaustedKeys, bool allowUsingExpiredKeys,
+                                 int64_t* _aidl_return) {
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
 
-    selectedAuthKey_ = data_->selectAuthKey(allowUsingExhaustedKeys);
+    selectedAuthKey_ = data->selectAuthKey(allowUsingExhaustedKeys, allowUsingExpiredKeys);
     if (selectedAuthKey_ == nullptr) {
         return Status::fromServiceSpecificError(
             ICredentialStore::ERROR_NO_AUTHENTICATION_KEY_AVAILABLE,
@@ -174,16 +190,23 @@
                               const vector<RequestNamespaceParcel>& requestNamespaces,
                               const vector<uint8_t>& sessionTranscript,
                               const vector<uint8_t>& readerSignature, bool allowUsingExhaustedKeys,
-                              GetEntriesResultParcel* _aidl_return) {
+                              bool allowUsingExpiredKeys, GetEntriesResultParcel* _aidl_return) {
     GetEntriesResultParcel ret;
 
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+
     // Calculate requestCounts ahead of time and be careful not to include
     // elements that don't exist.
     //
     // Also go through and figure out which access control profiles to include
     // in the startRetrieval() call.
     vector<int32_t> requestCounts;
-    const vector<SecureAccessControlProfile>& allProfiles = data_->getSecureAccessControlProfiles();
+    const vector<SecureAccessControlProfile>& allProfiles = data->getSecureAccessControlProfiles();
 
     // We don't support ACP identifiers which isn't in the range 0 to 31. This
     // guarantee exists so it's feasible to implement the TA part of an Identity
@@ -202,13 +225,13 @@
     for (const RequestNamespaceParcel& rns : requestNamespaces) {
         size_t numEntriesInNsToRequest = 0;
         for (const RequestEntryParcel& rep : rns.entries) {
-            if (data_->hasEntryData(rns.namespaceName, rep.name)) {
+            if (data->hasEntryData(rns.namespaceName, rep.name)) {
                 numEntriesInNsToRequest++;
             }
 
-            optional<EntryData> data = data_->getEntryData(rns.namespaceName, rep.name);
-            if (data) {
-                for (int32_t id : data.value().accessControlProfileIds) {
+            optional<EntryData> eData = data->getEntryData(rns.namespaceName, rep.name);
+            if (eData) {
+                for (int32_t id : eData.value().accessControlProfileIds) {
                     if (id < 0 || id >= 32) {
                         LOG(ERROR) << "Invalid accessControlProfileId " << id << " for "
                                    << rns.namespaceName << ": " << rep.name;
@@ -282,7 +305,7 @@
     if (userAuthNeeded) {
         vector<uint8_t> authTokenBytes;
         vector<uint8_t> verificationTokenBytes;
-        if (!getTokensFromKeystore(selectedChallenge_, data_->getSecureUserId(),
+        if (!getTokensFromKeystore(selectedChallenge_, data->getSecureUserId(),
                                    authTokenMaxAgeMillis, authTokenBytes, verificationTokenBytes)) {
             LOG(ERROR) << "Error getting tokens from keystore";
             return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
@@ -331,7 +354,7 @@
     const AuthKeyData* authKey = selectedAuthKey_;
     if (sessionTranscript.size() > 0) {
         if (authKey == nullptr) {
-            authKey = data_->selectAuthKey(allowUsingExhaustedKeys);
+            authKey = data->selectAuthKey(allowUsingExhaustedKeys, allowUsingExpiredKeys);
             if (authKey == nullptr) {
                 return Status::fromServiceSpecificError(
                     ICredentialStore::ERROR_NO_AUTHENTICATION_KEY_AVAILABLE,
@@ -351,7 +374,7 @@
         RequestNamespace ns;
         ns.namespaceName = rns.namespaceName;
         for (const RequestEntryParcel& rep : rns.entries) {
-            optional<EntryData> entryData = data_->getEntryData(rns.namespaceName, rep.name);
+            optional<EntryData> entryData = data->getEntryData(rns.namespaceName, rep.name);
             if (entryData) {
                 RequestDataItem di;
                 di.name = rep.name;
@@ -406,16 +429,16 @@
             ResultEntryParcel resultEntryParcel;
             resultEntryParcel.name = rep.name;
 
-            optional<EntryData> data = data_->getEntryData(rns.namespaceName, rep.name);
-            if (!data) {
+            optional<EntryData> eData = data->getEntryData(rns.namespaceName, rep.name);
+            if (!eData) {
                 resultEntryParcel.status = STATUS_NO_SUCH_ENTRY;
                 resultNamespaceParcel.entries.push_back(resultEntryParcel);
                 continue;
             }
 
             status =
-                halBinder_->startRetrieveEntryValue(rns.namespaceName, rep.name, data.value().size,
-                                                    data.value().accessControlProfileIds);
+                halBinder_->startRetrieveEntryValue(rns.namespaceName, rep.name, eData.value().size,
+                                                    eData.value().accessControlProfileIds);
             if (!status.isOk() && status.exceptionCode() == binder::Status::EX_SERVICE_SPECIFIC) {
                 int code = status.serviceSpecificErrorCode();
                 if (code == IIdentityCredentialStore::STATUS_USER_AUTHENTICATION_FAILED) {
@@ -441,7 +464,7 @@
             }
 
             vector<uint8_t> value;
-            for (const auto& encryptedChunk : data.value().encryptedChunks) {
+            for (const auto& encryptedChunk : eData.value().encryptedChunks) {
                 vector<uint8_t> chunk;
                 status = halBinder_->retrieveEntryValue(encryptedChunk, &chunk);
                 if (!status.isOk()) {
@@ -467,7 +490,7 @@
 
     // Ensure useCount is updated on disk.
     if (authKey != nullptr) {
-        if (!data_->saveToDisk()) {
+        if (!data->saveToDisk()) {
             return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                     "Error saving data");
         }
@@ -479,11 +502,19 @@
 
 Status Credential::deleteCredential(vector<uint8_t>* _aidl_return) {
     vector<uint8_t> proofOfDeletionSignature;
+
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+
     Status status = halBinder_->deleteCredential(&proofOfDeletionSignature);
     if (!status.isOk()) {
         return halStatusToGenericError(status);
     }
-    if (!data_->deleteCredential()) {
+    if (!data->deleteCredential()) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error deleting credential data on disk");
     }
@@ -491,6 +522,47 @@
     return Status::ok();
 }
 
+Status Credential::deleteWithChallenge(const vector<uint8_t>& challenge,
+                                       vector<uint8_t>* _aidl_return) {
+    if (halApiVersion_ < 3) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_NOT_SUPPORTED,
+                                                "Not implemented by HAL");
+    }
+    vector<uint8_t> proofOfDeletionSignature;
+
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+
+    Status status = halBinder_->deleteCredentialWithChallenge(challenge, &proofOfDeletionSignature);
+    if (!status.isOk()) {
+        return halStatusToGenericError(status);
+    }
+    if (!data->deleteCredential()) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error deleting credential data on disk");
+    }
+    *_aidl_return = proofOfDeletionSignature;
+    return Status::ok();
+}
+
+Status Credential::proveOwnership(const vector<uint8_t>& challenge, vector<uint8_t>* _aidl_return) {
+    if (halApiVersion_ < 3) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_NOT_SUPPORTED,
+                                                "Not implemented by HAL");
+    }
+    vector<uint8_t> proofOfOwnershipSignature;
+    Status status = halBinder_->proveOwnership(challenge, &proofOfOwnershipSignature);
+    if (!status.isOk()) {
+        return halStatusToGenericError(status);
+    }
+    *_aidl_return = proofOfOwnershipSignature;
+    return Status::ok();
+}
+
 Status Credential::createEphemeralKeyPair(vector<uint8_t>* _aidl_return) {
     vector<uint8_t> keyPair;
     Status status = halBinder_->createEphemeralKeyPair(&keyPair);
@@ -522,8 +594,14 @@
 }
 
 Status Credential::setAvailableAuthenticationKeys(int32_t keyCount, int32_t maxUsesPerKey) {
-    data_->setAvailableAuthenticationKeys(keyCount, maxUsesPerKey);
-    if (!data_->saveToDisk()) {
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+    data->setAvailableAuthenticationKeys(keyCount, maxUsesPerKey);
+    if (!data->saveToDisk()) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error saving data");
     }
@@ -531,8 +609,14 @@
 }
 
 Status Credential::getAuthKeysNeedingCertification(vector<AuthKeyParcel>* _aidl_return) {
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
     optional<vector<vector<uint8_t>>> keysNeedingCert =
-        data_->getAuthKeysNeedingCertification(halBinder_);
+        data->getAuthKeysNeedingCertification(halBinder_);
     if (!keysNeedingCert) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error getting auth keys neededing certification");
@@ -543,7 +627,7 @@
         authKeyParcel.x509cert = key;
         authKeyParcels.push_back(authKeyParcel);
     }
-    if (!data_->saveToDisk()) {
+    if (!data->saveToDisk()) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error saving data");
     }
@@ -553,13 +637,48 @@
 
 Status Credential::storeStaticAuthenticationData(const AuthKeyParcel& authenticationKey,
                                                  const vector<uint8_t>& staticAuthData) {
-    if (!data_->storeStaticAuthenticationData(authenticationKey.x509cert, staticAuthData)) {
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+    if (!data->storeStaticAuthenticationData(authenticationKey.x509cert,
+                                             std::numeric_limits<int64_t>::max(), staticAuthData)) {
         return Status::fromServiceSpecificError(
             ICredentialStore::ERROR_AUTHENTICATION_KEY_NOT_FOUND,
             "Error finding authentication key to store static "
             "authentication data for");
     }
-    if (!data_->saveToDisk()) {
+    if (!data->saveToDisk()) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error saving data");
+    }
+    return Status::ok();
+}
+
+Status
+Credential::storeStaticAuthenticationDataWithExpiration(const AuthKeyParcel& authenticationKey,
+                                                        int64_t expirationDateMillisSinceEpoch,
+                                                        const vector<uint8_t>& staticAuthData) {
+    if (halApiVersion_ < 3) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_NOT_SUPPORTED,
+                                                "Not implemented by HAL");
+    }
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+    if (!data->storeStaticAuthenticationData(authenticationKey.x509cert,
+                                             expirationDateMillisSinceEpoch, staticAuthData)) {
+        return Status::fromServiceSpecificError(
+            ICredentialStore::ERROR_AUTHENTICATION_KEY_NOT_FOUND,
+            "Error finding authentication key to store static "
+            "authentication data for");
+    }
+    if (!data->saveToDisk()) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error saving data");
     }
@@ -567,7 +686,13 @@
 }
 
 Status Credential::getAuthenticationDataUsageCount(vector<int32_t>* _aidl_return) {
-    const vector<AuthKeyData>& authKeyDatas = data_->getAuthKeyDatas();
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+    const vector<AuthKeyData>& authKeyDatas = data->getAuthKeyDatas();
     vector<int32_t> ret;
     for (const AuthKeyData& authKeyData : authKeyDatas) {
         ret.push_back(authKeyData.useCount);
@@ -576,6 +701,80 @@
     return Status::ok();
 }
 
+optional<string> extractDocType(const vector<uint8_t>& credentialData) {
+    auto [item, _ /* newPos */, message] = cppbor::parse(credentialData);
+    if (item == nullptr) {
+        LOG(ERROR) << "CredentialData is not valid CBOR: " << message;
+        return {};
+    }
+    const cppbor::Array* array = item->asArray();
+    if (array == nullptr || array->size() < 1) {
+        LOG(ERROR) << "CredentialData array with at least one element";
+        return {};
+    }
+    const cppbor::Tstr* tstr = ((*array)[0])->asTstr();
+    if (tstr == nullptr) {
+        LOG(ERROR) << "First item in CredentialData is not a string";
+        return {};
+    }
+    return tstr->value();
+}
+
+Status Credential::update(sp<IWritableCredential>* _aidl_return) {
+    if (halApiVersion_ < 3) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_NOT_SUPPORTED,
+                                                "Not implemented by HAL");
+    }
+    sp<CredentialData> data = new CredentialData(dataPath_, callingUid_, credentialName_);
+    if (!data->loadFromDisk()) {
+        LOG(ERROR) << "Error loading data for credential";
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Error loading data for credential");
+    }
+
+    sp<IWritableIdentityCredential> halWritableCredential;
+    Status status = halBinder_->updateCredential(&halWritableCredential);
+    if (!status.isOk()) {
+        return halStatusToGenericError(status);
+    }
+
+    optional<string> docType = extractDocType(data->getCredentialData());
+    if (!docType) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Unable to extract DocType from CredentialData");
+    }
+
+    // NOTE: The caller is expected to call WritableCredential::personalize() which will
+    // write brand new data to disk, specifically it will overwrite any data already
+    // have _including_ authentication keys.
+    //
+    // 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_);
+
+    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().
+    //
+    // The joys and pitfalls of mutable objects...
+    //
+    writableCredential->setCredentialUpdatedCallback([this] {
+        Status status = this->ensureOrReplaceHalBinder();
+        if (!status.isOk()) {
+            LOG(ERROR) << "Error loading credential";
+        }
+    });
+
+    *_aidl_return = writableCredential;
+    return Status::ok();
+}
+
 }  // namespace identity
 }  // namespace security
 }  // namespace android
diff --git a/identity/Credential.h b/identity/Credential.h
index e2880d9..7f08515 100644
--- a/identity/Credential.h
+++ b/identity/Credential.h
@@ -36,6 +36,7 @@
 using ::std::vector;
 
 using ::android::hardware::identity::CipherSuite;
+using ::android::hardware::identity::HardwareInformation;
 using ::android::hardware::identity::IIdentityCredential;
 using ::android::hardware::identity::IIdentityCredentialStore;
 using ::android::hardware::identity::RequestDataItem;
@@ -43,10 +44,12 @@
 
 class Credential : public BnCredential {
   public:
-    Credential(CipherSuite cipherSuite, const string& dataPath, const string& credentialName);
+    Credential(CipherSuite cipherSuite, const string& dataPath, const string& credentialName,
+               uid_t callingUid, HardwareInformation hwInfo,
+               sp<IIdentityCredentialStore> halStoreBinder, int halApiVersion);
     ~Credential();
 
-    Status loadCredential(sp<IIdentityCredentialStore> halStoreBinder);
+    Status ensureOrReplaceHalBinder();
 
     // ICredential overrides
     Status createEphemeralKeyPair(vector<uint8_t>* _aidl_return) override;
@@ -55,33 +58,47 @@
 
     Status deleteCredential(vector<uint8_t>* _aidl_return) override;
 
+    Status deleteWithChallenge(const vector<uint8_t>& challenge,
+                               vector<uint8_t>* _aidl_return) override;
+
+    Status proveOwnership(const vector<uint8_t>& challenge, vector<uint8_t>* _aidl_return) override;
+
     Status getCredentialKeyCertificateChain(vector<uint8_t>* _aidl_return) override;
 
-    Status selectAuthKey(bool allowUsingExhaustedKeys, int64_t* _aidl_return) override;
+    Status selectAuthKey(bool allowUsingExhaustedKeys, bool allowUsingExpiredKeys,
+                         int64_t* _aidl_return) override;
 
     Status getEntries(const vector<uint8_t>& requestMessage,
                       const vector<RequestNamespaceParcel>& requestNamespaces,
                       const vector<uint8_t>& sessionTranscript,
                       const vector<uint8_t>& readerSignature, bool allowUsingExhaustedKeys,
-                      GetEntriesResultParcel* _aidl_return) override;
+                      bool allowUsingExpiredKeys, GetEntriesResultParcel* _aidl_return) override;
 
     Status setAvailableAuthenticationKeys(int32_t keyCount, int32_t maxUsesPerKey) override;
     Status getAuthKeysNeedingCertification(vector<AuthKeyParcel>* _aidl_return) override;
     Status storeStaticAuthenticationData(const AuthKeyParcel& authenticationKey,
                                          const vector<uint8_t>& staticAuthData) override;
+    Status
+    storeStaticAuthenticationDataWithExpiration(const AuthKeyParcel& authenticationKey,
+                                                int64_t expirationDateMillisSinceEpoch,
+                                                const vector<uint8_t>& staticAuthData) override;
     Status getAuthenticationDataUsageCount(vector<int32_t>* _aidl_return) override;
 
+    Status update(sp<IWritableCredential>* _aidl_return) override;
+
   private:
     CipherSuite cipherSuite_;
     string dataPath_;
     string credentialName_;
+    uid_t callingUid_;
+    HardwareInformation hwInfo_;
+    sp<IIdentityCredentialStore> halStoreBinder_;
 
     const AuthKeyData* selectedAuthKey_ = nullptr;
     uint64_t selectedChallenge_ = 0;
 
-    sp<CredentialData> data_;
-
     sp<IIdentityCredential> halBinder_;
+    int halApiVersion_;
 
     ssize_t
     calcExpectedDeviceNameSpacesSize(const vector<uint8_t>& requestMessage,
diff --git a/identity/CredentialData.cpp b/identity/CredentialData.cpp
index b4e6641..96c436a 100644
--- a/identity/CredentialData.cpp
+++ b/identity/CredentialData.cpp
@@ -16,6 +16,8 @@
 
 #define LOG_TAG "CredentialData"
 
+#include <chrono>
+
 #include <fcntl.h>
 #include <stdlib.h>
 #include <sys/stat.h>
@@ -119,12 +121,15 @@
     cppbor::Array authKeyDatasArray;
     for (const AuthKeyData& data : authKeyDatas_) {
         cppbor::Array array;
+        // Fields 0-6 was in the original version in Android 11
         array.add(data.certificate);
         array.add(data.keyBlob);
         array.add(data.staticAuthenticationData);
         array.add(data.pendingCertificate);
         array.add(data.pendingKeyBlob);
         array.add(data.useCount);
+        // Field 7 was added in Android 12
+        array.add(data.expirationDateMillisSinceEpoch);
         authKeyDatasArray.add(std::move(array));
     }
     map.add("authKeyData", std::move(authKeyDatasArray));
@@ -183,9 +188,17 @@
         LOG(ERROR) << "One or more items in AuthKeyData array in CBOR is of wrong type";
         return {};
     }
+    // expirationDateMillisSinceEpoch was added as the 7th element for Android 12. If not
+    // present, default to longest possible expiration date.
+    int64_t expirationDateMillisSinceEpoch = INT64_MAX;
+    if (array->size() >= 7) {
+        const cppbor::Int* itemExpirationDateMillisSinceEpoch = ((*array)[6])->asInt();
+        expirationDateMillisSinceEpoch = itemExpirationDateMillisSinceEpoch->value();
+    }
     AuthKeyData authKeyData;
     authKeyData.certificate = itemCertificate->value();
     authKeyData.keyBlob = itemKeyBlob->value();
+    authKeyData.expirationDateMillisSinceEpoch = expirationDateMillisSinceEpoch;
     authKeyData.staticAuthenticationData = itemStaticAuthenticationData->value();
     authKeyData.pendingCertificate = itemPendingCertificate->value();
     authKeyData.pendingKeyBlob = itemPendingKeyBlob->value();
@@ -232,7 +245,6 @@
 }
 
 bool CredentialData::loadFromDisk() {
-
     // Reset all data.
     credentialData_.clear();
     attestationCertificate_.clear();
@@ -487,16 +499,28 @@
     return authKeyDatas_;
 }
 
-const AuthKeyData* CredentialData::selectAuthKey(bool allowUsingExhaustedKeys) {
+pair<int /* keyCount */, int /*maxUsersPerKey */> CredentialData::getAvailableAuthenticationKeys() {
+    return std::make_pair(keyCount_, maxUsesPerKey_);
+}
+
+AuthKeyData* CredentialData::findAuthKey_(bool allowUsingExhaustedKeys,
+                                          bool allowUsingExpiredKeys) {
     AuthKeyData* candidate = nullptr;
 
+    int64_t nowMilliSeconds =
+        std::chrono::system_clock::to_time_t(std::chrono::system_clock::now()) * 1000;
+
     int n = 0;
-    int candidateNum = -1;
     for (AuthKeyData& data : authKeyDatas_) {
+        if (nowMilliSeconds > data.expirationDateMillisSinceEpoch) {
+            if (!allowUsingExpiredKeys) {
+                continue;
+            }
+        }
         if (data.certificate.size() != 0) {
+            // Not expired, include in normal check
             if (candidate == nullptr || data.useCount < candidate->useCount) {
                 candidate = &data;
-                candidateNum = n;
             }
         }
         n++;
@@ -510,6 +534,28 @@
         return nullptr;
     }
 
+    return candidate;
+}
+
+const AuthKeyData* CredentialData::selectAuthKey(bool allowUsingExhaustedKeys,
+                                                 bool allowUsingExpiredKeys) {
+    AuthKeyData* candidate;
+
+    // First try to find a un-expired key..
+    candidate = findAuthKey_(allowUsingExhaustedKeys, false);
+    if (candidate == nullptr) {
+        // That didn't work, there are no un-expired keys and we don't allow using expired keys.
+        if (!allowUsingExpiredKeys) {
+            return nullptr;
+        }
+
+        // See if there's an expired key then...
+        candidate = findAuthKey_(allowUsingExhaustedKeys, true);
+        if (candidate == nullptr) {
+            return nullptr;
+        }
+    }
+
     candidate->useCount += 1;
     return candidate;
 }
@@ -519,8 +565,14 @@
 
     vector<vector<uint8_t>> keysNeedingCert;
 
+    int64_t nowMilliSeconds =
+        std::chrono::system_clock::to_time_t(std::chrono::system_clock::now()) * 1000;
+
     for (AuthKeyData& data : authKeyDatas_) {
-        bool newKeyNeeded = (data.certificate.size() == 0) || (data.useCount >= maxUsesPerKey_);
+        bool keyExceedUseCount = (data.useCount >= maxUsesPerKey_);
+        bool keyBeyondExpirationDate = (nowMilliSeconds > data.expirationDateMillisSinceEpoch);
+        bool newKeyNeeded =
+            (data.certificate.size() == 0) || keyExceedUseCount || keyBeyondExpirationDate;
         bool certificationPending = (data.pendingCertificate.size() > 0);
         if (newKeyNeeded && !certificationPending) {
             vector<uint8_t> signingKeyBlob;
@@ -543,11 +595,13 @@
 }
 
 bool CredentialData::storeStaticAuthenticationData(const vector<uint8_t>& authenticationKey,
+                                                   int64_t expirationDateMillisSinceEpoch,
                                                    const vector<uint8_t>& staticAuthData) {
     for (AuthKeyData& data : authKeyDatas_) {
         if (data.pendingCertificate == authenticationKey) {
             data.certificate = data.pendingCertificate;
             data.keyBlob = data.pendingKeyBlob;
+            data.expirationDateMillisSinceEpoch = expirationDateMillisSinceEpoch;
             data.staticAuthenticationData = staticAuthData;
             data.pendingCertificate.clear();
             data.pendingKeyBlob.clear();
diff --git a/identity/CredentialData.h b/identity/CredentialData.h
index 7995828..b037997 100644
--- a/identity/CredentialData.h
+++ b/identity/CredentialData.h
@@ -55,6 +55,7 @@
 
     vector<uint8_t> certificate;
     vector<uint8_t> keyBlob;
+    int64_t expirationDateMillisSinceEpoch;
     vector<uint8_t> staticAuthenticationData;
     vector<uint8_t> pendingCertificate;
     vector<uint8_t> pendingKeyBlob;
@@ -106,17 +107,22 @@
 
     const vector<AuthKeyData>& getAuthKeyDatas() const;
 
+    pair<int /* keyCount */, int /*maxUsersPerKey */> getAvailableAuthenticationKeys();
+
     // Returns |nullptr| if a suitable key cannot be found. Otherwise returns
     // the authentication and increases its use-count.
-    const AuthKeyData* selectAuthKey(bool allowUsingExhaustedKeys);
+    const AuthKeyData* selectAuthKey(bool allowUsingExhaustedKeys, bool allowUsingExpiredKeys);
 
     optional<vector<vector<uint8_t>>>
     getAuthKeysNeedingCertification(const sp<IIdentityCredential>& halBinder);
 
     bool storeStaticAuthenticationData(const vector<uint8_t>& authenticationKey,
+                                       int64_t expirationDateMillisSinceEpoch,
                                        const vector<uint8_t>& staticAuthData);
 
   private:
+    AuthKeyData* findAuthKey_(bool allowUsingExhaustedKeys, bool allowUsingExpiredKeys);
+
     // Set by constructor.
     //
     string dataPath_;
diff --git a/identity/CredentialStore.cpp b/identity/CredentialStore.cpp
index e3a825b..f77294e 100644
--- a/identity/CredentialStore.cpp
+++ b/identity/CredentialStore.cpp
@@ -41,11 +41,12 @@
         LOG(ERROR) << "Error getting hardware information: " << status.toString8();
         return false;
     }
+    halApiVersion_ = hal_->getInterfaceVersion();
 
-    LOG(INFO) << "Connected to Identity Credential HAL with name '" << hwInfo_.credentialStoreName
-              << "' authored by '" << hwInfo_.credentialStoreAuthorName << "' with chunk size "
-              << hwInfo_.dataChunkSize << " and directoAccess set to "
-              << (hwInfo_.isDirectAccess ? "true" : "false");
+    LOG(INFO) << "Connected to Identity Credential HAL with API version " << halApiVersion_
+              << " and name '" << hwInfo_.credentialStoreName << "' authored by '"
+              << hwInfo_.credentialStoreAuthorName << "' with chunk size " << hwInfo_.dataChunkSize
+              << " and directoAccess set to " << (hwInfo_.isDirectAccess ? "true" : "false");
     return true;
 }
 
@@ -89,7 +90,7 @@
     }
 
     sp<IWritableCredential> writableCredential = new WritableCredential(
-        dataPath_, credentialName, docType, hwInfo_.dataChunkSize, halWritableCredential);
+        dataPath_, credentialName, docType, false, hwInfo_, halWritableCredential, halApiVersion_);
     *_aidl_return = writableCredential;
     return Status::ok();
 }
@@ -112,9 +113,10 @@
 
     // Note: IdentityCredentialStore.java's CipherSuite enumeration and CipherSuite from the
     // HAL is manually kept in sync. So this cast is safe.
-    sp<Credential> credential = new Credential(CipherSuite(cipherSuite), dataPath_, credentialName);
+    sp<Credential> credential = new Credential(CipherSuite(cipherSuite), dataPath_, credentialName,
+                                               callingUid, hwInfo_, hal_, halApiVersion_);
 
-    Status loadStatus = credential->loadCredential(hal_);
+    Status loadStatus = credential->ensureOrReplaceHalBinder();
     if (!loadStatus.isOk()) {
         LOG(ERROR) << "Error loading credential";
     } else {
diff --git a/identity/CredentialStore.h b/identity/CredentialStore.h
index 24b2b4d..15da4eb 100644
--- a/identity/CredentialStore.h
+++ b/identity/CredentialStore.h
@@ -57,6 +57,7 @@
     string dataPath_;
 
     sp<IIdentityCredentialStore> hal_;
+    int halApiVersion_;
 
     HardwareInformation hwInfo_;
 };
diff --git a/identity/WritableCredential.cpp b/identity/WritableCredential.cpp
index a932dcf..d0688b8 100644
--- a/identity/WritableCredential.cpp
+++ b/identity/WritableCredential.cpp
@@ -39,13 +39,19 @@
 using ::android::hardware::identity::support::chunkVector;
 
 WritableCredential::WritableCredential(const string& dataPath, const string& credentialName,
-                                       const string& docType, size_t dataChunkSize,
-                                       sp<IWritableIdentityCredential> halBinder)
-    : dataPath_(dataPath), credentialName_(credentialName), docType_(docType),
-      dataChunkSize_(dataChunkSize), halBinder_(halBinder) {}
+                                       const string& docType, bool isUpdate,
+                                       HardwareInformation hwInfo,
+                                       sp<IWritableIdentityCredential> halBinder, int halApiVersion)
+    : dataPath_(dataPath), credentialName_(credentialName), docType_(docType), isUpdate_(isUpdate),
+      hwInfo_(std::move(hwInfo)), halBinder_(halBinder), halApiVersion_(halApiVersion) {}
 
 WritableCredential::~WritableCredential() {}
 
+void WritableCredential::setCredentialUpdatedCallback(
+    std::function<void()>&& onCredentialUpdatedCallback) {
+    onCredentialUpdatedCallback_ = onCredentialUpdatedCallback;
+}
+
 Status WritableCredential::ensureAttestationCertificateExists(const vector<uint8_t>& challenge) {
     if (!attestationCertificate_.empty()) {
         return Status::ok();
@@ -79,7 +85,10 @@
 
 Status WritableCredential::getCredentialKeyCertificateChain(const vector<uint8_t>& challenge,
                                                             vector<uint8_t>* _aidl_return) {
-
+    if (isUpdate_) {
+        return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
+                                                "Cannot be called for an update");
+    }
     Status ensureStatus = ensureAttestationCertificateExists(challenge);
     if (!ensureStatus.isOk()) {
         return ensureStatus;
@@ -89,6 +98,15 @@
     return Status::ok();
 }
 
+void WritableCredential::setAttestationCertificate(const vector<uint8_t>& attestationCertificate) {
+    attestationCertificate_ = attestationCertificate;
+}
+
+void WritableCredential::setAvailableAuthenticationKeys(int keyCount, int maxUsesPerKey) {
+    keyCount_ = keyCount;
+    maxUsesPerKey_ = maxUsesPerKey;
+}
+
 ssize_t WritableCredential::calcExpectedProofOfProvisioningSize(
     const vector<AccessControlProfileParcel>& accessControlProfiles,
     const vector<EntryNamespaceParcel>& entryNamespaces) {
@@ -149,9 +167,12 @@
 WritableCredential::personalize(const vector<AccessControlProfileParcel>& accessControlProfiles,
                                 const vector<EntryNamespaceParcel>& entryNamespaces,
                                 int64_t secureUserId, vector<uint8_t>* _aidl_return) {
-    Status ensureStatus = ensureAttestationCertificateExists({0x00});  // Challenge cannot be empty.
-    if (!ensureStatus.isOk()) {
-        return ensureStatus;
+    if (!isUpdate_) {
+        Status ensureStatus =
+            ensureAttestationCertificateExists({0x00});  // Challenge cannot be empty.
+        if (!ensureStatus.isOk()) {
+            return ensureStatus;
+        }
     }
 
     uid_t callingUid = android::IPCThreadState::self()->getCallingUid();
@@ -203,7 +224,7 @@
 
     for (const EntryNamespaceParcel& ensParcel : entryNamespaces) {
         for (const EntryParcel& eParcel : ensParcel.entries) {
-            vector<vector<uint8_t>> chunks = chunkVector(eParcel.value, dataChunkSize_);
+            vector<vector<uint8_t>> chunks = chunkVector(eParcel.value, hwInfo_.dataChunkSize);
 
             vector<int32_t> ids;
             std::copy(eParcel.accessControlProfileIds.begin(),
@@ -240,11 +261,15 @@
     }
     data.setCredentialData(credentialData);
 
+    data.setAvailableAuthenticationKeys(keyCount_, maxUsesPerKey_);
+
     if (!data.saveToDisk()) {
         return Status::fromServiceSpecificError(ICredentialStore::ERROR_GENERIC,
                                                 "Error saving credential data to disk");
     }
 
+    onCredentialUpdatedCallback_();
+
     *_aidl_return = proofOfProvisioningSignature;
     return Status::ok();
 }
diff --git a/identity/WritableCredential.h b/identity/WritableCredential.h
index eb63aca..6ff31ae 100644
--- a/identity/WritableCredential.h
+++ b/identity/WritableCredential.h
@@ -29,6 +29,7 @@
 namespace identity {
 
 using ::android::binder::Status;
+using ::android::hardware::identity::HardwareInformation;
 using ::android::hardware::identity::IWritableIdentityCredential;
 using ::std::string;
 using ::std::vector;
@@ -36,9 +37,15 @@
 class WritableCredential : public BnWritableCredential {
   public:
     WritableCredential(const string& dataPath, const string& credentialName, const string& docType,
-                       size_t dataChunkSize, sp<IWritableIdentityCredential> halBinder);
+                       bool isUpdate, HardwareInformation hwInfo,
+                       sp<IWritableIdentityCredential> halBinder, int halApiVersion);
     ~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);
+
     // IWritableCredential overrides
     Status getCredentialKeyCertificateChain(const vector<uint8_t>& challenge,
                                             vector<uint8_t>* _aidl_return) override;
@@ -51,9 +58,16 @@
     string dataPath_;
     string credentialName_;
     string docType_;
-    size_t dataChunkSize_;
+    bool isUpdate_;
+    HardwareInformation hwInfo_;
     sp<IWritableIdentityCredential> halBinder_;
+    int halApiVersion_;
+
     vector<uint8_t> attestationCertificate_;
+    int keyCount_ = 0;
+    int maxUsesPerKey_ = 1;
+
+    std::function<void()> onCredentialUpdatedCallback_ = []() {};
 
     ssize_t calcExpectedProofOfProvisioningSize(
         const vector<AccessControlProfileParcel>& accessControlProfiles,
diff --git a/identity/binder/android/security/identity/ICredential.aidl b/identity/binder/android/security/identity/ICredential.aidl
index 7bd0df7..2165810 100644
--- a/identity/binder/android/security/identity/ICredential.aidl
+++ b/identity/binder/android/security/identity/ICredential.aidl
@@ -16,6 +16,8 @@
 
 package android.security.identity;
 
+import android.security.identity.IWritableCredential;
+
 import android.security.identity.RequestNamespaceParcel;
 import android.security.identity.GetEntriesResultParcel;
 import android.security.identity.AuthKeyParcel;
@@ -40,23 +42,35 @@
     void setReaderEphemeralPublicKey(in byte[] publicKey);
 
     byte[] deleteCredential();
+    byte[] deleteWithChallenge(in byte[] challenge);
+
+    byte[] proveOwnership(in byte[] challenge);
 
     byte[] getCredentialKeyCertificateChain();
 
-    long selectAuthKey(in boolean allowUsingExhaustedKeys);
+    long selectAuthKey(in boolean allowUsingExhaustedKeys,
+                       in boolean allowUsingExpiredKeys);
 
     GetEntriesResultParcel getEntries(in byte[] requestMessage,
                                       in RequestNamespaceParcel[] requestNamespaces,
                                       in byte[] sessionTranscript,
                                       in byte[] readerSignature,
-                                      in boolean allowUsingExhaustedKeys);
+                                      in boolean allowUsingExhaustedKeys,
+                                      in boolean allowUsingExpiredKeys);
 
     void setAvailableAuthenticationKeys(in int keyCount, in int maxUsesPerKey);
 
     AuthKeyParcel[] getAuthKeysNeedingCertification();
 
-    void storeStaticAuthenticationData(in AuthKeyParcel authenticationKey, in byte[] staticAuthData);
+    void storeStaticAuthenticationData(in AuthKeyParcel authenticationKey,
+                                       in byte[] staticAuthData);
+
+    void storeStaticAuthenticationDataWithExpiration(in AuthKeyParcel authenticationKey,
+                                       in long expirationDateMillisSinceEpoch,
+                                       in byte[] staticAuthData);
 
     int[] getAuthenticationDataUsageCount();
+
+    IWritableCredential update();
 }
 
diff --git a/identity/binder/android/security/identity/ICredentialStore.aidl b/identity/binder/android/security/identity/ICredentialStore.aidl
index 1039831..8357f47 100644
--- a/identity/binder/android/security/identity/ICredentialStore.aidl
+++ b/identity/binder/android/security/identity/ICredentialStore.aidl
@@ -39,6 +39,7 @@
     const int ERROR_AUTHENTICATION_KEY_NOT_FOUND = 9;
     const int ERROR_INVALID_ITEMS_REQUEST_MESSAGE = 10;
     const int ERROR_SESSION_TRANSCRIPT_MISMATCH = 11;
+    const int ERROR_NOT_SUPPORTED = 12;
 
     SecurityHardwareInfoParcel getSecurityHardwareInfo();
 
diff --git a/keystore2/Android.bp b/keystore2/Android.bp
index cf37fca..f9295ca 100644
--- a/keystore2/Android.bp
+++ b/keystore2/Android.bp
@@ -18,14 +18,15 @@
     srcs: ["src/lib.rs"],
 
     rustlibs: [
-        "android.hardware.security.keymint-rust",
-        "android.hardware.security.secureclock-rust",
+        "android.hardware.security.keymint-V1-rust",
+        "android.hardware.security.secureclock-V1-rust",
         "android.security.apc-rust",
         "android.security.authorization-rust",
         "android.security.compat-rust",
-        "android.system.keystore2-rust",
+        "android.system.keystore2-V1-rust",
         "libanyhow",
         "libbinder_rs",
+        "libkeystore2_aaid-rust",
         "libkeystore2_apc_compat-rust",
         "libkeystore2_crypto_rust",
         "libkeystore2_km_compat",
@@ -40,6 +41,16 @@
     ],
 }
 
+rust_library {
+    name: "libkeystore2_test_utils",
+    crate_name: "keystore2_test_utils",
+    srcs: ["test_utils/lib.rs"],
+    rustlibs: [
+        "liblog_rust",
+        "librand",
+    ]
+}
+
 rust_test {
     name: "keystore2_test",
     crate_name: "keystore2",
@@ -47,19 +58,21 @@
     test_suites: ["general-tests"],
     auto_gen_config: true,
     rustlibs: [
-        "android.hardware.security.keymint-rust",
-        "android.hardware.security.secureclock-rust",
+        "android.hardware.security.keymint-V1-rust",
+        "android.hardware.security.secureclock-V1-rust",
         "android.security.apc-rust",
         "android.security.authorization-rust",
         "android.security.compat-rust",
-        "android.system.keystore2-rust",
+        "android.system.keystore2-V1-rust",
         "libandroid_logger",
         "libanyhow",
         "libbinder_rs",
+        "libkeystore2_aaid-rust",
         "libkeystore2_apc_compat-rust",
         "libkeystore2_crypto_rust",
         "libkeystore2_km_compat",
         "libkeystore2_selinux",
+        "libkeystore2_test_utils",
         "liblazy_static",
         "liblibc",
         "liblibsqlite3_sys",
diff --git a/keystore2/aaid/Android.bp b/keystore2/aaid/Android.bp
new file mode 100644
index 0000000..2329400
--- /dev/null
+++ b/keystore2/aaid/Android.bp
@@ -0,0 +1,50 @@
+// Copyright 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.
+
+cc_library {
+    name: "libkeystore2_aaid",
+    srcs: [
+        "aaid.cpp",
+    ],
+    shared_libs: [
+        "libkeystore-attestation-application-id"
+    ],
+}
+
+rust_bindgen {
+    name: "libkeystore2_aaid_bindgen",
+    wrapper_src: "aaid.hpp",
+    crate_name: "keystore2_aaid_bindgen",
+    source_stem: "bindings",
+
+    bindgen_flags: [
+        "--size_t-is-usize",
+        "--whitelist-function=aaid_keystore_attestation_id",
+        "--whitelist-var=KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE",
+    ],
+}
+
+rust_library {
+    name: "libkeystore2_aaid-rust",
+    crate_name: "keystore2_aaid",
+    srcs: [
+        "lib.rs",
+    ],
+    rustlibs: [
+        "libkeystore2_aaid_bindgen",
+    ],
+    shared_libs: [
+        "libkeystore2_aaid",
+    ],
+}
diff --git a/keystore2/aaid/aaid.cpp b/keystore2/aaid/aaid.cpp
new file mode 100644
index 0000000..fd3faf6
--- /dev/null
+++ b/keystore2/aaid/aaid.cpp
@@ -0,0 +1,42 @@
+/*
+ * 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.
+ */
+
+#include "aaid.hpp"
+
+#include <keystore/keystore_attestation_id.h>
+
+using android::security::gather_attestation_application_id;
+
+uint32_t aaid_keystore_attestation_id(uint32_t uid, uint8_t* aaid, size_t* aaid_size) {
+    static_assert(sizeof(uint32_t) == sizeof(uid_t), "uid_t has unexpected size");
+    static_assert(sizeof(uint32_t) == sizeof(android::status_t), "status_t has unexpected size");
+    static_assert(KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE ==
+                      android::security::KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE,
+                  "KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE sizes don't match.");
+    auto result = gather_attestation_application_id(uid);
+    if (!result.isOk()) {
+        return result.status();
+    }
+    if (result.value().size() > KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE) {
+        return ::android::NO_MEMORY;
+    }
+    if (*aaid_size != KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE) {
+        return ::android::BAD_VALUE;
+    }
+    std::copy(result.value().begin(), result.value().end(), aaid);
+    *aaid_size = result.value().size();
+    return ::android::OK;
+}
diff --git a/keystore2/aaid/aaid.hpp b/keystore2/aaid/aaid.hpp
new file mode 100644
index 0000000..9e2c148
--- /dev/null
+++ b/keystore2/aaid/aaid.hpp
@@ -0,0 +1,44 @@
+/*
+ * 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 <stdint.h>
+#include <stddef.h>
+
+/**
+ * This is a redefinition of KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE in
+ * system/security/keystore/keystore_attestation_id.h and must be kept in sync.
+ * There is a static assert in aaid.cpp to assure that they are in sync.
+ * We redefine this here to avoid unnecessary build dependencies for
+ * the rust bindgen target.
+ */
+constexpr const size_t KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE = 1024;
+
+extern "C" {
+    /**
+     * Fills the buffer at aaid with the attestation application id of the app uid.
+     * The buffer must be exactly KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE bytes in size.
+     * *aaid_size is set to the number of bytes written to aaid.
+     *
+     * @param uid the uid of the app to retrieve the aaid for.
+     * @param aaid output buffer for the attestation id.
+     * @param aaid_size must be set to the size of the output buffer, which must be exactly
+     *          KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE bytes in size, by the caller. On success
+     *          it is set to the number of bytes written.
+     * @return OK on success.
+     */
+    uint32_t aaid_keystore_attestation_id(uint32_t uid, uint8_t* aaid, size_t* aaid_size);
+}
diff --git a/keystore2/aaid/lib.rs b/keystore2/aaid/lib.rs
new file mode 100644
index 0000000..3187198
--- /dev/null
+++ b/keystore2/aaid/lib.rs
@@ -0,0 +1,35 @@
+// Copyright 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.
+
+//! Rust binding for getting the attestation application id.
+
+use keystore2_aaid_bindgen::{
+    aaid_keystore_attestation_id, KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE,
+};
+
+/// Returns the attestation application id for the given uid or an error code
+/// corresponding to ::android::status_t.
+pub fn get_aaid(uid: u32) -> Result<Vec<u8>, u32> {
+    let mut buffer = [0u8; KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE];
+    let mut size = KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE;
+    // Safety:
+    // aaid_keystore_attestation_id expects a buffer of exactly
+    // KEY_ATTESTATION_APPLICATION_ID_MAX_SIZE bytes and returns the number of bytes written
+    // in the second pointer argument.
+    let status = unsafe { aaid_keystore_attestation_id(uid, buffer.as_mut_ptr(), &mut size) };
+    match status {
+        0 => Ok(buffer[0..size as usize].to_vec()),
+        status => Err(status),
+    }
+}
diff --git a/keystore2/src/crypto/certificate_utils.cpp b/keystore2/src/crypto/certificate_utils.cpp
index 500600f..4b0dca4 100644
--- a/keystore2/src/crypto/certificate_utils.cpp
+++ b/keystore2/src/crypto/certificate_utils.cpp
@@ -544,6 +544,9 @@
 
     bssl::UniquePtr<uint8_t> free_cert_buf(cert_buf);
     auto signature = sign(cert_buf, buf_len);
+    if (signature.empty()) {
+        return CertUtilsError::SignatureFailed;
+    }
 
     if (!ASN1_STRING_set(certificate->signature, signature.data(), signature.size())) {
         return CertUtilsError::BoringSsl;
diff --git a/keystore2/src/crypto/include/certificate_utils.h b/keystore2/src/crypto/include/certificate_utils.h
index 9d41eb8..1e80d80 100644
--- a/keystore2/src/crypto/include/certificate_utils.h
+++ b/keystore2/src/crypto/include/certificate_utils.h
@@ -52,6 +52,7 @@
         MemoryAllocation,
         InvalidArgument,
         UnexpectedNullPointer,
+        SignatureFailed,
     };
 
   private:
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index db30e07..785847d 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -42,7 +42,7 @@
 //! callbacks.
 
 use crate::db_utils::{self, SqlField};
-use crate::error::{Error as KsError, ResponseCode};
+use crate::error::{Error as KsError, ErrorCode, ResponseCode};
 use crate::impl_metadata; // This is in db_utils.rs
 use crate::key_parameter::{KeyParameter, Tag};
 use crate::permission::KeyPermSet;
@@ -1356,6 +1356,43 @@
         Ok(parameters)
     }
 
+    /// Decrements the usage count of a limited use key. This function first checks whether the
+    /// usage has been exhausted, if not, decreases the usage count. If the usage count reaches
+    /// zero, the key also gets marked unreferenced and scheduled for deletion.
+    /// Returns Ok(true) if the key was marked unreferenced as a hint to the garbage collector.
+    pub fn check_and_update_key_usage_count(&mut self, key_id: i64) -> Result<bool> {
+        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+            let limit: Option<i32> = tx
+                .query_row(
+                    "SELECT data FROM persistent.keyparameter WHERE keyentryid = ? AND tag = ?;",
+                    params![key_id, Tag::USAGE_COUNT_LIMIT.0],
+                    |row| row.get(0),
+                )
+                .optional()
+                .context("Trying to load usage count")?;
+
+            let limit = limit
+                .ok_or(KsError::Km(ErrorCode::INVALID_KEY_BLOB))
+                .context("The Key no longer exists. Key is exhausted.")?;
+
+            tx.execute(
+                "UPDATE persistent.keyparameter
+                 SET data = data - 1
+                 WHERE keyentryid = ? AND tag = ? AND data > 0;",
+                params![key_id, Tag::USAGE_COUNT_LIMIT.0],
+            )
+            .context("Failed to update key usage count.")?;
+
+            match limit {
+                1 => Self::mark_unreferenced(tx, key_id)
+                    .context("Trying to mark limited use key for deletion."),
+                0 => Err(KsError::Km(ErrorCode::INVALID_KEY_BLOB)).context("Key is exhausted."),
+                _ => Ok(false),
+            }
+        })
+        .context("In check_and_update_key_usage_count.")
+    }
+
     /// Load a key entry by the given key descriptor.
     /// It uses the `check_permission` callback to verify if the access is allowed
     /// given the key access tuple read from the database using `load_access_tuple`.
@@ -1779,7 +1816,7 @@
     };
     use crate::key_perm_set;
     use crate::permission::{KeyPerm, KeyPermSet};
-    use crate::test::utils::TempDir;
+    use keystore2_test_utils::TempDir;
     use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
         HardwareAuthToken::HardwareAuthToken,
         HardwareAuthenticatorType::HardwareAuthenticatorType as kmhw_authenticator_type,
@@ -2247,7 +2284,7 @@
     #[test]
     fn test_insert_and_load_full_keyentry_domain_app() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_domain_app")?
             .0;
         let (_key_guard, key_entry) = db
@@ -2264,7 +2301,7 @@
                 |_k, _av| Ok(()),
             )
             .unwrap();
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(
             KeyDescriptor {
@@ -2304,7 +2341,7 @@
     #[test]
     fn test_insert_and_load_full_keyentry_domain_selinux() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_domain_selinux")?
             .0;
         let (_key_guard, key_entry) = db
@@ -2321,7 +2358,7 @@
                 |_k, _av| Ok(()),
             )
             .unwrap();
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(
             KeyDescriptor {
@@ -2361,7 +2398,7 @@
     #[test]
     fn test_insert_and_load_full_keyentry_domain_key_id() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_domain_key_id")?
             .0;
         let (_, key_entry) = db
@@ -2374,7 +2411,7 @@
             )
             .unwrap();
 
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(
             KeyDescriptor { domain: Domain::KEY_ID, nspace: key_id, alias: None, blob: None },
@@ -2402,9 +2439,57 @@
     }
 
     #[test]
+    fn test_check_and_update_key_usage_count_with_limited_use_key() -> Result<()> {
+        let mut db = new_test_db()?;
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, Some(123))
+            .context("test_check_and_update_key_usage_count_with_limited_use_key")?
+            .0;
+        // Update the usage count of the limited use key.
+        db.check_and_update_key_usage_count(key_id)?;
+
+        let (_key_guard, key_entry) = db.load_key_entry(
+            KeyDescriptor { domain: Domain::KEY_ID, nspace: key_id, alias: None, blob: None },
+            KeyType::Client,
+            KeyEntryLoadBits::BOTH,
+            1,
+            |_k, _av| Ok(()),
+        )?;
+
+        // The usage count is decremented now.
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, Some(122)));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_check_and_update_key_usage_count_with_exhausted_limited_use_key() -> Result<()> {
+        let mut db = new_test_db()?;
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, Some(1))
+            .context("test_check_and_update_key_usage_count_with_exhausted_limited_use_key")?
+            .0;
+        // Update the usage count of the limited use key.
+        db.check_and_update_key_usage_count(key_id).expect(concat!(
+            "In test_check_and_update_key_usage_count_with_exhausted_limited_use_key: ",
+            "This should succeed."
+        ));
+
+        // Try to update the exhausted limited use key.
+        let e = db.check_and_update_key_usage_count(key_id).expect_err(concat!(
+            "In test_check_and_update_key_usage_count_with_exhausted_limited_use_key: ",
+            "This should fail."
+        ));
+        assert_eq!(
+            &KsError::Km(ErrorCode::INVALID_KEY_BLOB),
+            e.root_cause().downcast_ref::<KsError>().unwrap()
+        );
+
+        Ok(())
+    }
+
+    #[test]
     fn test_insert_and_load_full_keyentry_from_grant() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_from_grant")?
             .0;
 
@@ -2439,7 +2524,7 @@
             )
             .unwrap();
 
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(granted_key.clone(), KeyType::Client, 2, |_, _| Ok(())).unwrap();
 
@@ -2552,7 +2637,7 @@
             let temp_dir = Arc::new(TempDir::new("id_lock_test")?);
             let temp_dir_clone = temp_dir.clone();
             let mut db = KeystoreDB::new(temp_dir.path())?;
-            let key_id = make_test_key_entry(&mut db, Domain::APP, 33, KEY_LOCK_TEST_ALIAS)
+            let key_id = make_test_key_entry(&mut db, Domain::APP, 33, KEY_LOCK_TEST_ALIAS, None)
                 .context("test_insert_and_load_full_keyentry_domain_app")?
                 .0;
             let (_key_guard, key_entry) = db
@@ -2569,7 +2654,7 @@
                     |_k, _av| Ok(()),
                 )
                 .unwrap();
-            assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+            assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
             let state = Arc::new(AtomicU8::new(1));
             let state2 = state.clone();
 
@@ -2651,8 +2736,8 @@
         let list_o_keys: Vec<(i64, i64)> = LIST_O_ENTRIES
             .iter()
             .map(|(domain, ns, alias)| {
-                let entry =
-                    make_test_key_entry(&mut db, *domain, *ns, *alias).unwrap_or_else(|e| {
+                let entry = make_test_key_entry(&mut db, *domain, *ns, *alias, None)
+                    .unwrap_or_else(|e| {
                         panic!("Failed to insert {:?} {} {}. Error {:?}", domain, ns, alias, e)
                     });
                 (entry.id(), *ns)
@@ -2760,8 +2845,8 @@
     // Note: The parameters and SecurityLevel associations are nonsensical. This
     // collection is only used to check if the parameters are preserved as expected by the
     // database.
-    fn make_test_params() -> Vec<KeyParameter> {
-        vec![
+    fn make_test_params(max_usage_count: Option<i32>) -> Vec<KeyParameter> {
+        let mut params = vec![
             KeyParameter::new(KeyParameterValue::Invalid, SecurityLevel::TRUSTED_ENVIRONMENT),
             KeyParameter::new(
                 KeyParameterValue::KeyPurpose(KeyPurpose::SIGN),
@@ -2976,7 +3061,14 @@
                 KeyParameterValue::ConfirmationToken(vec![5u8, 5u8, 5u8, 5u8]),
                 SecurityLevel::TRUSTED_ENVIRONMENT,
             ),
-        ]
+        ];
+        if let Some(value) = max_usage_count {
+            params.push(KeyParameter::new(
+                KeyParameterValue::UsageCountLimit(value),
+                SecurityLevel::SOFTWARE,
+            ));
+        }
+        params
     }
 
     fn make_test_key_entry(
@@ -2984,12 +3076,16 @@
         domain: Domain,
         namespace: i64,
         alias: &str,
+        max_usage_count: Option<i32>,
     ) -> Result<KeyIdGuard> {
         let key_id = db.create_key_entry(domain, namespace)?;
         db.insert_blob(&key_id, SubComponentType::KEY_BLOB, TEST_KEY_BLOB)?;
         db.insert_blob(&key_id, SubComponentType::CERT, TEST_CERT_BLOB)?;
         db.insert_blob(&key_id, SubComponentType::CERT_CHAIN, TEST_CERT_CHAIN_BLOB)?;
-        db.insert_keyparameter(&key_id, &make_test_params())?;
+
+        let params = make_test_params(max_usage_count);
+        db.insert_keyparameter(&key_id, &params)?;
+
         let mut metadata = KeyMetaData::new();
         metadata.add(KeyMetaEntry::EncryptedBy(EncryptedBy::Password));
         metadata.add(KeyMetaEntry::Salt(vec![1, 2, 3]));
@@ -3000,7 +3096,9 @@
         Ok(key_id)
     }
 
-    fn make_test_key_entry_test_vector(key_id: i64) -> KeyEntry {
+    fn make_test_key_entry_test_vector(key_id: i64, max_usage_count: Option<i32>) -> KeyEntry {
+        let params = make_test_params(max_usage_count);
+
         let mut metadata = KeyMetaData::new();
         metadata.add(KeyMetaEntry::EncryptedBy(EncryptedBy::Password));
         metadata.add(KeyMetaEntry::Salt(vec![1, 2, 3]));
@@ -3013,7 +3111,7 @@
             cert: Some(TEST_CERT_BLOB.to_vec()),
             cert_chain: Some(TEST_CERT_CHAIN_BLOB.to_vec()),
             sec_level: SecurityLevel::TRUSTED_ENVIRONMENT,
-            parameters: make_test_params(),
+            parameters: params,
             metadata,
         }
     }
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 22b4bed..3fd2b19 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -14,10 +14,13 @@
 
 //! This is the Keystore 2.0 Enforcements module.
 // TODO: more description to follow.
-use crate::database::{AuthTokenEntry, MonotonicRawTime};
 use crate::error::{map_binder_status, Error, ErrorCode};
 use crate::globals::{get_timestamp_service, ASYNC_TASK, DB, ENFORCEMENTS};
 use crate::key_parameter::{KeyParameter, KeyParameterValue};
+use crate::{
+    database::{AuthTokenEntry, MonotonicRawTime},
+    gc::Gc,
+};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     Algorithm::Algorithm, ErrorCode::ErrorCode as Ec, HardwareAuthToken::HardwareAuthToken,
     HardwareAuthenticatorType::HardwareAuthenticatorType,
@@ -128,6 +131,8 @@
 #[derive(Debug)]
 pub struct AuthInfo {
     state: DeferredAuthState,
+    /// An optional key id required to update the usage count if the key usage is limited.
+    key_usage_limited: Option<i64>,
 }
 
 struct TokenReceiverMap {
@@ -251,14 +256,46 @@
         }
     }
 
+    /// This function is the authorization hook called before operation update.
+    /// It returns the auth tokens required by the operation to commence update.
+    pub fn before_update(&mut self) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
+        self.get_auth_tokens()
+    }
+
+    /// This function is the authorization hook called before operation finish.
+    /// It returns the auth tokens required by the operation to commence finish.
+    pub fn before_finish(&mut self) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
+        self.get_auth_tokens()
+    }
+
+    /// This function is the authorization hook called after finish succeeded.
+    /// As of this writing it checks if the key was a limited use key. If so it updates the
+    /// use counter of the key in the database. When the use counter is depleted, the key gets
+    /// marked for deletion and the garbage collector is notified.
+    pub fn after_finish(&self) -> Result<()> {
+        if let Some(key_id) = self.key_usage_limited {
+            // On the last successful use, the key gets deleted. In this case we
+            // have to notify the garbage collector.
+            let need_gc = DB
+                .with(|db| {
+                    db.borrow_mut()
+                        .check_and_update_key_usage_count(key_id)
+                        .context("Trying to update key usage count.")
+                })
+                .context("In after_finish.")?;
+            if need_gc {
+                Gc::notify_gc();
+            }
+        }
+        Ok(())
+    }
+
     /// This function returns the auth tokens as needed by the ongoing operation or fails
     /// with ErrorCode::KEY_USER_NOT_AUTHENTICATED. If this was called for the first time
     /// after a deferred authorization was requested by finalize_create_authorization, this
     /// function may block on the generation of a time stamp token. It then moves the
     /// tokens into the DeferredAuthState::Token state for future use.
-    pub fn get_auth_tokens(
-        &mut self,
-    ) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
+    fn get_auth_tokens(&mut self) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
         let deferred_tokens = if let DeferredAuthState::Waiting(ref auth_request) = self.state {
             let mut state = auth_request.lock().unwrap();
             Some(state.get_auth_tokens().context("In AuthInfo::get_auth_tokens.")?)
@@ -325,14 +362,18 @@
     pub fn authorize_create(
         &self,
         purpose: KeyPurpose,
-        key_params: Option<&[KeyParameter]>,
+        key_properties: Option<&(i64, Vec<KeyParameter>)>,
         op_params: &[KmKeyParameter],
         requires_timestamp: bool,
     ) -> Result<(Option<HardwareAuthToken>, AuthInfo)> {
-        let key_params = if let Some(k) = key_params {
-            k
-        } else {
-            return Ok((None, AuthInfo { state: DeferredAuthState::NoAuthRequired }));
+        let (key_id, key_params) = match key_properties {
+            Some((key_id, key_params)) => (*key_id, key_params),
+            None => {
+                return Ok((
+                    None,
+                    AuthInfo { state: DeferredAuthState::NoAuthRequired, key_usage_limited: None },
+                ))
+            }
         };
 
         match purpose {
@@ -377,6 +418,7 @@
         let mut key_time_out: Option<i64> = None;
         let mut allow_while_on_body = false;
         let mut unlocked_device_required = false;
+        let mut key_usage_limited: Option<i64> = None;
 
         // iterate through key parameters, recording information we need for authorization
         // enforcements later, or enforcing authorizations in place, where applicable
@@ -392,13 +434,10 @@
                     user_auth_type = Some(*a);
                 }
                 KeyParameterValue::KeyPurpose(p) => {
-                    // Note: if there can be multiple KeyPurpose key parameters (TODO: confirm this),
-                    // following check has the effect of key_params.contains(purpose)
+                    // The following check has the effect of key_params.contains(purpose)
                     // Also, authorizing purpose can not be completed here, if there can be multiple
-                    // key parameters for KeyPurpose
-                    if !key_purpose_authorized && *p == purpose {
-                        key_purpose_authorized = true;
-                    }
+                    // key parameters for KeyPurpose.
+                    key_purpose_authorized = key_purpose_authorized || *p == purpose;
                 }
                 KeyParameterValue::CallerNonce => {
                     caller_nonce_allowed = true;
@@ -437,6 +476,12 @@
                 KeyParameterValue::AllowWhileOnBody => {
                     allow_while_on_body = true;
                 }
+                KeyParameterValue::UsageCountLimit(_) => {
+                    // We don't examine the limit here because this is enforced on finish.
+                    // Instead, we store the key_id so that finish can look up the key
+                    // in the database again and check and update the counter.
+                    key_usage_limited = Some(key_id);
+                }
                 // NOTE: as per offline discussion, sanitizing key parameters and rejecting
                 // create operation if any non-allowed tags are present, is not done in
                 // authorize_create (unlike in legacy keystore where AuthorizeBegin is rejected if
@@ -491,7 +536,10 @@
         }
 
         if !unlocked_device_required && no_auth_required {
-            return Ok((None, AuthInfo { state: DeferredAuthState::NoAuthRequired }));
+            return Ok((
+                None,
+                AuthInfo { state: DeferredAuthState::NoAuthRequired, key_usage_limited },
+            ));
         }
 
         let has_sids = !user_secure_ids.is_empty();
@@ -565,7 +613,7 @@
             (None, _, true) => (None, DeferredAuthState::OpAuthRequired),
             (None, _, false) => (None, DeferredAuthState::NoAuthRequired),
         })
-        .map(|(hat, state)| (hat, AuthInfo { state }))
+        .map(|(hat, state)| (hat, AuthInfo { state, key_usage_limited }))
     }
 
     fn find_auth_token<F>(p: F) -> Result<Option<(AuthTokenEntry, MonotonicRawTime)>>
diff --git a/keystore2/src/key_parameter.rs b/keystore2/src/key_parameter.rs
index 71a17fe..93de6f2 100644
--- a/keystore2/src/key_parameter.rs
+++ b/keystore2/src/key_parameter.rs
@@ -840,6 +840,9 @@
     /// Maximum number of times that a key may be used between system reboots
     #[key_param(tag = MAX_USES_PER_BOOT, field = Integer)]
     MaxUsesPerBoot(i32),
+    /// The number of times that a limited use key can be used
+    #[key_param(tag = USAGE_COUNT_LIMIT, field = Integer)]
+    UsageCountLimit(i32),
     /// ID of the Android user that is permitted to use the key
     #[key_param(tag = USER_ID, field = Integer)]
     UserID(i32),
diff --git a/keystore2/src/km_compat/Android.bp b/keystore2/src/km_compat/Android.bp
index a5da5a6..126aeff 100644
--- a/keystore2/src/km_compat/Android.bp
+++ b/keystore2/src/km_compat/Android.bp
@@ -18,7 +18,7 @@
     srcs: ["lib.rs"],
 
     rustlibs: [
-        "android.hardware.security.keymint-rust",
+        "android.hardware.security.keymint-V1-rust",
         "android.security.compat-rust",
     ],
     shared_libs: [
@@ -33,7 +33,7 @@
     test_suites: ["general-tests"],
     auto_gen_config: true,
     rustlibs: [
-        "android.hardware.security.keymint-rust",
+        "android.hardware.security.keymint-V1-rust",
         "android.security.compat-rust",
     ],
     shared_libs: [
diff --git a/keystore2/src/km_compat/km_compat.cpp b/keystore2/src/km_compat/km_compat.cpp
index d965922..a27bfd1 100644
--- a/keystore2/src/km_compat/km_compat.cpp
+++ b/keystore2/src/km_compat/km_compat.cpp
@@ -19,6 +19,7 @@
 #include "km_compat_type_conversion.h"
 #include <aidl/android/hardware/security/keymint/Algorithm.h>
 #include <aidl/android/hardware/security/keymint/Digest.h>
+#include <aidl/android/hardware/security/keymint/ErrorCode.h>
 #include <aidl/android/hardware/security/keymint/PaddingMode.h>
 #include <aidl/android/system/keystore2/ResponseCode.h>
 #include <android-base/logging.h>
@@ -59,6 +60,14 @@
     return ScopedAStatus::fromServiceSpecificError(static_cast<int32_t>(result));
 }
 
+static V4_0_ErrorCode toErrorCode(const ScopedAStatus& status) {
+    if (status.getExceptionCode() == EX_SERVICE_SPECIFIC) {
+        return static_cast<V4_0_ErrorCode>(status.getServiceSpecificError());
+    } else {
+        return V4_0_ErrorCode::UNKNOWN_ERROR;
+    }
+}
+
 static std::vector<V4_0::KeyParameter>
 convertKeyParametersToLegacy(const std::vector<KeyParameter>& kps) {
     std::vector<V4_0::KeyParameter> legacyKps(kps.size());
@@ -660,22 +669,23 @@
             BeginResult beginResult;
             auto error = begin(KeyPurpose::SIGN, keyBlob, kps, HardwareAuthToken(), &beginResult);
             if (!error.isOk()) {
-                errorCode = static_cast<V4_0_ErrorCode>(error.getServiceSpecificError());
+                errorCode = toErrorCode(error);
                 return std::vector<uint8_t>();
             }
             std::optional<KeyParameterArray> outParams;
             std::optional<ByteArray> outByte;
             int32_t status;
-            beginResult.operation->update(std::nullopt, dataVec, std::nullopt, std::nullopt,
-                                          &outParams, &outByte, &status);
-            if (!status) {
+            error = beginResult.operation->update(std::nullopt, dataVec, std::nullopt, std::nullopt,
+                                                  &outParams, &outByte, &status);
+            if (!error.isOk()) {
+                errorCode = toErrorCode(error);
                 return std::vector<uint8_t>();
             }
             std::vector<uint8_t> result;
             error = beginResult.operation->finish(std::nullopt, std::nullopt, std::nullopt,
                                                   std::nullopt, std::nullopt, &outParams, &result);
             if (!error.isOk()) {
-                errorCode = static_cast<V4_0_ErrorCode>(error.getServiceSpecificError());
+                errorCode = toErrorCode(error);
                 return std::vector<uint8_t>();
             }
             return result;
@@ -725,6 +735,7 @@
                 }
             });
         if (!result.isOk()) {
+            LOG(ERROR) << __func__ << ": Call to attestKey failed.";
             return V4_0_ErrorCode::UNKNOWN_ERROR;
         }
         if (errorCode != V4_0_ErrorCode::OK) {
@@ -743,6 +754,7 @@
     // setIssuer
     auto error = keystore::setIssuer(&*cert, &*cert, false);
     if (error) {
+        LOG(ERROR) << __func__ << ": Set issuer failed.";
         return V4_0_ErrorCode::UNKNOWN_ERROR;
     }
 
@@ -755,26 +767,27 @@
             return false;
         }) != keyParams.end();
     auto noAuthRequired = containsParam(keyParams, KMV1::TAG_NO_AUTH_REQUIRED);
-    if (canSelfSign && noAuthRequired) {
-        auto errorCode = signCertificate(keyParams, keyBlob, &*cert);
-        if (errorCode.has_value()) {
-            return errorCode.value();
-        }
-    } else {
+    // If we cannot sign because of purpose or authorization requirement,
+    if (!(canSelfSign && noAuthRequired)
+        // or if self signing fails for any other reason,
+        || signCertificate(keyParams, keyBlob, &*cert).has_value()) {
+        // we sign with ephemeral key.
         keystore::EVP_PKEY_CTX_Ptr pkey_ctx(EVP_PKEY_CTX_new_id(EVP_PKEY_EC, NULL));
         EVP_PKEY_keygen_init(pkey_ctx.get());
         EVP_PKEY_CTX_set_ec_paramgen_curve_nid(pkey_ctx.get(), NID_X9_62_prime256v1);
         EVP_PKEY* pkey_ptr = nullptr;
         EVP_PKEY_keygen(pkey_ctx.get(), &pkey_ptr);
         error = keystore::signCert(&*cert, pkey_ptr);
-    }
-    if (error) {
-        return V4_0_ErrorCode::UNKNOWN_ERROR;
+        if (error) {
+            LOG(ERROR) << __func__ << ": signCert failed.";
+            return V4_0_ErrorCode::UNKNOWN_ERROR;
+        }
     }
 
     // encodeCert
     auto encodedCertOrError = keystore::encodeCert(&*cert);
     if (std::holds_alternative<keystore::CertUtilsError>(encodedCertOrError)) {
+        LOG(ERROR) << __func__ << ": encodeCert failed.";
         return V4_0_ErrorCode::UNKNOWN_ERROR;
     }
 
diff --git a/keystore2/src/km_compat/km_compat_type_conversion.h b/keystore2/src/km_compat/km_compat_type_conversion.h
index 97c61ce..db9d2a4 100644
--- a/keystore2/src/km_compat/km_compat_type_conversion.h
+++ b/keystore2/src/km_compat/km_compat_type_conversion.h
@@ -23,7 +23,7 @@
 namespace V4_1 = ::android::hardware::keymaster::V4_1;
 namespace KMV1 = ::aidl::android::hardware::security::keymint;
 
-static V4_0::KeyPurpose convert(KMV1::KeyPurpose p) {
+static std::optional<V4_0::KeyPurpose> convert(KMV1::KeyPurpose p) {
     switch (p) {
     case KMV1::KeyPurpose::ENCRYPT:
         return V4_0::KeyPurpose::ENCRYPT;
@@ -35,6 +35,9 @@
         return V4_0::KeyPurpose::VERIFY;
     case KMV1::KeyPurpose::WRAP_KEY:
         return V4_0::KeyPurpose::WRAP_KEY;
+    default:
+        // Can end up here because KeyMint may have KeyPurpose values not in KM4.
+        return {};
     }
 }
 
@@ -288,7 +291,10 @@
         break;
     case KMV1::Tag::PURPOSE:
         if (auto v = KMV1::authorizationValue(KMV1::TAG_PURPOSE, kp)) {
-            return V4_0::makeKeyParameter(V4_0::TAG_PURPOSE, convert(v->get()));
+            std::optional<V4_0::KeyPurpose> purpose = convert(v->get());
+            if (purpose) {
+                return V4_0::makeKeyParameter(V4_0::TAG_PURPOSE, purpose.value());
+            }
         }
         break;
     case KMV1::Tag::ALGORITHM:
@@ -389,6 +395,9 @@
             return V4_0::makeKeyParameter(V4_0::TAG_MAX_USES_PER_BOOT, v->get());
         }
         break;
+    case KMV1::Tag::USAGE_COUNT_LIMIT:
+        // Does not exist in KM < KeyMint 1.0.
+        break;
     case KMV1::Tag::USER_ID:
         if (auto v = KMV1::authorizationValue(KMV1::TAG_USER_ID, kp)) {
             return V4_0::makeKeyParameter(V4_0::TAG_USER_ID, v->get());
diff --git a/keystore2/src/legacy_blob.rs b/keystore2/src/legacy_blob.rs
index 34a0eca..230a82c 100644
--- a/keystore2/src/legacy_blob.rs
+++ b/keystore2/src/legacy_blob.rs
@@ -775,7 +775,7 @@
     mod legacy_blob_test_vectors;
     use crate::error;
     use crate::legacy_blob::test::legacy_blob_test_vectors::*;
-    use crate::test::utils::TempDir;
+    use keystore2_test_utils::TempDir;
 
     #[test]
     fn decode_encode_alias_test() {
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 0475d6f..811db91 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -34,8 +34,3 @@
 mod db_utils;
 mod gc;
 mod super_key;
-
-#[cfg(test)]
-mod test {
-    pub mod utils;
-}
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 30e6d55..829987d 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -341,7 +341,7 @@
             .auth_info
             .lock()
             .unwrap()
-            .get_auth_tokens()
+            .before_update()
             .context("In update_aad: Trying to get auth tokens.")?;
 
         self.update_outcome(
@@ -377,7 +377,7 @@
             .auth_info
             .lock()
             .unwrap()
-            .get_auth_tokens()
+            .before_update()
             .context("In update: Trying to get auth tokens.")?;
 
         self.update_outcome(
@@ -423,7 +423,7 @@
             .auth_info
             .lock()
             .unwrap()
-            .get_auth_tokens()
+            .before_finish()
             .context("In finish: Trying to get auth tokens.")?;
 
         let output = self
@@ -440,6 +440,8 @@
             )
             .context("In finish: KeyMint::finish failed.")?;
 
+        self.auth_info.lock().unwrap().after_finish().context("In finish.")?;
+
         // At this point the operation concluded successfully.
         *outcome = Outcome::Success;
 
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index d8787bd..12b75bf 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -34,7 +34,7 @@
 use crate::globals::ENFORCEMENTS;
 use crate::key_parameter::KeyParameter as KsKeyParam;
 use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
-use crate::utils::{check_key_permission, Asp};
+use crate::utils::{check_key_permission, uid_to_android_user, Asp};
 use crate::{database::KeyIdGuard, globals::DB};
 use crate::{
     database::{DateTime, KeyMetaData, KeyMetaEntry, KeyType},
@@ -48,9 +48,8 @@
 use crate::{
     error::{self, map_km_error, map_or_log_err, Error, ErrorCode},
     utils::key_characteristics_to_internal,
-    utils::uid_to_android_user,
 };
-use anyhow::{Context, Result};
+use anyhow::{anyhow, Context, Result};
 use binder::{IBinder, Interface, ThreadState};
 
 /// Implementation of the IKeystoreSecurityLevel Interface.
@@ -173,7 +172,7 @@
         // so that we can use it by reference like the blob provided by the key descriptor.
         // Otherwise, we would have to clone the blob from the key descriptor.
         let scoping_blob: Vec<u8>;
-        let (km_blob, key_id_guard, key_parameters) = match key.domain {
+        let (km_blob, key_properties, key_id_guard) = match key.domain {
             Domain::BLOB => {
                 check_key_permission(KeyPerm::use_(), key, &None)
                     .context("In create_operation: checking use permission for Domain::BLOB.")?;
@@ -212,7 +211,11 @@
                         ))
                     }
                 };
-                (&scoping_blob, Some(key_id_guard), Some(key_entry.into_key_parameters()))
+                (
+                    &scoping_blob,
+                    Some((key_id_guard.id(), key_entry.into_key_parameters())),
+                    Some(key_id_guard),
+                )
             }
         };
 
@@ -229,8 +232,8 @@
         let (immediate_hat, mut auth_info) = ENFORCEMENTS
             .authorize_create(
                 purpose,
-                key_parameters.as_deref(),
-                operation_parameters,
+                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.
@@ -293,6 +296,20 @@
         })
     }
 
+    fn add_attestation_parameters(uid: u32, params: &[KeyParameter]) -> Result<Vec<KeyParameter>> {
+        let mut result = params.to_vec();
+        if params.iter().any(|kp| kp.tag == Tag::ATTESTATION_CHALLENGE) {
+            let aaid = keystore2_aaid::get_aaid(uid).map_err(|e| {
+                anyhow!(format!("In add_attestation_parameters: get_aaid returned status {}.", e))
+            })?;
+            result.push(KeyParameter {
+                tag: Tag::ATTESTATION_APPLICATION_ID,
+                value: KeyParameterValue::Blob(aaid),
+            });
+        }
+        Ok(result)
+    }
+
     fn generate_key(
         &self,
         key: &KeyDescriptor,
@@ -320,9 +337,14 @@
         // generate_key requires the rebind permission.
         check_key_permission(KeyPerm::rebind(), &key, &None).context("In generate_key.")?;
 
+        let params = Self::add_attestation_parameters(caller_uid, params)
+            .context("In generate_key: Trying to get aaid.")?;
+
         let km_dev: Box<dyn IKeyMintDevice> = self.keymint.get_interface()?;
-        map_km_error(km_dev.addRngEntropy(entropy))?;
-        let creation_result = map_km_error(km_dev.generateKey(&params))?;
+        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))
+            .context("In generate_key: While generating Key")?;
 
         let user_id = uid_to_android_user(caller_uid);
         self.store_new_key(key, creation_result, user_id).context("In generate_key.")
@@ -355,6 +377,9 @@
         // import_key requires the rebind permission.
         check_key_permission(KeyPerm::rebind(), &key, &None).context("In import_key.")?;
 
+        let params = Self::add_attestation_parameters(caller_uid, params)
+            .context("In import_key: Trying to get aaid.")?;
+
         let format = params
             .iter()
             .find(|p| p.tag == Tag::ALGORITHM)
@@ -371,8 +396,10 @@
             })
             .context("In import_key.")?;
 
-        let km_dev: Box<dyn IKeyMintDevice> = self.keymint.get_interface()?;
-        let creation_result = map_km_error(km_dev.importKey(&params, format, key_data))?;
+        let km_dev: Box<dyn IKeyMintDevice> =
+            self.keymint.get_interface().context("In import_key: Trying to get the KM device")?;
+        let creation_result = map_km_error(km_dev.importKey(&params, format, key_data))
+            .context("In import_key: Trying to call importKey")?;
 
         let user_id = uid_to_android_user(caller_uid);
         self.store_new_key(key, creation_result, user_id).context("In import_key.")
@@ -386,7 +413,7 @@
         params: &[KeyParameter],
         authenticators: &[AuthenticatorSpec],
     ) -> Result<KeyMetadata> {
-        if key.domain != Domain::BLOB && key.alias.is_none() {
+        if !(key.domain == Domain::BLOB && key.alias.is_some()) {
             return Err(error::Error::Km(ErrorCode::INVALID_ARGUMENT))
                 .context("In import_wrapped_key: Alias must be specified.");
         }
@@ -466,26 +493,29 @@
         let masking_key = masking_key.unwrap_or(ZERO_BLOB_32);
 
         let km_dev: Box<dyn IKeyMintDevice> = self.keymint.get_interface()?;
-        let (creation_result, _) = self.upgrade_keyblob_if_required_with(
-            &*km_dev,
-            Some(wrapping_key_id_guard),
-            wrapping_key_blob,
-            &[],
-            |wrapping_blob| {
-                let creation_result = map_km_error(km_dev.importWrappedKey(
-                    wrapped_data,
-                    wrapping_key_blob,
-                    masking_key,
-                    &params,
-                    pw_sid,
-                    fp_sid,
-                ))?;
-                Ok(creation_result)
-            },
-        )?;
+        let (creation_result, _) = self
+            .upgrade_keyblob_if_required_with(
+                &*km_dev,
+                Some(wrapping_key_id_guard),
+                wrapping_key_blob,
+                &[],
+                |wrapping_blob| {
+                    let creation_result = map_km_error(km_dev.importWrappedKey(
+                        wrapped_data,
+                        wrapping_key_blob,
+                        masking_key,
+                        &params,
+                        pw_sid,
+                        fp_sid,
+                    ))?;
+                    Ok(creation_result)
+                },
+            )
+            .context("In import_wrapped_key.")?;
 
         let user_id = uid_to_android_user(caller_uid);
-        self.store_new_key(key, creation_result, user_id).context("In import_wrapped_key.")
+        self.store_new_key(key, creation_result, user_id)
+            .context("In import_wrapped_key: Trying to store the new key.")
     }
 
     fn upgrade_keyblob_if_required_with<T, F>(
diff --git a/keystore2/src/test/utils.rs b/keystore2/test_utils/lib.rs
similarity index 67%
rename from keystore2/src/test/utils.rs
rename to keystore2/test_utils/lib.rs
index 8c93859..627af20 100644
--- a/keystore2/src/test/utils.rs
+++ b/keystore2/test_utils/lib.rs
@@ -12,11 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+//! Implements TempDir which aids in creating an cleaning up temporary directories for testing.
+
 use std::fs::{create_dir, remove_dir_all};
 use std::io::ErrorKind;
 use std::path::{Path, PathBuf};
 use std::{env::temp_dir, ops::Deref};
 
+/// Represents the lifecycle of a temporary directory for testing.
 #[derive(Debug)]
 pub struct TempDir {
     path: std::path::PathBuf,
@@ -24,6 +27,11 @@
 }
 
 impl TempDir {
+    /// Creates a temporary directory with a name of the form <prefix>_NNNNN where NNNNN is a zero
+    /// padded random number with 5 figures. The prefix must not contain file system separators.
+    /// The location of the directory cannot be chosen.
+    /// The directory with all of its content is removed from the file system when the resulting
+    /// object gets dropped.
     pub fn new(prefix: &str) -> std::io::Result<Self> {
         let tmp = loop {
             let mut tmp = temp_dir();
@@ -40,10 +48,20 @@
         Ok(Self { path: tmp, do_drop: true })
     }
 
+    /// Returns the absolute path of the temporary directory.
     pub fn path(&self) -> &Path {
         &self.path
     }
 
+    /// Returns a path builder for convenient extension of the path.
+    ///
+    /// ## Example:
+    ///
+    /// ```
+    /// let tdir = TempDir::new("my_test")?;
+    /// let temp_foo_bar = tdir.build().push("foo").push("bar");
+    /// ```
+    /// `temp_foo_bar` derefs to a Path that represents "<tdir.path()>/foo/bar"
     pub fn build(&self) -> PathBuilder {
         PathBuilder(self.path.clone())
     }
@@ -66,9 +84,11 @@
     }
 }
 
+/// Allows for convenient building of paths from a TempDir. See TempDir.build() for more details.
 pub struct PathBuilder(PathBuf);
 
 impl PathBuilder {
+    /// Adds another segment to the end of the path. Consumes, modifies and returns self.
     pub fn push(mut self, segment: &str) -> Self {
         self.0.push(segment);
         self