diff --git a/diced/open_dice_cbor/lib.rs b/diced/open_dice_cbor/lib.rs
index ffb8a48..78f5c81 100644
--- a/diced/open_dice_cbor/lib.rs
+++ b/diced/open_dice_cbor/lib.rs
@@ -24,7 +24,7 @@
 //!     code_hash: [3u8, dice::HASH_SIZE],
 //!     config: dice::ConfigOwned::Descriptor("My descriptor".as_bytes().to_vec()),
 //!     authority_hash: [0u8, dice::HASH_SIZE],
-//!     mode: dice::Mode::Normal,
+//!     mode: dice::DiceMode::kDiceModeNormal,
 //!     hidden: [0u8, dice::HIDDEN_SIZE],
 //! };
 //! let (cdi_attest, cdi_seal, cert_chain) = context
@@ -33,20 +33,17 @@
 
 use keystore2_crypto::{zvec, ZVec};
 use open_dice_bcc_bindgen::BccMainFlow;
+pub use open_dice_cbor_bindgen::DiceMode;
 use open_dice_cbor_bindgen::{
     DiceConfigType, DiceDeriveCdiCertificateId, DiceDeriveCdiPrivateKeySeed,
     DiceGenerateCertificate, DiceHash, DiceInputValues, DiceKdf, DiceKeypairFromSeed, DiceMainFlow,
-    DiceMode, DiceResult, DiceSign, DiceVerify, DICE_CDI_SIZE, DICE_HASH_SIZE, DICE_HIDDEN_SIZE,
+    DiceResult, DiceSign, DiceVerify, DICE_CDI_SIZE, DICE_HASH_SIZE, DICE_HIDDEN_SIZE,
     DICE_ID_SIZE, DICE_INLINE_CONFIG_SIZE, DICE_PRIVATE_KEY_SEED_SIZE, DICE_PRIVATE_KEY_SIZE,
     DICE_PUBLIC_KEY_SIZE, DICE_SIGNATURE_SIZE,
 };
 use open_dice_cbor_bindgen::{
     DiceConfigType_kDiceConfigTypeDescriptor as DICE_CONFIG_TYPE_DESCRIPTOR,
     DiceConfigType_kDiceConfigTypeInline as DICE_CONFIG_TYPE_INLINE,
-    DiceMode_kDiceModeDebug as DICE_MODE_DEBUG,
-    DiceMode_kDiceModeMaintenance as DICE_MODE_RECOVERY,
-    DiceMode_kDiceModeNormal as DICE_MODE_NORMAL,
-    DiceMode_kDiceModeNotInitialized as DICE_MODE_NOT_CONFIGURED,
     DiceResult_kDiceResultBufferTooSmall as DICE_RESULT_BUFFER_TOO_SMALL,
     DiceResult_kDiceResultInvalidInput as DICE_RESULT_INVALID_INPUT,
     DiceResult_kDiceResultOk as DICE_RESULT_OK,
@@ -175,31 +172,6 @@
     }
 }
 
-/// DICE modes as defined here:
-/// https://pigweed.googlesource.com/open-dice/+/refs/heads/main/docs/specification.md#mode-value-details
-#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
-pub enum Mode {
-    /// See documentation linked above.
-    NotConfigured = 0,
-    /// See documentation linked above.
-    Normal = 1,
-    /// See documentation linked above.
-    Debug = 2,
-    /// See documentation linked above.
-    Recovery = 3,
-}
-
-impl Mode {
-    fn get_internal(&self) -> DiceMode {
-        match self {
-            Self::NotConfigured => DICE_MODE_NOT_CONFIGURED,
-            Self::Normal => DICE_MODE_NORMAL,
-            Self::Debug => DICE_MODE_DEBUG,
-            Self::Recovery => DICE_MODE_RECOVERY,
-        }
-    }
-}
-
 /// This trait allows API users to supply DICE input values without copying.
 pub trait InputValues {
     /// Returns the code hash.
@@ -211,7 +183,7 @@
     /// Returns the authority descriptor.
     fn authority_descriptor(&self) -> Option<&[u8]>;
     /// Returns the mode.
-    fn mode(&self) -> Mode;
+    fn mode(&self) -> DiceMode;
     /// Returns the hidden value.
     fn hidden(&self) -> &[u8; HIDDEN_SIZE];
 }
@@ -222,7 +194,7 @@
     config: ConfigOwned,
     authority_hash: [u8; HASH_SIZE],
     authority_descriptor: Option<Vec<u8>>,
-    mode: Mode,
+    mode: DiceMode,
     hidden: [u8; HIDDEN_SIZE],
 }
 
@@ -233,7 +205,7 @@
         config: Config,
         authority_hash: [u8; HASH_SIZE],
         authority_descriptor: Option<Vec<u8>>,
-        mode: Mode,
+        mode: DiceMode,
         hidden: [u8; HIDDEN_SIZE],
     ) -> Self {
         Self {
@@ -263,7 +235,7 @@
     fn authority_descriptor(&self) -> Option<&[u8]> {
         self.authority_descriptor.as_deref()
     }
-    fn mode(&self) -> Mode {
+    fn mode(&self) -> DiceMode {
         self.mode
     }
     fn hidden(&self) -> &[u8; HIDDEN_SIZE] {
@@ -288,7 +260,7 @@
             .authority_descriptor()
             .map_or_else(std::ptr::null, <[u8]>::as_ptr),
         authority_descriptor_size: input_values.authority_descriptor().map_or(0, <[u8]>::len),
-        mode: input_values.mode().get_internal(),
+        mode: input_values.mode(),
         hidden: *input_values.hidden(),
     };
 
diff --git a/diced/src/utils.rs b/diced/src/utils.rs
index 03e8969..3cb4ba2 100644
--- a/diced/src/utils.rs
+++ b/diced/src/utils.rs
@@ -19,7 +19,7 @@
     Mode::Mode as BinderMode,
 };
 use anyhow::{Context, Result};
-use dice::ContextImpl;
+use dice::{ContextImpl, DiceMode};
 use diced_open_dice_cbor as dice;
 use keystore2_crypto::ZVec;
 use std::convert::TryInto;
@@ -63,13 +63,13 @@
         self.0.authorityDescriptor.as_deref()
     }
 
-    fn mode(&self) -> dice::Mode {
+    fn mode(&self) -> DiceMode {
         match self.0.mode {
-            BinderMode::NOT_INITIALIZED => dice::Mode::NotConfigured,
-            BinderMode::NORMAL => dice::Mode::Normal,
-            BinderMode::DEBUG => dice::Mode::Debug,
-            BinderMode::RECOVERY => dice::Mode::Recovery,
-            _ => dice::Mode::NotConfigured,
+            BinderMode::NOT_INITIALIZED => DiceMode::kDiceModeNotInitialized,
+            BinderMode::NORMAL => DiceMode::kDiceModeNormal,
+            BinderMode::DEBUG => DiceMode::kDiceModeDebug,
+            BinderMode::RECOVERY => DiceMode::kDiceModeMaintenance,
+            _ => DiceMode::kDiceModeNotInitialized,
         }
     }
 
@@ -256,18 +256,17 @@
     pub fn encode_header(t: u8, n: u64, buffer: &mut dyn Write) -> Result<()> {
         match n {
             n if n < 24 => {
-                let written = buffer
-                    .write(&u8::to_be_bytes(((t as u8) << 5) | (n as u8 & 0x1F)))
-                    .with_context(|| {
-                    format!("In encode_header: Failed to write header ({}, {})", t, n)
-                })?;
+                let written =
+                    buffer.write(&u8::to_be_bytes((t << 5) | (n as u8 & 0x1F))).with_context(
+                        || format!("In encode_header: Failed to write header ({}, {})", t, n),
+                    )?;
                 if written != 1 {
                     return Err(anyhow!("In encode_header: Buffer to small. ({}, {})", t, n));
                 }
             }
             n if n <= 0xFF => {
                 let written =
-                    buffer.write(&u8::to_be_bytes(((t as u8) << 5) | (24u8 & 0x1F))).with_context(
+                    buffer.write(&u8::to_be_bytes((t << 5) | (24u8 & 0x1F))).with_context(
                         || format!("In encode_header: Failed to write header ({}, {})", t, n),
                     )?;
                 if written != 1 {
@@ -286,7 +285,7 @@
             }
             n if n <= 0xFFFF => {
                 let written =
-                    buffer.write(&u8::to_be_bytes(((t as u8) << 5) | (25u8 & 0x1F))).with_context(
+                    buffer.write(&u8::to_be_bytes((t << 5) | (25u8 & 0x1F))).with_context(
                         || format!("In encode_header: Failed to write header ({}, {})", t, n),
                     )?;
                 if written != 1 {
@@ -305,7 +304,7 @@
             }
             n if n <= 0xFFFFFFFF => {
                 let written =
-                    buffer.write(&u8::to_be_bytes(((t as u8) << 5) | (26u8 & 0x1F))).with_context(
+                    buffer.write(&u8::to_be_bytes((t << 5) | (26u8 & 0x1F))).with_context(
                         || format!("In encode_header: Failed to write header ({}, {})", t, n),
                     )?;
                 if written != 1 {
@@ -324,13 +323,13 @@
             }
             n => {
                 let written =
-                    buffer.write(&u8::to_be_bytes(((t as u8) << 5) | (27u8 & 0x1F))).with_context(
+                    buffer.write(&u8::to_be_bytes((t << 5) | (27u8 & 0x1F))).with_context(
                         || format!("In encode_header: Failed to write header ({}, {})", t, n),
                     )?;
                 if written != 1 {
                     return Err(anyhow!("In encode_header: Buffer to small. ({}, {})", t, n));
                 }
-                let written = buffer.write(&u64::to_be_bytes(n as u64)).with_context(|| {
+                let written = buffer.write(&u64::to_be_bytes(n)).with_context(|| {
                     format!("In encode_header: Failed to write size ({}, {})", t, n)
                 })?;
                 if written != 8 {
diff --git a/identity/Android.bp b/identity/Android.bp
index 7b81685..4f203e6 100644
--- a/identity/Android.bp
+++ b/identity/Android.bp
@@ -47,7 +47,7 @@
         "android.hardware.identity-support-lib",
         "android.hardware.keymaster@4.0",
         "android.security.authorization-ndk",
-	"android.security.remoteprovisioning-cpp",
+        "android.security.remoteprovisioning-cpp",
         "libbase",
         "libbinder",
         "libbinder_ndk",
@@ -59,7 +59,6 @@
         "libutils",
         "libutilscallstack",
         "libvintf",
-	"server_configurable_flags",
     ],
     static_libs: [
         "android.hardware.security.rkp-V3-cpp",
diff --git a/identity/CredentialData.cpp b/identity/CredentialData.cpp
index fb08333..1bf1527 100644
--- a/identity/CredentialData.cpp
+++ b/identity/CredentialData.cpp
@@ -581,13 +581,17 @@
 
     vector<vector<uint8_t>> keysNeedingCert;
 
-    int64_t nowMilliSeconds =
-        std::chrono::system_clock::to_time_t(std::chrono::system_clock::now()) * 1000;
+    time_t now = std::chrono::system_clock::to_time_t(std::chrono::system_clock::now());
+    int64_t nowMilliseconds;
+    if (__builtin_mul_overflow(int64_t(now), int64_t(1000), &nowMilliseconds)) {
+        LOG(ERROR) << "Overflow converting " << now << " to milliseconds";
+        return {};
+    }
 
     for (AuthKeyData& data : authKeyDatas_) {
         bool keyExceedUseCount = (data.useCount >= maxUsesPerKey_);
         int64_t expirationDateAdjusted = data.expirationDateMillisSinceEpoch - minValidTimeMillis_;
-        bool keyBeyondAdjustedExpirationDate = (nowMilliSeconds > expirationDateAdjusted);
+        bool keyBeyondAdjustedExpirationDate = (nowMilliseconds > expirationDateAdjusted);
         bool newKeyNeeded =
             (data.certificate.size() == 0) || keyExceedUseCount || keyBeyondAdjustedExpirationDate;
         bool certificationPending = (data.pendingCertificate.size() > 0);
diff --git a/identity/CredentialStore.cpp b/identity/CredentialStore.cpp
index eb9bdb6..fea4df9 100644
--- a/identity/CredentialStore.cpp
+++ b/identity/CredentialStore.cpp
@@ -20,18 +20,19 @@
 #include <optional>
 
 #include <android-base/logging.h>
+#include <android-base/properties.h>
 #include <android/hardware/security/keymint/IRemotelyProvisionedComponent.h>
 #include <android/hardware/security/keymint/RpcHardwareInfo.h>
 #include <android/security/remoteprovisioning/IRemotelyProvisionedKeyPool.h>
 #include <android/security/remoteprovisioning/RemotelyProvisionedKey.h>
 #include <binder/IPCThreadState.h>
 #include <binder/IServiceManager.h>
-#include <server_configurable_flags/get_flags.h>
 #include <vintf/VintfObject.h>
 
 #include "Credential.h"
 #include "CredentialData.h"
 #include "CredentialStore.h"
+#include "RemotelyProvisionedKey.h"
 #include "Session.h"
 #include "Util.h"
 #include "WritableCredential.h"
@@ -45,10 +46,8 @@
 using ::android::security::rkp::IRemoteProvisioning;
 
 bool useRkpd() {
-    std::string useRkpdFlagValue = server_configurable_flags::GetServerConfigurableFlag(
-        "remote_key_provisioning_native", "enable_rkpd",
-        /*default_value=*/"false");
-    return useRkpdFlagValue == "true";
+    return android::base::GetBoolProperty("remote_provisioning.enable_rkpd",
+                                          /*default_value=*/false);
 }
 
 }  // namespace
@@ -70,31 +69,14 @@
             LOG(ERROR) << "Error getting remotely provisioned component: " << status;
             return false;
         }
-        useRkpd_ = useRkpd();
-
-        if (useRkpd_) {
-            uid_t callingUid = android::IPCThreadState::self()->getCallingUid();
-            auto rpcKeyFuture = getRpcKeyFuture(rpc_, callingUid);
-            if (!rpcKeyFuture) {
-                LOG(ERROR) << "Error in getRpcKeyFuture()";
-                return false;
-            }
-            rpcKeyFuture_ = std::move(*rpcKeyFuture);
-        } else {
-            keyPool_ = android::waitForService<IRemotelyProvisionedKeyPool>(
-                IRemotelyProvisionedKeyPool::descriptor);
-            if (!keyPool_) {
-                LOG(ERROR) << "Error getting IRemotelyProvisionedKeyPool HAL with service name '"
-                           << IRemotelyProvisionedKeyPool::descriptor << "'";
-                return 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");
+              << " directoAccess set to " << (hwInfo_.isDirectAccess ? "true" : "false")
+              << " and remote key provisioning support "
+              << (hwInfo_.isRemoteKeyProvisioningSupported ? "enabled" : "disabled");
     return true;
 }
 
@@ -140,7 +122,9 @@
     if (hwInfo_.isRemoteKeyProvisioningSupported) {
         status = setRemotelyProvisionedAttestationKey(halWritableCredential.get());
         if (!status.isOk()) {
-            return halStatusToGenericError(status);
+            LOG(WARNING) << status.toString8()
+                         << "\nUnable to fetch remotely provisioned attestation key, falling back "
+                         << "to the factory-provisioned attestation key.";
         }
     }
 
@@ -205,13 +189,21 @@
     std::vector<uint8_t> encodedCertChain;
     Status status;
 
-    if (useRkpd_) {
-        if (rpcKeyFuture_.wait_for(std::chrono::seconds(10)) != std::future_status::ready) {
+    if (useRkpd()) {
+        LOG(INFO) << "Fetching attestation key from RKPD";
+
+        uid_t callingUid = android::IPCThreadState::self()->getCallingUid();
+        auto rpcKeyFuture = getRpcKeyFuture(rpc_, callingUid);
+        if (!rpcKeyFuture) {
+            return Status::fromServiceSpecificError(ERROR_GENERIC, "Error in getRpcKeyFuture()");
+        }
+
+        if (rpcKeyFuture->wait_for(std::chrono::seconds(10)) != std::future_status::ready) {
             return Status::fromServiceSpecificError(
                 ERROR_GENERIC, "Waiting for remotely provisioned attestation key timed out");
         }
 
-        std::optional<::android::security::rkp::RemotelyProvisionedKey> key = rpcKeyFuture_.get();
+        std::optional<::android::security::rkp::RemotelyProvisionedKey> key = rpcKeyFuture->get();
         if (!key) {
             return Status::fromServiceSpecificError(
                 ERROR_GENERIC, "Failed to get remotely provisioned attestation key");
@@ -225,6 +217,16 @@
         keyBlob = std::move(key->keyBlob);
         encodedCertChain = std::move(key->encodedCertChain);
     } else {
+        LOG(INFO) << "Fetching attestation key from remotely provisioned key pool.";
+
+        sp<IRemotelyProvisionedKeyPool> keyPool =
+            android::waitForService<IRemotelyProvisionedKeyPool>(
+                IRemotelyProvisionedKeyPool::descriptor);
+        if (!keyPool) {
+            return Status::fromServiceSpecificError(
+                ERROR_GENERIC, "Error getting IRemotelyProvisionedKeyPool HAL");
+        }
+
         std::optional<std::string> rpcId = getRpcId(rpc_);
         if (!rpcId) {
             return Status::fromServiceSpecificError(
@@ -233,11 +235,9 @@
 
         uid_t callingUid = android::IPCThreadState::self()->getCallingUid();
         ::android::security::remoteprovisioning::RemotelyProvisionedKey key;
-        Status status = keyPool_->getAttestationKey(callingUid, *rpcId, &key);
+        Status status = keyPool->getAttestationKey(callingUid, *rpcId, &key);
         if (!status.isOk()) {
-            LOG(WARNING) << "Unable to fetch remotely provisioned attestation key, falling back "
-                         << "to the factory-provisioned attestation key.";
-            return Status::ok();
+            return status;
         }
 
         keyBlob = std::move(key.keyBlob);
diff --git a/identity/CredentialStore.h b/identity/CredentialStore.h
index 495841b..57c94e0 100644
--- a/identity/CredentialStore.h
+++ b/identity/CredentialStore.h
@@ -17,7 +17,6 @@
 #ifndef SYSTEM_SECURITY_CREDENTIAL_STORE_H_
 #define SYSTEM_SECURITY_CREDENTIAL_STORE_H_
 
-#include <future>
 #include <string>
 #include <vector>
 
@@ -26,8 +25,6 @@
 #include <android/security/remoteprovisioning/IRemotelyProvisionedKeyPool.h>
 #include <android/security/rkp/IRemoteProvisioning.h>
 
-#include "RemotelyProvisionedKey.h"
-
 namespace android {
 namespace security {
 namespace identity {
@@ -80,10 +77,7 @@
 
     HardwareInformation hwInfo_;
 
-    bool useRkpd_;
     sp<IRemotelyProvisionedComponent> rpc_;
-    sp<IRemotelyProvisionedKeyPool> keyPool_;
-    std::future<std::optional<RemotelyProvisionedKey>> rpcKeyFuture_;
 };
 
 }  // namespace identity
diff --git a/identity/RemotelyProvisionedKey.cpp b/identity/RemotelyProvisionedKey.cpp
index 46a42f4..7e90d63 100644
--- a/identity/RemotelyProvisionedKey.cpp
+++ b/identity/RemotelyProvisionedKey.cpp
@@ -42,6 +42,8 @@
 using ::android::security::rkp::IRemoteProvisioning;
 using ::android::security::rkp::RemotelyProvisionedKey;
 
+constexpr const char* kRemoteProvisioningServiceName = "remote_provisioning";
+
 std::optional<String16> findRpcNameById(std::string_view targetRpcId) {
     auto deviceManifest = vintf::VintfObject::GetDeviceHalManifest();
     auto instances = deviceManifest->getAidlInstances("android.hardware.security.keymint",
@@ -182,7 +184,7 @@
     }
 
     sp<IRemoteProvisioning> remoteProvisioning =
-        android::waitForService<IRemoteProvisioning>(IRemoteProvisioning::descriptor);
+        android::waitForService<IRemoteProvisioning>(String16(kRemoteProvisioningServiceName));
     if (!remoteProvisioning) {
         LOG(ERROR) << "Failed to get IRemoteProvisioning HAL";
         return std::nullopt;
diff --git a/identity/main.cpp b/identity/main.cpp
index 2559789..b3a41ec 100644
--- a/identity/main.cpp
+++ b/identity/main.cpp
@@ -23,6 +23,7 @@
 #include <android-base/logging.h>
 #include <binder/IPCThreadState.h>
 #include <binder/IServiceManager.h>
+#include <binder/ProcessState.h>
 
 #include "CredentialStoreFactory.h"
 
@@ -32,6 +33,7 @@
 
 using ::android::IPCThreadState;
 using ::android::IServiceManager;
+using ::android::ProcessState;
 using ::android::sp;
 using ::android::String16;
 using ::android::base::InitLogging;
@@ -53,8 +55,10 @@
     CHECK(ret == ::android::OK) << "Couldn't register binder service";
     LOG(INFO) << "Registered binder service";
 
-    // Credstore is a single-threaded process. So devote the main thread
-    // to handling binder messages.
+    // Credstore needs one thread to handle binder messages and one to handle
+    // asynchronous responses from RKPD.
+    ProcessState::self()->setThreadPoolMaxThreadCount(2);
+    ProcessState::self()->startThreadPool();
     IPCThreadState::self()->joinThreadPool();
 
     return 0;
diff --git a/keystore2/Android.bp b/keystore2/Android.bp
index fa80563..d8c0081 100644
--- a/keystore2/Android.bp
+++ b/keystore2/Android.bp
@@ -44,7 +44,6 @@
         "android.security.rkp_aidl-rust",
         "libanyhow",
         "libbinder_rs",
-        "libfutures",
         "libkeystore2_aaid-rust",
         "libkeystore2_apc_compat-rust",
         "libkeystore2_crypto_rust",
@@ -60,6 +59,7 @@
         "libserde",
         "libserde_cbor",
         "libthiserror",
+        "libtokio",
     ],
     shared_libs: [
         "libcutils",
@@ -158,6 +158,7 @@
         "watchdog",
         "keystore2_blob_test_utils",
     ],
+    require_root: true,
 }
 
 rust_defaults {
diff --git a/keystore2/aaid/lib.rs b/keystore2/aaid/lib.rs
index 3187198..8f6a09e 100644
--- a/keystore2/aaid/lib.rs
+++ b/keystore2/aaid/lib.rs
@@ -29,7 +29,7 @@
     // 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()),
+        0 => Ok(buffer[0..size].to_vec()),
         status => Err(status),
     }
 }
diff --git a/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl b/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
index 3f33431..e3b7d11 100644
--- a/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
+++ b/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
@@ -41,8 +41,26 @@
 
     /**
      * Unlocks the keystore for the given user id.
+     *
      * Callers require 'Unlock' permission.
-     * If a password was set, a password must be given on unlock or the operation fails.
+     *
+     * Super-Encryption Key:
+     * When the device is unlocked (and password is non-null), Keystore stores in memory
+     * a super-encryption key derived from the password that protects UNLOCKED_DEVICE_REQUIRED
+     * keys; this key is wiped from memory when the device is locked.
+     *
+     * If unlockingSids is non-empty on lock, then before the super-encryption key is wiped from
+     * memory, a copy of it is stored in memory encrypted with a fresh AES key. This key is then
+     * imported into KM, tagged such that it can be used given a valid, recent auth token for any
+     * of the unlockingSids.
+     *
+     * Options for unlock:
+     *  - If the password is non-null, the super-encryption key is re-derived as above.
+     *  - If the password is null, then if a suitable auth token to access the encrypted
+     *    Super-encryption key stored in KM has been sent to keystore (via addAuthToken), the
+     *    encrypted super-encryption key is recovered so that UNLOCKED_DEVICE_REQUIRED keys can
+     *    be used once again.
+     *  - If neither of these are met, then the operation fails.
      *
      * ## Error conditions:
      * `ResponseCode::PERMISSION_DENIED` - if the callers do not have the 'Unlock' permission.
@@ -50,33 +68,10 @@
      * `ResponseCode::VALUE_CORRUPTED` - if the super key can not be decrypted.
      * `ResponseCode::KEY_NOT_FOUND` - if the super key is not found.
      *
-     * @lockScreenEvent - Indicates what happened.
-     *                    * LockScreenEvent.UNLOCK if the screen was unlocked.
-     *                    * LockScreenEvent.LOCK if the screen was locked.
-     *
-     * @param userId - Android user id
-     *
-     * @param password - synthetic password derived by the user denoted by the user id
-     *
-     * @param unlockingSids - list of biometric SIDs for this user. This will be null when
-     *                        lockScreenEvent is UNLOCK, but may be non-null when
-     *                        lockScreenEvent is LOCK.
-     *
-     *                        When the device is unlocked, Keystore stores in memory
-     *                        a super-encryption key that protects UNLOCKED_DEVICE_REQUIRED
-     *                        keys; this key is wiped from memory when the device is locked.
-     *
-     *                        If unlockingSids is non-empty on lock, then before the
-     *                        super-encryption key is wiped from memory, a copy of it
-     *                        is stored in memory encrypted with a fresh AES key.
-     *                        This key is then imported into KM, tagged such that it can be
-     *                        used given a valid, recent auth token for any of the
-     *                        unlockingSids.
-     *
-     *                        Then, when the device is unlocked again, if a suitable auth token
-     *                        has been sent to keystore, it is used to recover the
-     *                        super-encryption key, so that UNLOCKED_DEVICE_REQUIRED keys can
-     *                        be used once again.
+     * @param lockScreenEvent whether the lock screen locked or unlocked
+     * @param userId android user id
+     * @param password synthetic password derived from the user's LSKF, must be null on lock
+     * @param unlockingSids list of biometric SIDs for this user, ignored on unlock
      */
     void onLockScreenEvent(in LockScreenEvent lockScreenEvent, in int userId,
                            in @nullable byte[] password, in @nullable long[] unlockingSids);
diff --git a/keystore2/apc_compat/apc_compat.rs b/keystore2/apc_compat/apc_compat.rs
index 9f44927..480f14d 100644
--- a/keystore2/apc_compat/apc_compat.rs
+++ b/keystore2/apc_compat/apc_compat.rs
@@ -94,7 +94,7 @@
             // If the pointer and size is not nullptr and not 0 respectively, the C/C++
             // implementation must pass a valid pointer to an allocation of at least size bytes,
             // and the pointer must be valid until this function returns.
-            unsafe { slice::from_raw_parts(tbs_message, s as usize) },
+            unsafe { slice::from_raw_parts(tbs_message, s) },
         ),
     };
     let confirmation_token = match (confirmation_token.is_null(), confirmation_token_size) {
@@ -104,7 +104,7 @@
             // If the pointer and size is not nullptr and not 0 respectively, the C/C++
             // implementation must pass a valid pointer to an allocation of at least size bytes,
             // and the pointer must be valid until this function returns.
-            unsafe { slice::from_raw_parts(confirmation_token, s as usize) },
+            unsafe { slice::from_raw_parts(confirmation_token, s) },
         ),
     };
     hal_cb(rc, tbs_message, confirmation_token)
@@ -178,7 +178,7 @@
                 cb,
                 prompt_text.as_ptr(),
                 extra_data.as_ptr(),
-                extra_data.len() as usize,
+                extra_data.len(),
                 locale.as_ptr(),
                 ui_opts,
             )
diff --git a/keystore2/legacykeystore/lib.rs b/keystore2/legacykeystore/lib.rs
index ed5bd4f..464f0a2 100644
--- a/keystore2/legacykeystore/lib.rs
+++ b/keystore2/legacykeystore/lib.rs
@@ -502,10 +502,8 @@
     ) -> Result<bool> {
         let blob = legacy_loader
             .read_legacy_keystore_entry(uid, alias, |ciphertext, iv, tag, _salt, _key_size| {
-                if let Some(key) = SUPER_KEY
-                    .read()
-                    .unwrap()
-                    .get_per_boot_key_by_user_id(uid_to_android_user(uid as u32))
+                if let Some(key) =
+                    SUPER_KEY.read().unwrap().get_per_boot_key_by_user_id(uid_to_android_user(uid))
                 {
                     key.decrypt(ciphertext, iv, tag)
                 } else {
diff --git a/keystore2/src/async_task.rs b/keystore2/src/async_task.rs
index 0515c8f..6548445 100644
--- a/keystore2/src/async_task.rs
+++ b/keystore2/src/async_task.rs
@@ -67,7 +67,7 @@
     pub fn get_mut<T: Any + Send + Default>(&mut self) -> &mut T {
         self.0
             .entry(TypeId::of::<T>())
-            .or_insert_with(|| Box::new(T::default()) as Box<dyn Any + Send>)
+            .or_insert_with(|| Box::<T>::default() as Box<dyn Any + Send>)
             .downcast_mut::<T>()
             .unwrap()
     }
diff --git a/keystore2/src/attestation_key_utils.rs b/keystore2/src/attestation_key_utils.rs
index d01cf86..d31fa82 100644
--- a/keystore2/src/attestation_key_utils.rs
+++ b/keystore2/src/attestation_key_utils.rs
@@ -30,6 +30,7 @@
 };
 use anyhow::{Context, Result};
 use keystore2_crypto::parse_subject_from_certificate;
+use rustutils::system_properties;
 
 /// KeyMint takes two different kinds of attestation keys. Remote provisioned keys
 /// and those that have been generated by the user. Unfortunately, they need to be
@@ -53,9 +54,11 @@
 }
 
 fn use_rkpd() -> bool {
-    let property_name = "persist.device_config.remote_key_provisioning_native.enable_rkpd";
+    let mutable_property = "persist.device_config.remote_key_provisioning_native.enable_rkpd";
+    let fixed_property = "remote_provisioning.enable_rkpd";
     let default_value = false;
-    rustutils::system_properties::read_bool(property_name, default_value).unwrap_or(default_value)
+    system_properties::read_bool(mutable_property, default_value).unwrap_or(default_value)
+        || system_properties::read_bool(fixed_property, default_value).unwrap_or(default_value)
 }
 
 /// This function loads and, optionally, assigns the caller's remote provisioned
diff --git a/keystore2/src/crypto/lib.rs b/keystore2/src/crypto/lib.rs
index 7ba47c8..08b7589 100644
--- a/keystore2/src/crypto/lib.rs
+++ b/keystore2/src/crypto/lib.rs
@@ -360,8 +360,7 @@
     // Safety: the key is valid.
     // This will not write past the specified length of the buffer; if the
     // len above is too short, it returns 0.
-    let written_len =
-        unsafe { ECKEYMarshalPrivateKey(key.0, buf.as_mut_ptr(), buf.len()) } as usize;
+    let written_len = unsafe { ECKEYMarshalPrivateKey(key.0, buf.as_mut_ptr(), buf.len()) };
     if written_len == len {
         Ok(buf)
     } else {
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 62fd579..c9c28f6 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -830,7 +830,7 @@
     pub fn satisfies(&self, user_secure_ids: &[i64], auth_type: HardwareAuthenticatorType) -> bool {
         user_secure_ids.iter().any(|&sid| {
             (sid == self.auth_token.userId || sid == self.auth_token.authenticatorId)
-                && (((auth_type.0 as i32) & (self.auth_token.authenticatorType.0 as i32)) != 0)
+                && ((auth_type.0 & self.auth_token.authenticatorType.0) != 0)
         })
     }
 
@@ -1859,7 +1859,8 @@
                 let (_, hw_info) = get_keymint_dev_by_uuid(km_uuid)
                     .context("Error in retrieving keymint device by UUID.")?;
                 log_rkp_error_stats(MetricsRkpError::OUT_OF_KEYS, &hw_info.securityLevel);
-                return Err(KsError::Rc(ResponseCode::OUT_OF_KEYS)).context("Out of keys.");
+                return Err(KsError::Rc(ResponseCode::OUT_OF_KEYS_TRANSIENT_ERROR))
+                    .context("Out of keys.");
             } else if result > 1 {
                 return Err(KsError::sys())
                     .context(format!("Expected to update 1 entry, instead updated {}", result));
@@ -4562,7 +4563,7 @@
                 DESTINATION_UID,
                 |k, av| {
                     assert_eq!(Domain::SELINUX, k.domain);
-                    assert_eq!(DESTINATION_NAMESPACE as i64, k.nspace);
+                    assert_eq!(DESTINATION_NAMESPACE, k.nspace);
                     assert!(av.is_none());
                     Ok(())
                 },
diff --git a/keystore2/src/error.rs b/keystore2/src/error.rs
index b60b64f..d1d58a4 100644
--- a/keystore2/src/error.rs
+++ b/keystore2/src/error.rs
@@ -72,9 +72,9 @@
         Error::Rc(ResponseCode::PERMISSION_DENIED)
     }
 
-    /// Short hand for `Error::Rc(ResponseCode::OUT_OF_KEYS)`
+    /// Short hand for `Error::Rc(ResponseCode::OUT_OF_KEYS_TRANSIENT_ERROR)`
     pub fn out_of_keys() -> Self {
-        Error::Rc(ResponseCode::OUT_OF_KEYS)
+        Error::Rc(ResponseCode::OUT_OF_KEYS_TRANSIENT_ERROR)
     }
 }
 
diff --git a/keystore2/src/legacy_blob.rs b/keystore2/src/legacy_blob.rs
index 7cf1819..2ffcc71 100644
--- a/keystore2/src/legacy_blob.rs
+++ b/keystore2/src/legacy_blob.rs
@@ -321,7 +321,7 @@
                     acc.push(c as char);
                 }
                 c => {
-                    acc.push((b'+' + (c as u8 >> 6)) as char);
+                    acc.push((b'+' + (c >> 6)) as char);
                     acc.push((b'0' + (c & 0x3F)) as char);
                 }
             };
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index cb2962a..1a83339 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -138,8 +138,8 @@
     /// (2) if remote provisioning is present and enabled on the system. If these conditions are
     /// met, it makes an attempt to fetch the attestation key assigned to the `caller_uid`.
     ///
-    /// It returns the ResponseCode `OUT_OF_KEYS` if there is not one key currently assigned to the
-    /// `caller_uid` and there are none available to assign.
+    /// It returns the ResponseCode `OUT_OF_KEYS_TRANSIENT_ERROR` if there is not one key currently
+    /// assigned to the `caller_uid` and there are none available to assign.
     pub fn get_remotely_provisioned_attestation_key_and_certs(
         &self,
         key: &KeyDescriptor,
@@ -490,7 +490,7 @@
 /// Fetches a remote provisioning attestation key and certificate chain inside of the
 /// returned `CertificateChain` struct if one exists for the given caller_uid. If one has not
 /// been assigned, this function will assign it. If there are no signed attestation keys
-/// available to be assigned, it will return the ResponseCode `OUT_OF_KEYS`
+/// available to be assigned, it will return the ResponseCode `OUT_OF_KEYS_TRANSIENT_ERROR`
 fn get_rem_prov_attest_key(
     domain: Domain,
     caller_uid: u32,
@@ -645,7 +645,7 @@
     /// Fetches a remotely provisioned certificate chain and key for the given client uid that
     /// was provisioned using the IRemotelyProvisionedComponent with the given id. The same key
     /// will be returned for a given caller_uid on every request. If there are no attestation keys
-    /// available, `OUT_OF_KEYS` is returned.
+    /// available, `OUT_OF_KEYS_TRANSIENT_ERROR` is returned.
     fn get_attestation_key(
         &self,
         db: &mut KeystoreDB,
@@ -671,7 +671,7 @@
             }),
             // It should be impossible to get `None`, but handle it just in case as a
             // precaution against future behavioral changes in `get_rem_prov_attest_key`.
-            None => Err(error::Error::Rc(ResponseCode::OUT_OF_KEYS))
+            None => Err(error::Error::Rc(ResponseCode::OUT_OF_KEYS_TRANSIENT_ERROR))
                 .context(ks_err!("No available attestation keys")),
         }
     }
@@ -958,7 +958,7 @@
                 .unwrap_err()
                 .downcast::<error::Error>()
                 .unwrap(),
-            error::Error::Rc(ResponseCode::OUT_OF_KEYS)
+            error::Error::Rc(ResponseCode::OUT_OF_KEYS_TRANSIENT_ERROR)
         );
     }
 
@@ -1023,7 +1023,7 @@
                 .unwrap_err()
                 .downcast::<error::Error>()
                 .unwrap(),
-            error::Error::Rc(ResponseCode::OUT_OF_KEYS)
+            error::Error::Rc(ResponseCode::OUT_OF_KEYS_TRANSIENT_ERROR)
         );
     }
 
diff --git a/keystore2/src/rkpd_client.rs b/keystore2/src/rkpd_client.rs
index b426440..bfc7847 100644
--- a/keystore2/src/rkpd_client.rs
+++ b/keystore2/src/rkpd_client.rs
@@ -13,8 +13,9 @@
 // limitations under the License.
 
 //! Helper wrapper around RKPD interface.
+// TODO(b/264891956): Return RKP specific errors.
 
-use crate::error::{map_binder_status_code, Error, ErrorCode};
+use crate::error::{map_binder_status_code, Error};
 use crate::globals::get_remotely_provisioned_component_name;
 use crate::ks_err;
 use crate::utils::watchdog as wd;
@@ -29,10 +30,21 @@
     RemotelyProvisionedKey::RemotelyProvisionedKey,
 };
 use android_security_rkp_aidl::binder::{BinderFeatures, Interface, Strong};
+use android_system_keystore2::aidl::android::system::keystore2::ResponseCode::ResponseCode;
 use anyhow::{Context, Result};
-use futures::channel::oneshot;
-use futures::executor::block_on;
 use std::sync::Mutex;
+use std::time::Duration;
+use tokio::sync::oneshot;
+use tokio::time::timeout;
+
+// Normally, we block indefinitely when making calls outside of keystore and rely on watchdog to
+// report deadlocks. However, RKPD is mainline updatable. Also, calls to RKPD may wait on network
+// for certificates. So, we err on the side of caution and timeout instead.
+static RKPD_TIMEOUT: Duration = Duration::from_secs(10);
+
+fn tokio_rt() -> tokio::runtime::Runtime {
+    tokio::runtime::Builder::new_current_thread().enable_all().build().unwrap()
+}
 
 /// Thread-safe channel for sending a value once and only once. If a value has
 /// already been send, subsequent calls to send will noop.
@@ -79,7 +91,7 @@
         let _wp = wd::watch_millis("IGetRegistrationCallback::onCancel", 500);
         log::warn!("IGetRegistrationCallback cancelled");
         self.registration_tx.send(
-            Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS))
                 .context(ks_err!("GetRegistrationCallback cancelled.")),
         );
         Ok(())
@@ -88,7 +100,7 @@
         let _wp = wd::watch_millis("IGetRegistrationCallback::onError", 500);
         log::error!("IGetRegistrationCallback failed: '{error}'");
         self.registration_tx.send(
-            Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS))
                 .context(ks_err!("GetRegistrationCallback failed: {:?}", error)),
         );
         Ok(())
@@ -113,7 +125,12 @@
         .getRegistration(&rpc_name, &cb)
         .context(ks_err!("Trying to get registration."))?;
 
-    rx.await.unwrap()
+    match timeout(RKPD_TIMEOUT, rx).await {
+        Err(e) => {
+            Err(Error::Rc(ResponseCode::SYSTEM_ERROR)).context(ks_err!("Waiting for RKPD: {:?}", e))
+        }
+        Ok(v) => v.unwrap(),
+    }
 }
 
 struct GetKeyCallback {
@@ -144,8 +161,7 @@
         let _wp = wd::watch_millis("IGetKeyCallback::onCancel", 500);
         log::warn!("IGetKeyCallback cancelled");
         self.key_tx.send(
-            Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
-                .context(ks_err!("GetKeyCallback cancelled.")),
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS)).context(ks_err!("GetKeyCallback cancelled.")),
         );
         Ok(())
     }
@@ -153,13 +169,31 @@
         let _wp = wd::watch_millis("IGetKeyCallback::onError", 500);
         log::error!("IGetKeyCallback failed: {error}");
         self.key_tx.send(
-            Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS))
                 .context(ks_err!("GetKeyCallback failed: {:?}", error)),
         );
         Ok(())
     }
 }
 
+async fn get_rkpd_attestation_key_from_registration_async(
+    registration: &Strong<dyn IRegistration>,
+    caller_uid: u32,
+) -> Result<RemotelyProvisionedKey> {
+    let (tx, rx) = oneshot::channel();
+    let cb = GetKeyCallback::new_native_binder(tx);
+
+    registration
+        .getKey(caller_uid.try_into().unwrap(), &cb)
+        .context(ks_err!("Trying to get key."))?;
+
+    match timeout(RKPD_TIMEOUT, rx).await {
+        Err(e) => Err(Error::Rc(ResponseCode::OUT_OF_KEYS))
+            .context(ks_err!("Waiting for RKPD key timed out: {:?}", e)),
+        Ok(v) => v.unwrap(),
+    }
+}
+
 async fn get_rkpd_attestation_key_async(
     security_level: &SecurityLevel,
     caller_uid: u32,
@@ -167,15 +201,7 @@
     let registration = get_rkpd_registration(security_level)
         .await
         .context(ks_err!("Trying to get to IRegistration service."))?;
-
-    let (tx, rx) = oneshot::channel();
-    let cb = GetKeyCallback::new_native_binder(tx);
-
-    registration
-        .getKey(caller_uid.try_into().unwrap(), &cb)
-        .context(ks_err!("Trying to get key."))?;
-
-    rx.await.unwrap()
+    get_rkpd_attestation_key_from_registration_async(&registration, caller_uid).await
 }
 
 struct StoreUpgradedKeyCallback {
@@ -204,13 +230,32 @@
         let _wp = wd::watch_millis("IGetRegistrationCallback::onError", 500);
         log::error!("IGetRegistrationCallback failed: {error}");
         self.completer.send(
-            Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
+            Err(Error::Rc(ResponseCode::SYSTEM_ERROR))
                 .context(ks_err!("Failed to store upgraded key: {:?}", error)),
         );
         Ok(())
     }
 }
 
+async fn store_rkpd_attestation_key_with_registration_async(
+    registration: &Strong<dyn IRegistration>,
+    key_blob: &[u8],
+    upgraded_blob: &[u8],
+) -> Result<()> {
+    let (tx, rx) = oneshot::channel();
+    let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+
+    registration
+        .storeUpgradedKeyAsync(key_blob, upgraded_blob, &cb)
+        .context(ks_err!("Failed to store upgraded blob with RKPD."))?;
+
+    match timeout(RKPD_TIMEOUT, rx).await {
+        Err(e) => Err(Error::Rc(ResponseCode::SYSTEM_ERROR))
+            .context(ks_err!("Waiting for RKPD to complete storing key: {:?}", e)),
+        Ok(v) => v.unwrap(),
+    }
+}
+
 async fn store_rkpd_attestation_key_async(
     security_level: &SecurityLevel,
     key_blob: &[u8],
@@ -219,15 +264,7 @@
     let registration = get_rkpd_registration(security_level)
         .await
         .context(ks_err!("Trying to get to IRegistration service."))?;
-
-    let (tx, rx) = oneshot::channel();
-    let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
-
-    registration
-        .storeUpgradedKeyAsync(key_blob, upgraded_blob, &cb)
-        .context(ks_err!("Failed to store upgraded blob with RKPD."))?;
-
-    rx.await.unwrap()
+    store_rkpd_attestation_key_with_registration_async(&registration, key_blob, upgraded_blob).await
 }
 
 /// Get attestation key from RKPD.
@@ -236,7 +273,7 @@
     caller_uid: u32,
 ) -> Result<RemotelyProvisionedKey> {
     let _wp = wd::watch_millis("Calling get_rkpd_attestation_key()", 500);
-    block_on(get_rkpd_attestation_key_async(security_level, caller_uid))
+    tokio_rt().block_on(get_rkpd_attestation_key_async(security_level, caller_uid))
 }
 
 /// Store attestation key in RKPD.
@@ -246,26 +283,41 @@
     upgraded_blob: &[u8],
 ) -> Result<()> {
     let _wp = wd::watch_millis("Calling store_rkpd_attestation_key()", 500);
-    block_on(store_rkpd_attestation_key_async(security_level, key_blob, upgraded_blob))
+    tokio_rt().block_on(store_rkpd_attestation_key_async(security_level, key_blob, upgraded_blob))
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::error::map_km_error;
+    use crate::globals::get_keymint_device;
+    use crate::utils::upgrade_keyblob_if_required_with;
+    use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+        Algorithm::Algorithm, AttestationKey::AttestationKey, KeyParameter::KeyParameter,
+        KeyParameterValue::KeyParameterValue, Tag::Tag,
+    };
     use android_security_rkp_aidl::aidl::android::security::rkp::IRegistration::BnRegistration;
-    use std::sync::Arc;
+    use keystore2_crypto::parse_subject_from_certificate;
+    use std::sync::atomic::{AtomicU32, Ordering};
 
     #[derive(Default)]
-    struct MockRegistrationValues {
-        _key: RemotelyProvisionedKey,
+    struct MockRegistration {
+        key: RemotelyProvisionedKey,
+        latency: Option<Duration>,
     }
 
-    #[derive(Default)]
-    struct MockRegistration(Arc<Mutex<MockRegistrationValues>>);
-
     impl MockRegistration {
-        pub fn new_native_binder() -> Strong<dyn IRegistration> {
-            let result: Self = Default::default();
+        pub fn new_native_binder(
+            key: &RemotelyProvisionedKey,
+            latency: Option<Duration>,
+        ) -> Strong<dyn IRegistration> {
+            let result = Self {
+                key: RemotelyProvisionedKey {
+                    keyBlob: key.keyBlob.clone(),
+                    encodedCertChain: key.encodedCertChain.clone(),
+                },
+                latency,
+            };
             BnRegistration::new_binder(result, BinderFeatures::default())
         }
     }
@@ -273,8 +325,22 @@
     impl Interface for MockRegistration {}
 
     impl IRegistration for MockRegistration {
-        fn getKey(&self, _: i32, _: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
-            todo!()
+        fn getKey(&self, _: i32, cb: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
+            let key = RemotelyProvisionedKey {
+                keyBlob: self.key.keyBlob.clone(),
+                encodedCertChain: self.key.encodedCertChain.clone(),
+            };
+            let latency = self.latency;
+            let get_key_cb = cb.clone();
+
+            // Need a separate thread to trigger timeout in the caller.
+            std::thread::spawn(move || {
+                if let Some(duration) = latency {
+                    std::thread::sleep(duration);
+                }
+                get_key_cb.onSuccess(&key).unwrap();
+            });
+            Ok(())
         }
 
         fn cancelGetKey(&self, _: &Strong<dyn IGetKeyCallback>) -> binder::Result<()> {
@@ -285,24 +351,46 @@
             &self,
             _: &[u8],
             _: &[u8],
-            _: &Strong<dyn IStoreUpgradedKeyCallback>,
+            cb: &Strong<dyn IStoreUpgradedKeyCallback>,
         ) -> binder::Result<()> {
-            todo!()
+            // We are primarily concerned with timing out correctly. Storing the key in this mock
+            // registration isn't particularly interesting, so skip that part.
+            let store_cb = cb.clone();
+            let latency = self.latency;
+
+            std::thread::spawn(move || {
+                if let Some(duration) = latency {
+                    std::thread::sleep(duration);
+                }
+                store_cb.onSuccess().unwrap();
+            });
+            Ok(())
         }
     }
 
-    fn get_mock_registration() -> Result<binder::Strong<dyn IRegistration>> {
+    fn get_mock_registration(
+        key: &RemotelyProvisionedKey,
+        latency: Option<Duration>,
+    ) -> Result<binder::Strong<dyn IRegistration>> {
         let (tx, rx) = oneshot::channel();
         let cb = GetRegistrationCallback::new_native_binder(tx);
-        let mock_registration = MockRegistration::new_native_binder();
+        let mock_registration = MockRegistration::new_native_binder(key, latency);
 
         assert!(cb.onSuccess(&mock_registration).is_ok());
-        block_on(rx).unwrap()
+        tokio_rt().block_on(rx).unwrap()
+    }
+
+    // Using the same key ID makes test cases race with each other. So, we use separate key IDs for
+    // different test cases.
+    fn get_next_key_id() -> u32 {
+        static ID: AtomicU32 = AtomicU32::new(0);
+        ID.fetch_add(1, Ordering::Relaxed)
     }
 
     #[test]
     fn test_get_registration_cb_success() {
-        let registration = get_mock_registration();
+        let key: RemotelyProvisionedKey = Default::default();
+        let registration = get_mock_registration(&key, /*latency=*/ None);
         assert!(registration.is_ok());
     }
 
@@ -312,10 +400,10 @@
         let cb = GetRegistrationCallback::new_native_binder(tx);
         assert!(cb.onCancel().is_ok());
 
-        let result = block_on(rx).unwrap();
+        let result = tokio_rt().block_on(rx).unwrap();
         assert_eq!(
             result.unwrap_err().downcast::<Error>().unwrap(),
-            Error::Km(ErrorCode::OPERATION_CANCELLED)
+            Error::Rc(ResponseCode::OUT_OF_KEYS)
         );
     }
 
@@ -325,10 +413,10 @@
         let cb = GetRegistrationCallback::new_native_binder(tx);
         assert!(cb.onError("error").is_ok());
 
-        let result = block_on(rx).unwrap();
+        let result = tokio_rt().block_on(rx).unwrap();
         assert_eq!(
             result.unwrap_err().downcast::<Error>().unwrap(),
-            Error::Km(ErrorCode::UNKNOWN_ERROR)
+            Error::Rc(ResponseCode::OUT_OF_KEYS)
         );
     }
 
@@ -340,7 +428,7 @@
         let cb = GetKeyCallback::new_native_binder(tx);
         assert!(cb.onSuccess(&mock_key).is_ok());
 
-        let key = block_on(rx).unwrap().unwrap();
+        let key = tokio_rt().block_on(rx).unwrap().unwrap();
         assert_eq!(key, mock_key);
     }
 
@@ -350,10 +438,10 @@
         let cb = GetKeyCallback::new_native_binder(tx);
         assert!(cb.onCancel().is_ok());
 
-        let result = block_on(rx).unwrap();
+        let result = tokio_rt().block_on(rx).unwrap();
         assert_eq!(
             result.unwrap_err().downcast::<Error>().unwrap(),
-            Error::Km(ErrorCode::OPERATION_CANCELLED)
+            Error::Rc(ResponseCode::OUT_OF_KEYS)
         );
     }
 
@@ -363,10 +451,10 @@
         let cb = GetKeyCallback::new_native_binder(tx);
         assert!(cb.onError("error").is_ok());
 
-        let result = block_on(rx).unwrap();
+        let result = tokio_rt().block_on(rx).unwrap();
         assert_eq!(
             result.unwrap_err().downcast::<Error>().unwrap(),
-            Error::Km(ErrorCode::UNKNOWN_ERROR)
+            Error::Rc(ResponseCode::OUT_OF_KEYS)
         );
     }
 
@@ -376,7 +464,7 @@
         let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
         assert!(cb.onSuccess().is_ok());
 
-        block_on(rx).unwrap().unwrap();
+        tokio_rt().block_on(rx).unwrap().unwrap();
     }
 
     #[test]
@@ -385,17 +473,73 @@
         let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
         assert!(cb.onError("oh no! it failed").is_ok());
 
-        let result = block_on(rx).unwrap();
+        let result = tokio_rt().block_on(rx).unwrap();
         assert_eq!(
             result.unwrap_err().downcast::<Error>().unwrap(),
-            Error::Km(ErrorCode::UNKNOWN_ERROR)
+            Error::Rc(ResponseCode::SYSTEM_ERROR)
+        );
+    }
+
+    #[test]
+    fn test_get_mock_key_success() {
+        let mock_key =
+            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+        let registration = get_mock_registration(&mock_key, /*latency=*/ None).unwrap();
+
+        let key = tokio_rt()
+            .block_on(get_rkpd_attestation_key_from_registration_async(&registration, 0))
+            .unwrap();
+        assert_eq!(key, mock_key);
+    }
+
+    #[test]
+    fn test_get_mock_key_timeout() {
+        let mock_key =
+            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+        let latency = RKPD_TIMEOUT + Duration::from_secs(10);
+        let registration = get_mock_registration(&mock_key, Some(latency)).unwrap();
+
+        let result =
+            tokio_rt().block_on(get_rkpd_attestation_key_from_registration_async(&registration, 0));
+        assert_eq!(
+            result.unwrap_err().downcast::<Error>().unwrap(),
+            Error::Rc(ResponseCode::OUT_OF_KEYS)
+        );
+    }
+
+    #[test]
+    fn test_store_mock_key_success() {
+        let mock_key =
+            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+        let registration = get_mock_registration(&mock_key, /*latency=*/ None).unwrap();
+        tokio_rt()
+            .block_on(store_rkpd_attestation_key_with_registration_async(&registration, &[], &[]))
+            .unwrap();
+    }
+
+    #[test]
+    fn test_store_mock_key_timeout() {
+        let mock_key =
+            RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
+        let latency = RKPD_TIMEOUT + Duration::from_secs(10);
+        let registration = get_mock_registration(&mock_key, Some(latency)).unwrap();
+
+        let result = tokio_rt().block_on(store_rkpd_attestation_key_with_registration_async(
+            &registration,
+            &[],
+            &[],
+        ));
+        assert_eq!(
+            result.unwrap_err().downcast::<Error>().unwrap(),
+            Error::Rc(ResponseCode::SYSTEM_ERROR)
         );
     }
 
     #[test]
     fn test_get_rkpd_attestation_key() {
         binder::ProcessState::start_thread_pool();
-        let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, 0).unwrap();
+        let key_id = get_next_key_id();
+        let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, key_id).unwrap();
         assert!(!key.keyBlob.is_empty());
         assert!(!key.encodedCertChain.is_empty());
     }
@@ -404,11 +548,11 @@
     fn test_get_rkpd_attestation_key_same_caller() {
         binder::ProcessState::start_thread_pool();
         let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
-        let caller_uid = 0;
+        let key_id = get_next_key_id();
 
         // Multiple calls should return the same key.
-        let first_key = get_rkpd_attestation_key(&sec_level, caller_uid).unwrap();
-        let second_key = get_rkpd_attestation_key(&sec_level, caller_uid).unwrap();
+        let first_key = get_rkpd_attestation_key(&sec_level, key_id).unwrap();
+        let second_key = get_rkpd_attestation_key(&sec_level, key_id).unwrap();
 
         assert_eq!(first_key.keyBlob, second_key.keyBlob);
         assert_eq!(first_key.encodedCertChain, second_key.encodedCertChain);
@@ -418,21 +562,118 @@
     fn test_get_rkpd_attestation_key_different_caller() {
         binder::ProcessState::start_thread_pool();
         let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
+        let first_key_id = get_next_key_id();
+        let second_key_id = get_next_key_id();
 
         // Different callers should be getting different keys.
-        let first_key = get_rkpd_attestation_key(&sec_level, 1).unwrap();
-        let second_key = get_rkpd_attestation_key(&sec_level, 2).unwrap();
+        let first_key = get_rkpd_attestation_key(&sec_level, first_key_id).unwrap();
+        let second_key = get_rkpd_attestation_key(&sec_level, second_key_id).unwrap();
 
         assert_ne!(first_key.keyBlob, second_key.keyBlob);
         assert_ne!(first_key.encodedCertChain, second_key.encodedCertChain);
     }
 
     #[test]
+    // Couple of things to note:
+    // 1. This test must never run with UID of keystore. Otherwise, it can mess up keys stored by
+    //    keystore.
+    // 2. Storing and reading the stored key is prone to race condition. So, we only do this in one
+    //    test case.
     fn test_store_rkpd_attestation_key() {
         binder::ProcessState::start_thread_pool();
         let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
-        let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, 0).unwrap();
+        let key_id = get_next_key_id();
+        let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, key_id).unwrap();
+        let new_blob: [u8; 8] = rand::random();
 
-        assert!(store_rkpd_attestation_key(&sec_level, &key.keyBlob, &key.keyBlob).is_ok());
+        assert!(store_rkpd_attestation_key(&sec_level, &key.keyBlob, &new_blob).is_ok());
+
+        let new_key =
+            get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, key_id).unwrap();
+        assert_eq!(new_key.keyBlob, new_blob);
+    }
+
+    #[test]
+    // This is a helper for a manual test. We want to check that after a system upgrade RKPD
+    // attestation keys can also be upgraded and stored again with RKPD. The steps are:
+    // 1. Run this test and check in stdout that no key upgrade happened.
+    // 2. Perform a system upgrade.
+    // 3. Run this test and check in stdout that key upgrade did happen.
+    //
+    // Note that this test must be run with that same UID every time. Running as root, i.e. UID 0,
+    // should do the trick. Also, use "--nocapture" flag to get stdout.
+    fn test_rkpd_attestation_key_upgrade() {
+        binder::ProcessState::start_thread_pool();
+        let security_level = SecurityLevel::TRUSTED_ENVIRONMENT;
+        let (keymint, _, _) = get_keymint_device(&security_level).unwrap();
+        let key_id = get_next_key_id();
+        let mut key_upgraded = false;
+
+        let key = get_rkpd_attestation_key(&security_level, key_id).unwrap();
+        assert!(!key.keyBlob.is_empty());
+        assert!(!key.encodedCertChain.is_empty());
+
+        upgrade_keyblob_if_required_with(
+            &*keymint,
+            &key.keyBlob,
+            /*upgrade_params=*/ &[],
+            /*km_op=*/
+            |blob| {
+                let params = vec![
+                    KeyParameter {
+                        tag: Tag::ALGORITHM,
+                        value: KeyParameterValue::Algorithm(Algorithm::AES),
+                    },
+                    KeyParameter { tag: Tag::KEY_SIZE, value: KeyParameterValue::Integer(128) },
+                ];
+                let attestation_key = AttestationKey {
+                    keyBlob: blob.to_vec(),
+                    attestKeyParams: vec![],
+                    issuerSubjectName: parse_subject_from_certificate(&key.encodedCertChain)
+                        .unwrap(),
+                };
+
+                map_km_error(keymint.generateKey(&params, Some(&attestation_key)))
+            },
+            /*new_blob_handler=*/
+            |new_blob| {
+                // This handler is only executed if a key upgrade was performed.
+                key_upgraded = true;
+                store_rkpd_attestation_key(&security_level, &key.keyBlob, new_blob).unwrap();
+                Ok(())
+            },
+        )
+        .unwrap();
+
+        if key_upgraded {
+            println!("RKPD key was upgraded and stored with RKPD.");
+        } else {
+            println!("RKPD key was NOT upgraded.");
+        }
+    }
+
+    #[test]
+    #[ignore] // b/266607003
+    fn test_stress_get_rkpd_attestation_key() {
+        binder::ProcessState::start_thread_pool();
+        let key_id = get_next_key_id();
+        let mut threads = vec![];
+        const NTHREADS: u32 = 10;
+        const NCALLS: u32 = 1000;
+
+        for _ in 0..NTHREADS {
+            threads.push(std::thread::spawn(move || {
+                for _ in 0..NCALLS {
+                    let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, key_id)
+                        .unwrap();
+                    assert!(!key.keyBlob.is_empty());
+                    assert!(!key.encodedCertChain.is_empty());
+                }
+            }));
+        }
+
+        for t in threads {
+            assert!(t.join().is_ok());
+        }
     }
 }
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index 75d98e2..7bc548e 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -209,6 +209,7 @@
     parameters.into_iter().map(|p| p.into_authorization()).collect()
 }
 
+#[allow(clippy::unnecessary_cast)]
 /// This returns the current time (in milliseconds) as an instance of a monotonic clock,
 /// by invoking the system call since Rust does not support getting monotonic time instance
 /// as an integer.
@@ -276,7 +277,40 @@
     );
     result.sort_unstable();
     result.dedup();
-    Ok(result)
+
+    let mut items_to_return = 0;
+    let mut returned_bytes: usize = 0;
+    const RESPONSE_SIZE_LIMIT: usize = 358400;
+    // Estimate the transaction size to avoid returning more items than what
+    // could fit in a binder transaction.
+    for kd in result.iter() {
+        // 4 bytes for the Domain enum
+        // 8 bytes for the Namespace long.
+        returned_bytes += 4 + 8;
+        // Size of the alias string. Includes 4 bytes for length encoding.
+        if let Some(alias) = &kd.alias {
+            returned_bytes += 4 + alias.len();
+        }
+        // Size of the blob. Includes 4 bytes for length encoding.
+        if let Some(blob) = &kd.blob {
+            returned_bytes += 4 + blob.len();
+        }
+        // The binder transaction size limit is 1M. Empirical measurements show
+        // that the binder overhead is 60% (to be confirmed). So break after
+        // 350KB and return a partial list.
+        if returned_bytes > RESPONSE_SIZE_LIMIT {
+            log::warn!(
+                "Key descriptors list ({} items) may exceed binder \
+                       size, returning {} items est {} bytes.",
+                result.len(),
+                items_to_return,
+                returned_bytes
+            );
+            break;
+        }
+        items_to_return += 1;
+    }
+    Ok(result[..items_to_return].to_vec())
 }
 
 /// This module provides helpers for simplified use of the watchdog module.
diff --git a/keystore2/src/watchdog.rs b/keystore2/src/watchdog.rs
index a26b632..01043c5 100644
--- a/keystore2/src/watchdog.rs
+++ b/keystore2/src/watchdog.rs
@@ -141,7 +141,7 @@
             },
         );
         // Put the groups back into a vector.
-        let mut groups: Vec<Vec<(&Index, &Record)>> = groups.into_iter().map(|(_, v)| v).collect();
+        let mut groups: Vec<Vec<(&Index, &Record)>> = groups.into_values().collect();
         // Sort the groups by start time of the most recent (.last()) of each group.
         // It is panic safe to use unwrap() here because we never add empty vectors to
         // the map.
diff --git a/keystore2/tests/keystore2_client_delete_key_tests.rs b/keystore2/tests/keystore2_client_delete_key_tests.rs
new file mode 100644
index 0000000..2a06edb
--- /dev/null
+++ b/keystore2/tests/keystore2_client_delete_key_tests.rs
@@ -0,0 +1,150 @@
+// Copyright 2022, 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.
+
+use nix::unistd::getuid;
+
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    ErrorCode::ErrorCode, SecurityLevel::SecurityLevel,
+};
+use android_system_keystore2::aidl::android::system::keystore2::{
+    Domain::Domain, KeyDescriptor::KeyDescriptor, ResponseCode::ResponseCode,
+};
+
+use keystore2_test_utils::{get_keystore_service, key_generations, key_generations::Error};
+
+/// Generate a key and delete it using keystore2 service `deleteKey` API. Test should successfully
+/// delete the generated key.
+#[test]
+fn keystore2_delete_key_success() {
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+    let alias = "delete_key_success_key";
+
+    let key_metadata = key_generations::generate_ec_p256_signing_key(
+        &sec_level,
+        Domain::APP,
+        -1,
+        Some(alias.to_string()),
+        None,
+    )
+    .unwrap();
+
+    keystore2.deleteKey(&key_metadata.key).expect("Failed to delete a key.");
+
+    // Check wehther deleted key is removed from keystore.
+    let result = key_generations::map_ks_error(keystore2.getKeyEntry(&key_metadata.key));
+    assert!(result.is_err());
+    assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+}
+
+/// Try to delete non-existing key with domain other than BLOB using keystore2 service `deleteKey`
+/// API. Test should fail with an error code `KEY_NOT_FOUND`.
+#[test]
+fn keystore2_delete_key_fail() {
+    let test_alias = "delete_key_failure_key";
+    let keystore2 = get_keystore_service();
+
+    let result = key_generations::map_ks_error(keystore2.deleteKey(&KeyDescriptor {
+        domain: Domain::SELINUX,
+        nspace: key_generations::SELINUX_SHELL_NAMESPACE,
+        alias: Some(test_alias.to_string()),
+        blob: None,
+    }));
+    assert!(result.is_err());
+    assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+}
+
+/// Generate a key with `Domain::BLOB`. Try to delete a key with `Domain::BLOB` using keystore2
+/// service `deleteKey` API. Test should fail to delete a key with domain BLOB with an error code
+/// `INVALID_ARGUMENT`.
+#[test]
+fn keystore2_delete_key_with_blob_domain_fail() {
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+    let alias = "delete_key_blob_fail_key";
+
+    let key_metadata = key_generations::generate_ec_p256_signing_key(
+        &sec_level,
+        Domain::BLOB,
+        key_generations::SELINUX_SHELL_NAMESPACE,
+        Some(alias.to_string()),
+        None,
+    )
+    .unwrap();
+
+    let result = key_generations::map_ks_error(keystore2.deleteKey(&key_metadata.key));
+    assert!(result.is_err());
+    assert_eq!(Error::Rc(ResponseCode::INVALID_ARGUMENT), result.unwrap_err());
+}
+
+/// Generate a key with `Domain::BLOB`. Delete generated key with `Domain::BLOB` using underlying
+/// security level `deleteKey` API. Test should delete the key successfully.
+#[test]
+fn keystore2_delete_key_blob_success() {
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+    let alias = "delete_key_blob_success_key";
+
+    let key_metadata = key_generations::generate_ec_p256_signing_key(
+        &sec_level,
+        Domain::BLOB,
+        key_generations::SELINUX_SHELL_NAMESPACE,
+        Some(alias.to_string()),
+        None,
+    )
+    .unwrap();
+
+    let result = sec_level.deleteKey(&key_metadata.key);
+    assert!(result.is_ok());
+}
+
+/// Try to delete a key with `Domain::BLOB` without providing key-blob. Test should fail to delete a
+/// key with error code `INVALID_ARGUMENT`.
+#[test]
+fn keystore2_delete_key_fails_with_missing_key_blob() {
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+
+    let result = key_generations::map_ks_error(sec_level.deleteKey(&KeyDescriptor {
+        domain: Domain::BLOB,
+        nspace: key_generations::SELINUX_SHELL_NAMESPACE,
+        alias: None,
+        blob: None,
+    }));
+    assert!(result.is_err());
+    assert_eq!(Error::Km(ErrorCode::INVALID_ARGUMENT), result.unwrap_err());
+}
+
+/// Try to delete a key with domain other than `Domain::BLOB` using underlying security-level
+/// `deleteKey` API. Test should fail to delete a key-blob from underlying security-level backend
+/// with error code `INVALID_ARGUMENT`.
+#[test]
+fn keystore2_delete_key_blob_fail() {
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+    let alias = format!("ks_delete_keyblob_test_key_{}", getuid());
+
+    let key_metadata = key_generations::generate_ec_p256_signing_key(
+        &sec_level,
+        Domain::APP,
+        -1,
+        Some(alias),
+        None,
+    )
+    .unwrap();
+
+    let result = key_generations::map_ks_error(sec_level.deleteKey(&key_metadata.key));
+    assert!(result.is_err());
+    assert_eq!(Error::Km(ErrorCode::INVALID_ARGUMENT), result.unwrap_err());
+}
diff --git a/keystore2/tests/keystore2_client_grant_key_tests.rs b/keystore2/tests/keystore2_client_grant_key_tests.rs
index 7c75734..bde872d 100644
--- a/keystore2/tests/keystore2_client_grant_key_tests.rs
+++ b/keystore2/tests/keystore2_client_grant_key_tests.rs
@@ -19,7 +19,8 @@
     Digest::Digest, KeyPurpose::KeyPurpose, SecurityLevel::SecurityLevel,
 };
 use android_system_keystore2::aidl::android::system::keystore2::{
-    Domain::Domain, KeyDescriptor::KeyDescriptor, KeyPermission::KeyPermission,
+    Domain::Domain, IKeystoreSecurityLevel::IKeystoreSecurityLevel,
+    IKeystoreService::IKeystoreService, KeyDescriptor::KeyDescriptor, KeyPermission::KeyPermission,
     ResponseCode::ResponseCode,
 };
 
@@ -27,7 +28,9 @@
     authorizations, get_keystore_service, key_generations, key_generations::Error, run_as,
 };
 
-use crate::keystore2_client_test_utils::perform_sample_sign_operation;
+use crate::keystore2_client_test_utils::{
+    generate_ec_key_and_grant_to_users, perform_sample_sign_operation,
+};
 
 /// Generate an EC signing key and grant it to the user with given access vector.
 fn generate_ec_key_and_grant_to_user(
@@ -50,6 +53,36 @@
     keystore2.grant(&key_metadata.key, grantee_uid, access_vector)
 }
 
+fn load_grant_key_and_perform_sign_operation(
+    keystore2: &binder::Strong<dyn IKeystoreService>,
+    sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
+    grant_key_nspace: i64,
+) -> Result<(), binder::Status> {
+    let key_entry_response = keystore2.getKeyEntry(&KeyDescriptor {
+        domain: Domain::GRANT,
+        nspace: grant_key_nspace,
+        alias: None,
+        blob: None,
+    })?;
+
+    // Perform sample crypto operation using granted key.
+    let op_response = sec_level.createOperation(
+        &key_entry_response.metadata.key,
+        &authorizations::AuthSetBuilder::new().purpose(KeyPurpose::SIGN).digest(Digest::SHA_2_256),
+        false,
+    )?;
+
+    assert!(op_response.iOperation.is_some());
+    assert_eq!(
+        Ok(()),
+        key_generations::map_ks_error(perform_sample_sign_operation(
+            &op_response.iOperation.unwrap()
+        ))
+    );
+
+    Ok(())
+}
+
 /// Try to grant a key with permission that does not map to any of the `KeyPermission` values.
 /// An error is expected with values that does not map to set of permissions listed in
 /// `KeyPermission`.
@@ -203,3 +236,520 @@
         )
     };
 }
+
+/// Grant a key to the user with DELETE access. In grantee context load the key and delete it.
+/// Verify that grantee should succeed in deleting the granted key and in grantor context test
+/// should fail to find the key with error response `KEY_NOT_FOUND`.
+#[test]
+fn keystore2_grant_delete_key_success() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+    const USER_ID: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+    static ALIAS: &str = "ks_grant_key_delete_success";
+
+    // Generate a key and grant it to a user with DELETE permission.
+    let grant_key_nspace = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let access_vector = KeyPermission::DELETE.0;
+            let mut grant_keys = generate_ec_key_and_grant_to_users(
+                &keystore2,
+                &sec_level,
+                Some(ALIAS.to_string()),
+                vec![GRANTEE_UID.try_into().unwrap()],
+                access_vector,
+            )
+            .unwrap();
+
+            grant_keys.remove(0)
+        })
+    };
+
+    // Grantee context, delete the key.
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_UID),
+            Gid::from_raw(GRANTEE_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+                keystore2
+                    .deleteKey(&KeyDescriptor {
+                        domain: Domain::GRANT,
+                        nspace: grant_key_nspace,
+                        alias: None,
+                        blob: None,
+                    })
+                    .unwrap();
+            },
+        )
+    };
+
+    // Verify whether key got deleted in grantor's context.
+    unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), move || {
+            let keystore2_inst = get_keystore_service();
+            let result =
+                key_generations::map_ks_error(keystore2_inst.getKeyEntry(&KeyDescriptor {
+                    domain: Domain::APP,
+                    nspace: -1,
+                    alias: Some(ALIAS.to_string()),
+                    blob: None,
+                }));
+            assert!(result.is_err());
+            assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+        })
+    };
+}
+
+/// Grant a key to the user. In grantee context load the granted key and try to grant it to second
+/// user. Test should fail with a response code `PERMISSION_DENIED` to grant a key to second user
+/// from grantee context. Test should make sure second grantee should not have a access to granted
+/// key.
+#[test]
+fn keystore2_grant_key_fails_with_permission_denied() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+    const USER_ID: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+
+    const SEC_USER_ID: u32 = 98;
+    const SEC_APPLICATION_ID: u32 = 10001;
+    static SEC_GRANTEE_UID: u32 = SEC_USER_ID * AID_USER_OFFSET + SEC_APPLICATION_ID;
+    static SEC_GRANTEE_GID: u32 = SEC_GRANTEE_UID;
+
+    // Generate a key and grant it to a user with GET_INFO permission.
+    let grant_key_nspace = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let access_vector = KeyPermission::GET_INFO.0;
+            let alias = format!("ks_grant_perm_denied_key_{}", getuid());
+            let mut grant_keys = generate_ec_key_and_grant_to_users(
+                &keystore2,
+                &sec_level,
+                Some(alias),
+                vec![GRANTEE_UID.try_into().unwrap()],
+                access_vector,
+            )
+            .unwrap();
+
+            grant_keys.remove(0)
+        })
+    };
+
+    // Grantee context, load the granted key and try to grant it to `SEC_GRANTEE_UID` grantee.
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_UID),
+            Gid::from_raw(GRANTEE_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+                let access_vector = KeyPermission::GET_INFO.0;
+
+                let key_entry_response = keystore2
+                    .getKeyEntry(&KeyDescriptor {
+                        domain: Domain::GRANT,
+                        nspace: grant_key_nspace,
+                        alias: None,
+                        blob: None,
+                    })
+                    .unwrap();
+
+                let result = key_generations::map_ks_error(keystore2.grant(
+                    &key_entry_response.metadata.key,
+                    SEC_GRANTEE_UID.try_into().unwrap(),
+                    access_vector,
+                ));
+                assert!(result.is_err());
+                assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+            },
+        )
+    };
+
+    // Make sure second grantee shouldn't have access to the above granted key.
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(SEC_GRANTEE_UID),
+            Gid::from_raw(SEC_GRANTEE_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+
+                let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
+                    domain: Domain::GRANT,
+                    nspace: grant_key_nspace,
+                    alias: None,
+                    blob: None,
+                }));
+
+                assert!(result.is_err());
+                assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+            },
+        )
+    };
+}
+
+/// Try to grant a key with `GRANT` access. Keystore2 system shouldn't allow to grant a key with
+/// `GRANT` access. Test should fail to grant a key with `PERMISSION_DENIED` error response code.
+#[test]
+fn keystore2_grant_key_fails_with_grant_perm_expect_perm_denied() {
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+    let access_vector = KeyPermission::GRANT.0;
+    let alias = format!("ks_grant_access_vec_key_{}", getuid());
+    let user_id = 98;
+    let application_id = 10001;
+    let grantee_uid = user_id * AID_USER_OFFSET + application_id;
+
+    let result = key_generations::map_ks_error(generate_ec_key_and_grant_to_users(
+        &keystore2,
+        &sec_level,
+        Some(alias),
+        vec![grantee_uid.try_into().unwrap()],
+        access_vector,
+    ));
+    assert!(result.is_err());
+    assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+}
+
+/// Try to grant a non-existing key to the user. Test should fail with `KEY_NOT_FOUND` error
+/// response.
+#[test]
+fn keystore2_grant_fails_with_non_existing_key_expect_key_not_found_err() {
+    let keystore2 = get_keystore_service();
+    let alias = format!("ks_grant_test_non_existing_key_5_{}", getuid());
+    let user_id = 98;
+    let application_id = 10001;
+    let grantee_uid = user_id * AID_USER_OFFSET + application_id;
+    let access_vector = KeyPermission::GET_INFO.0;
+
+    let result = key_generations::map_ks_error(keystore2.grant(
+        &KeyDescriptor {
+            domain: Domain::SELINUX,
+            nspace: key_generations::SELINUX_SHELL_NAMESPACE,
+            alias: Some(alias),
+            blob: None,
+        },
+        grantee_uid.try_into().unwrap(),
+        access_vector,
+    ));
+    assert!(result.is_err());
+    assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+}
+
+/// Grant a key to the user and immediately ungrant the granted key. In grantee context try to load
+/// the key. Grantee should fail to load the ungranted key with `KEY_NOT_FOUND` error response.
+#[test]
+fn keystore2_ungrant_key_success() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+    const USER_ID: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+
+    // Generate a key and grant it to a user with GET_INFO permission.
+    let grant_key_nspace = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let alias = format!("ks_ungrant_test_key_1{}", getuid());
+            let access_vector = KeyPermission::GET_INFO.0;
+            let mut grant_keys = generate_ec_key_and_grant_to_users(
+                &keystore2,
+                &sec_level,
+                Some(alias.to_string()),
+                vec![GRANTEE_UID.try_into().unwrap()],
+                access_vector,
+            )
+            .unwrap();
+
+            let grant_key_nspace = grant_keys.remove(0);
+
+            //Ungrant above granted key.
+            keystore2
+                .ungrant(
+                    &KeyDescriptor {
+                        domain: Domain::APP,
+                        nspace: -1,
+                        alias: Some(alias),
+                        blob: None,
+                    },
+                    GRANTEE_UID.try_into().unwrap(),
+                )
+                .unwrap();
+
+            grant_key_nspace
+        })
+    };
+
+    // Grantee context, try to load the ungranted key.
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_UID),
+            Gid::from_raw(GRANTEE_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+                let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
+                    domain: Domain::GRANT,
+                    nspace: grant_key_nspace,
+                    alias: None,
+                    blob: None,
+                }));
+                assert!(result.is_err());
+                assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+            },
+        )
+    };
+}
+
+/// Generate a key, grant it to the user and then delete the granted key. Try to ungrant
+/// a deleted key. Test should fail to ungrant a non-existing key with `KEY_NOT_FOUND` error
+/// response. Generate a new key with the same alias and try to access the previously granted
+/// key in grantee context. Test should fail to load the granted key in grantee context as the
+/// associated key is deleted from grantor context.
+#[test]
+fn keystore2_ungrant_fails_with_non_existing_key_expect_key_not_found_error() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+
+    const APPLICATION_ID: u32 = 10001;
+    const USER_ID: u32 = 99;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+
+    let grant_key_nspace = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let alias = format!("{}{}", "ks_grant_delete_ungrant_test_key_1", getuid());
+
+            let key_metadata = key_generations::generate_ec_p256_signing_key(
+                &sec_level,
+                Domain::SELINUX,
+                key_generations::SELINUX_SHELL_NAMESPACE,
+                Some(alias.to_string()),
+                None,
+            )
+            .unwrap();
+
+            let access_vector = KeyPermission::GET_INFO.0;
+            let grant_key = keystore2
+                .grant(&key_metadata.key, GRANTEE_UID.try_into().unwrap(), access_vector)
+                .unwrap();
+            assert_eq!(grant_key.domain, Domain::GRANT);
+
+            // Delete above granted key.
+            keystore2.deleteKey(&key_metadata.key).unwrap();
+
+            // Try to ungrant above granted key.
+            let result = key_generations::map_ks_error(
+                keystore2.ungrant(&key_metadata.key, GRANTEE_UID.try_into().unwrap()),
+            );
+            assert!(result.is_err());
+            assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+
+            // Generate a new key with the same alias and try to access the earlier granted key
+            // in grantee context.
+            let result = key_generations::generate_ec_p256_signing_key(
+                &sec_level,
+                Domain::SELINUX,
+                key_generations::SELINUX_SHELL_NAMESPACE,
+                Some(alias),
+                None,
+            );
+            assert!(result.is_ok());
+
+            grant_key.nspace
+        })
+    };
+
+    // Make sure grant did not persist, try to access the earlier granted key in grantee context.
+    // Grantee context should fail to load the granted key as its associated key is deleted in
+    // grantor context.
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_UID),
+            Gid::from_raw(GRANTEE_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+
+                let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
+                    domain: Domain::GRANT,
+                    nspace: grant_key_nspace,
+                    alias: None,
+                    blob: None,
+                }));
+                assert!(result.is_err());
+                assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+            },
+        )
+    };
+}
+
+/// Grant a key to multiple users. Verify that all grantees should succeed in loading the key and
+/// use it for performing an operation successfully.
+#[test]
+fn keystore2_grant_key_to_multi_users_success() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+
+    const APPLICATION_ID: u32 = 10001;
+    const USER_ID_1: u32 = 99;
+    static GRANTEE_1_UID: u32 = USER_ID_1 * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_1_GID: u32 = GRANTEE_1_UID;
+
+    const USER_ID_2: u32 = 98;
+    static GRANTEE_2_UID: u32 = USER_ID_2 * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_2_GID: u32 = GRANTEE_2_UID;
+
+    // Generate a key and grant it to multiple users with GET_INFO|USE permissions.
+    let mut grant_keys = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let alias = format!("ks_grant_test_key_2{}", getuid());
+            let access_vector = KeyPermission::GET_INFO.0 | KeyPermission::USE.0;
+
+            generate_ec_key_and_grant_to_users(
+                &keystore2,
+                &sec_level,
+                Some(alias),
+                vec![GRANTEE_1_UID.try_into().unwrap(), GRANTEE_2_UID.try_into().unwrap()],
+                access_vector,
+            )
+            .unwrap()
+        })
+    };
+
+    for (grantee_uid, grantee_gid) in
+        &[(GRANTEE_1_UID, GRANTEE_1_GID), (GRANTEE_2_UID, GRANTEE_2_GID)]
+    {
+        let grant_key_nspace = grant_keys.remove(0);
+        unsafe {
+            run_as::run_as(
+                GRANTEE_CTX,
+                Uid::from_raw(*grantee_uid),
+                Gid::from_raw(*grantee_gid),
+                move || {
+                    let keystore2 = get_keystore_service();
+                    let sec_level =
+                        keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+
+                    assert_eq!(
+                        Ok(()),
+                        key_generations::map_ks_error(load_grant_key_and_perform_sign_operation(
+                            &keystore2,
+                            &sec_level,
+                            grant_key_nspace
+                        ))
+                    );
+                },
+            )
+        };
+    }
+}
+
+/// Grant a key to multiple users with GET_INFO|DELETE permissions. In one of the grantee context
+/// use the key and delete it. Try to load the granted key in another grantee context. Test should
+/// fail to load the granted key with `KEY_NOT_FOUND` error response.
+#[test]
+fn keystore2_grant_key_to_multi_users_delete_fails_with_key_not_found_error() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+
+    const USER_ID_1: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_1_UID: u32 = USER_ID_1 * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_1_GID: u32 = GRANTEE_1_UID;
+
+    const USER_ID_2: u32 = 98;
+    static GRANTEE_2_UID: u32 = USER_ID_2 * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_2_GID: u32 = GRANTEE_2_UID;
+
+    // Generate a key and grant it to multiple users with GET_INFO permission.
+    let mut grant_keys = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let alias = format!("ks_grant_test_key_2{}", getuid());
+            let access_vector =
+                KeyPermission::GET_INFO.0 | KeyPermission::USE.0 | KeyPermission::DELETE.0;
+
+            generate_ec_key_and_grant_to_users(
+                &keystore2,
+                &sec_level,
+                Some(alias),
+                vec![GRANTEE_1_UID.try_into().unwrap(), GRANTEE_2_UID.try_into().unwrap()],
+                access_vector,
+            )
+            .unwrap()
+        })
+    };
+
+    // Grantee #1 context
+    let grant_key1_nspace = grant_keys.remove(0);
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_1_UID),
+            Gid::from_raw(GRANTEE_1_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+                let sec_level =
+                    keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+
+                assert_eq!(
+                    Ok(()),
+                    key_generations::map_ks_error(load_grant_key_and_perform_sign_operation(
+                        &keystore2,
+                        &sec_level,
+                        grant_key1_nspace
+                    ))
+                );
+
+                // Delete the granted key.
+                keystore2
+                    .deleteKey(&KeyDescriptor {
+                        domain: Domain::GRANT,
+                        nspace: grant_key1_nspace,
+                        alias: None,
+                        blob: None,
+                    })
+                    .unwrap();
+            },
+        )
+    };
+
+    // Grantee #2 context
+    let grant_key2_nspace = grant_keys.remove(0);
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_2_UID),
+            Gid::from_raw(GRANTEE_2_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+
+                let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
+                    domain: Domain::GRANT,
+                    nspace: grant_key2_nspace,
+                    alias: None,
+                    blob: None,
+                }));
+                assert!(result.is_err());
+                assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+            },
+        )
+    };
+}
diff --git a/keystore2/tests/keystore2_client_test_utils.rs b/keystore2/tests/keystore2_client_test_utils.rs
index 56995e4..58e6b7d 100644
--- a/keystore2/tests/keystore2_client_test_utils.rs
+++ b/keystore2/tests/keystore2_client_test_utils.rs
@@ -90,6 +90,29 @@
         .expect("Could not check for declared keymint interface")
 }
 
+/// Generate EC key and grant it to the list of users with given access vector.
+/// Returns the list of granted keys `nspace` values in the order of given grantee uids.
+pub fn generate_ec_key_and_grant_to_users(
+    keystore2: &binder::Strong<dyn IKeystoreService>,
+    sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
+    alias: Option<String>,
+    grantee_uids: Vec<i32>,
+    access_vector: i32,
+) -> Result<Vec<i64>, binder::Status> {
+    let key_metadata =
+        key_generations::generate_ec_p256_signing_key(sec_level, Domain::APP, -1, alias, None)?;
+
+    let mut granted_keys = Vec::new();
+
+    for uid in grantee_uids {
+        let granted_key = keystore2.grant(&key_metadata.key, uid, access_vector)?;
+        assert_eq!(granted_key.domain, Domain::GRANT);
+        granted_keys.push(granted_key.nspace);
+    }
+
+    Ok(granted_keys)
+}
+
 /// Generate a EC_P256 key using given domain, namespace and alias.
 /// Create an operation using the generated key and perform sample signing operation.
 pub fn create_signing_operation(
diff --git a/keystore2/tests/keystore2_client_tests.rs b/keystore2/tests/keystore2_client_tests.rs
index d705aa4..45c5fb7 100644
--- a/keystore2/tests/keystore2_client_tests.rs
+++ b/keystore2/tests/keystore2_client_tests.rs
@@ -16,6 +16,7 @@
 pub mod keystore2_client_3des_key_tests;
 pub mod keystore2_client_aes_key_tests;
 pub mod keystore2_client_attest_key_tests;
+pub mod keystore2_client_delete_key_tests;
 pub mod keystore2_client_ec_key_tests;
 pub mod keystore2_client_grant_key_tests;
 pub mod keystore2_client_hmac_key_tests;
@@ -25,3 +26,4 @@
 pub mod keystore2_client_operation_tests;
 pub mod keystore2_client_rsa_key_tests;
 pub mod keystore2_client_test_utils;
+pub mod keystore2_client_update_subcomponent_tests;
diff --git a/keystore2/tests/keystore2_client_update_subcomponent_tests.rs b/keystore2/tests/keystore2_client_update_subcomponent_tests.rs
new file mode 100644
index 0000000..c987f22
--- /dev/null
+++ b/keystore2/tests/keystore2_client_update_subcomponent_tests.rs
@@ -0,0 +1,230 @@
+// Copyright 2022, 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.
+
+use nix::unistd::{getuid, Gid, Uid};
+use rustutils::users::AID_USER_OFFSET;
+
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    ErrorCode::ErrorCode, SecurityLevel::SecurityLevel,
+};
+use android_system_keystore2::aidl::android::system::keystore2::{
+    Domain::Domain, KeyDescriptor::KeyDescriptor, KeyPermission::KeyPermission,
+    ResponseCode::ResponseCode,
+};
+
+use keystore2_test_utils::{get_keystore_service, key_generations, key_generations::Error, run_as};
+
+/// Generate a key and update its public certificate and certificate chain. Test should be able to
+/// load the key and able to verify whether its certificate and cert-chain are updated successfully.
+#[test]
+fn keystore2_update_subcomponent_success() {
+    let alias = "update_subcomponent_success_key";
+
+    let keystore2 = get_keystore_service();
+    let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+
+    let key_metadata = key_generations::generate_ec_p256_signing_key(
+        &sec_level,
+        Domain::SELINUX,
+        key_generations::SELINUX_SHELL_NAMESPACE,
+        Some(alias.to_string()),
+        None,
+    )
+    .unwrap();
+
+    let other_cert: [u8; 32] = [123; 32];
+    let other_cert_chain: [u8; 32] = [12; 32];
+
+    keystore2
+        .updateSubcomponent(&key_metadata.key, Some(&other_cert), Some(&other_cert_chain))
+        .expect("updateSubcomponent should have succeeded.");
+
+    let key_entry_response = keystore2.getKeyEntry(&key_metadata.key).unwrap();
+    assert_eq!(Some(other_cert.to_vec()), key_entry_response.metadata.certificate);
+    assert_eq!(Some(other_cert_chain.to_vec()), key_entry_response.metadata.certificateChain);
+}
+
+/// Try to update non-existing asymmetric key public cert and certificate chain. Test should fail
+/// to update with error response code `KEY_NOT_FOUND`.
+#[test]
+fn keystore2_update_subcomponent_fail() {
+    let alias = "update_component_failure_key";
+
+    let keystore2 = get_keystore_service();
+
+    let other_cert: [u8; 32] = [123; 32];
+    let other_cert_chain: [u8; 32] = [12; 32];
+
+    let result = key_generations::map_ks_error(keystore2.updateSubcomponent(
+        &KeyDescriptor {
+            domain: Domain::SELINUX,
+            nspace: key_generations::SELINUX_SHELL_NAMESPACE,
+            alias: Some(alias.to_string()),
+            blob: None,
+        },
+        Some(&other_cert),
+        Some(&other_cert_chain),
+    ));
+    assert!(result.is_err());
+    assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+}
+
+/// Generate a key and grant it to two users. For one user grant it with only `GET_INFO` access
+/// permission and for another user grant it with GET_INFO and UPDATE access permissions. In a
+/// grantee context where key is granted with only GET_INFO access permission, try to update
+/// key's public certificate and certificate chain. Test should fail to update with error response
+/// code `PERMISSION_DENIED` because grantee does not possess UPDATE access permission for the
+/// specified key. In a grantee context where key is granted with UPDATE and GET_INFO access
+/// permissions, test should be able to update public certificate and cert-chain successfully.
+#[test]
+fn keystore2_update_subcomponent_fails_permission_denied() {
+    static GRANTOR_SU_CTX: &str = "u:r:su:s0";
+    static GRANTEE_CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
+
+    const USER_ID_1: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_1_UID: u32 = USER_ID_1 * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_1_GID: u32 = GRANTEE_1_UID;
+
+    const USER_ID_2: u32 = 98;
+    static GRANTEE_2_UID: u32 = USER_ID_2 * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_2_GID: u32 = GRANTEE_2_UID;
+
+    // Generate a key and grant it to multiple users with different access permissions.
+    let mut granted_keys = unsafe {
+        run_as::run_as(GRANTOR_SU_CTX, Uid::from_raw(0), Gid::from_raw(0), || {
+            let keystore2 = get_keystore_service();
+            let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
+            let alias = format!("ks_update_subcompo_test_1_{}", getuid());
+            let mut granted_keys = Vec::new();
+
+            let key_metadata = key_generations::generate_ec_p256_signing_key(
+                &sec_level,
+                Domain::APP,
+                -1,
+                Some(alias),
+                None,
+            )
+            .unwrap();
+
+            // Grant a key without update permission.
+            let access_vector = KeyPermission::GET_INFO.0;
+            let granted_key = keystore2
+                .grant(&key_metadata.key, GRANTEE_1_UID.try_into().unwrap(), access_vector)
+                .unwrap();
+            assert_eq!(granted_key.domain, Domain::GRANT);
+            granted_keys.push(granted_key.nspace);
+
+            // Grant a key with update permission.
+            let access_vector = KeyPermission::GET_INFO.0 | KeyPermission::UPDATE.0;
+            let granted_key = keystore2
+                .grant(&key_metadata.key, GRANTEE_2_UID.try_into().unwrap(), access_vector)
+                .unwrap();
+            assert_eq!(granted_key.domain, Domain::GRANT);
+            granted_keys.push(granted_key.nspace);
+
+            granted_keys
+        })
+    };
+
+    // Grantee context, try to update the key public certs, permission denied error is expected.
+    let granted_key1_nspace = granted_keys.remove(0);
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_1_UID),
+            Gid::from_raw(GRANTEE_1_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+
+                let other_cert: [u8; 32] = [123; 32];
+                let other_cert_chain: [u8; 32] = [12; 32];
+
+                let result = key_generations::map_ks_error(keystore2.updateSubcomponent(
+                    &KeyDescriptor {
+                        domain: Domain::GRANT,
+                        nspace: granted_key1_nspace,
+                        alias: None,
+                        blob: None,
+                    },
+                    Some(&other_cert),
+                    Some(&other_cert_chain),
+                ));
+                assert!(result.is_err());
+                assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+            },
+        )
+    };
+
+    // Grantee context, update granted key public certs. Update should happen successfully.
+    let granted_key2_nspace = granted_keys.remove(0);
+    unsafe {
+        run_as::run_as(
+            GRANTEE_CTX,
+            Uid::from_raw(GRANTEE_2_UID),
+            Gid::from_raw(GRANTEE_2_GID),
+            move || {
+                let keystore2 = get_keystore_service();
+
+                let other_cert: [u8; 32] = [124; 32];
+                let other_cert_chain: [u8; 32] = [13; 32];
+
+                keystore2
+                    .updateSubcomponent(
+                        &KeyDescriptor {
+                            domain: Domain::GRANT,
+                            nspace: granted_key2_nspace,
+                            alias: None,
+                            blob: None,
+                        },
+                        Some(&other_cert),
+                        Some(&other_cert_chain),
+                    )
+                    .expect("updateSubcomponent should have succeeded.");
+
+                let key_entry_response = keystore2
+                    .getKeyEntry(&KeyDescriptor {
+                        domain: Domain::GRANT,
+                        nspace: granted_key2_nspace,
+                        alias: None,
+                        blob: None,
+                    })
+                    .unwrap();
+                assert_eq!(Some(other_cert.to_vec()), key_entry_response.metadata.certificate);
+                assert_eq!(
+                    Some(other_cert_chain.to_vec()),
+                    key_entry_response.metadata.certificateChain
+                );
+            },
+        )
+    };
+}
+
+#[test]
+fn keystore2_get_security_level_success() {
+    let keystore2 = get_keystore_service();
+    assert!(
+        keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).is_ok(),
+        "getSecurityLevel with SecurityLevel::TRUSTED_ENVIRONMENT should have succeeded."
+    );
+}
+
+#[test]
+fn keystore2_get_security_level_failure() {
+    let keystore2 = get_keystore_service();
+    let result = key_generations::map_ks_error(keystore2.getSecurityLevel(SecurityLevel::SOFTWARE));
+
+    assert!(result.is_err());
+    assert_eq!(Error::Km(ErrorCode::HARDWARE_TYPE_UNAVAILABLE), result.unwrap_err());
+}
diff --git a/prng_seeder/Android.bp b/prng_seeder/Android.bp
index 292535a..763aaa0 100644
--- a/prng_seeder/Android.bp
+++ b/prng_seeder/Android.bp
@@ -63,3 +63,20 @@
     installable: false,
     prefer_rlib: true,
 }
+
+rust_test {
+    name: "prng_seeder.test",
+    edition: "2021",
+    srcs: ["src/main.rs"],
+    rustlibs: [
+        "libanyhow",
+        "libbssl_ffi",
+        "libclap",
+        "libcutils_socket_bindgen",
+        "liblogger",
+        "liblog_rust",
+        "libnix",
+        "libtokio",
+    ],
+    test_suites: ["general-tests"],
+}
diff --git a/prng_seeder/src/main.rs b/prng_seeder/src/main.rs
index 3f698f6..924481a 100644
--- a/prng_seeder/src/main.rs
+++ b/prng_seeder/src/main.rs
@@ -37,7 +37,7 @@
 
 use crate::conditioner::ConditionerBuilder;
 
-#[derive(Debug, clap::Parser)]
+#[derive(Debug, Parser)]
 struct Cli {
     #[clap(long, default_value = "/dev/hw_random")]
     source: PathBuf,
@@ -135,3 +135,14 @@
     println!("prng_seeder: launch terminated: {:?}", e);
     std::process::exit(-1);
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    #[test]
+    fn verify_cli() {
+        Cli::command().debug_assert();
+    }
+}