diff --git a/keystore/keystore_cli_v2.cpp b/keystore/keystore_cli_v2.cpp
index d01c67d..ab3e22c 100644
--- a/keystore/keystore_cli_v2.cpp
+++ b/keystore/keystore_cli_v2.cpp
@@ -59,6 +59,16 @@
 
 constexpr const char keystore2_service_name[] = "android.system.keystore2.IKeystoreService/default";
 
+std::string string_replace_all(std::string str, const std::string& from,
+                                      const std::string& to) {
+    size_t start = 0;
+    while ((start = str.find(from, start)) != std::string::npos) {
+        str.replace(start, from.length(), to);
+        start += to.length();
+    }
+    return str;
+}
+
 int unwrapError(const ndk::ScopedAStatus& status) {
     if (status.isOk()) return 0;
     if (status.getExceptionCode() == EX_SERVICE_SPECIFIC) {
@@ -1028,7 +1038,8 @@
     auto listener = ndk::SharedRefBase::make<ConfirmationListener>();
 
     auto future = listener->get_future();
-    auto rc = apcService->presentPrompt(listener, promptText, extraData, locale, uiOptionsAsFlags);
+    auto rc = apcService->presentPrompt(listener, string_replace_all(promptText, "\\n", "\n"),
+                                        extraData, locale, uiOptionsAsFlags);
 
     if (!rc.isOk()) {
         std::cerr << "Presenting confirmation prompt failed: " << rc.getDescription() << std::endl;
diff --git a/keystore2/Android.bp b/keystore2/Android.bp
index fa80563..57038df 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",
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/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/fuzzers/keystore2_unsafe_fuzzer.rs b/keystore2/src/fuzzers/keystore2_unsafe_fuzzer.rs
index 4c2419a..1a385e7 100644
--- a/keystore2/src/fuzzers/keystore2_unsafe_fuzzer.rs
+++ b/keystore2/src/fuzzers/keystore2_unsafe_fuzzer.rs
@@ -16,8 +16,6 @@
 
 #![feature(slice_internals)]
 #![no_main]
-#[macro_use]
-extern crate libfuzzer_sys;
 
 use core::slice::memchr;
 use keystore2::{legacy_blob::LegacyBlobLoader, utils::ui_opts_2_compat};
@@ -31,7 +29,7 @@
 };
 use keystore2_selinux::{check_access, getpidcon, setcon, Backend, Context, KeystoreKeyBackend};
 use keystore2_vintf::{get_aidl_instances, get_hidl_instances};
-use libfuzzer_sys::arbitrary::Arbitrary;
+use libfuzzer_sys::{arbitrary::Arbitrary, fuzz_target};
 use std::{ffi::CString, sync::Arc};
 
 // Avoid allocating too much memory and crashing the fuzzer.
diff --git a/keystore2/src/rkpd_client.rs b/keystore2/src/rkpd_client.rs
index 2d1b23b..c4b0686 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, map_or_log_err, 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;
@@ -23,53 +24,58 @@
     IGetKeyCallback::BnGetKeyCallback, IGetKeyCallback::IGetKeyCallback,
     IGetRegistrationCallback::BnGetRegistrationCallback,
     IGetRegistrationCallback::IGetRegistrationCallback, IRegistration::IRegistration,
-    IRemoteProvisioning::IRemoteProvisioning, RemotelyProvisionedKey::RemotelyProvisionedKey,
+    IRemoteProvisioning::IRemoteProvisioning,
+    IStoreUpgradedKeyCallback::BnStoreUpgradedKeyCallback,
+    IStoreUpgradedKeyCallback::IStoreUpgradedKeyCallback,
+    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;
 
-type RegistrationSender = oneshot::Sender<Result<binder::Strong<dyn IRegistration>>>;
+// 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.
+struct SafeSender<T> {
+    inner: Mutex<Option<oneshot::Sender<T>>>,
+}
+
+impl<T> SafeSender<T> {
+    fn new(sender: oneshot::Sender<T>) -> Self {
+        Self { inner: Mutex::new(Some(sender)) }
+    }
+
+    fn send(&self, value: T) {
+        if let Some(inner) = self.inner.lock().unwrap().take() {
+            // assert instead of unwrap, because on failure send returns Err(value)
+            assert!(inner.send(value).is_ok(), "thread state is terminally broken");
+        }
+    }
+}
 
 struct GetRegistrationCallback {
-    registration_tx: Mutex<Option<RegistrationSender>>,
+    registration_tx: SafeSender<Result<binder::Strong<dyn IRegistration>>>,
 }
 
 impl GetRegistrationCallback {
     pub fn new_native_binder(
-        registration_tx: RegistrationSender,
-    ) -> Result<Strong<dyn IGetRegistrationCallback>> {
+        registration_tx: oneshot::Sender<Result<binder::Strong<dyn IRegistration>>>,
+    ) -> Strong<dyn IGetRegistrationCallback> {
         let result: Self =
-            GetRegistrationCallback { registration_tx: Mutex::new(Some(registration_tx)) };
-        Ok(BnGetRegistrationCallback::new_binder(result, BinderFeatures::default()))
-    }
-    fn on_success(&self, registration: &binder::Strong<dyn IRegistration>) -> Result<()> {
-        if let Some(tx) = self.registration_tx.lock().unwrap().take() {
-            tx.send(Ok(registration.clone())).unwrap();
-        }
-        Ok(())
-    }
-    fn on_cancel(&self) -> Result<()> {
-        if let Some(tx) = self.registration_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
-                    .context(ks_err!("GetRegistrationCallback cancelled.")),
-            )
-            .unwrap();
-        }
-        Ok(())
-    }
-    fn on_error(&self, error: &str) -> Result<()> {
-        if let Some(tx) = self.registration_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
-                    .context(ks_err!("GetRegistrationCallback failed: {:?}", error)),
-            )
-            .unwrap();
-        }
-        Ok(())
+            GetRegistrationCallback { registration_tx: SafeSender::new(registration_tx) };
+        BnGetRegistrationCallback::new_binder(result, BinderFeatures::default())
     }
 }
 
@@ -78,15 +84,26 @@
 impl IGetRegistrationCallback for GetRegistrationCallback {
     fn onSuccess(&self, registration: &Strong<dyn IRegistration>) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetRegistrationCallback::onSuccess", 500);
-        map_or_log_err(self.on_success(registration), Ok)
+        self.registration_tx.send(Ok(registration.clone()));
+        Ok(())
     }
     fn onCancel(&self) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetRegistrationCallback::onCancel", 500);
-        map_or_log_err(self.on_cancel(), Ok)
+        log::warn!("IGetRegistrationCallback cancelled");
+        self.registration_tx.send(
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS))
+                .context(ks_err!("GetRegistrationCallback cancelled.")),
+        );
+        Ok(())
     }
     fn onError(&self, error: &str) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetRegistrationCallback::onError", 500);
-        map_or_log_err(self.on_error(error), Ok)
+        log::error!("IGetRegistrationCallback failed: '{error}'");
+        self.registration_tx.send(
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS))
+                .context(ks_err!("GetRegistrationCallback failed: {:?}", error)),
+        );
+        Ok(())
     }
 }
 
@@ -102,56 +119,30 @@
         .context(ks_err!("Trying to get IRPC name."))?;
 
     let (tx, rx) = oneshot::channel();
-    let cb = GetRegistrationCallback::new_native_binder(tx)
-        .context(ks_err!("Trying to create a IGetRegistrationCallback."))?;
+    let cb = GetRegistrationCallback::new_native_binder(tx);
 
     remote_provisioning
         .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(),
+    }
 }
 
-type KeySender = oneshot::Sender<Result<RemotelyProvisionedKey>>;
-
 struct GetKeyCallback {
-    key_tx: Mutex<Option<KeySender>>,
+    key_tx: SafeSender<Result<RemotelyProvisionedKey>>,
 }
 
 impl GetKeyCallback {
-    pub fn new_native_binder(key_tx: KeySender) -> Result<Strong<dyn IGetKeyCallback>> {
-        let result: Self = GetKeyCallback { key_tx: Mutex::new(Some(key_tx)) };
-        Ok(BnGetKeyCallback::new_binder(result, BinderFeatures::default()))
-    }
-    fn on_success(&self, key: &RemotelyProvisionedKey) -> Result<()> {
-        if let Some(tx) = self.key_tx.lock().unwrap().take() {
-            tx.send(Ok(RemotelyProvisionedKey {
-                keyBlob: key.keyBlob.clone(),
-                encodedCertChain: key.encodedCertChain.clone(),
-            }))
-            .unwrap();
-        }
-        Ok(())
-    }
-    fn on_cancel(&self) -> Result<()> {
-        if let Some(tx) = self.key_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
-                    .context(ks_err!("GetKeyCallback cancelled.")),
-            )
-            .unwrap();
-        }
-        Ok(())
-    }
-    fn on_error(&self, error: &str) -> Result<()> {
-        if let Some(tx) = self.key_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
-                    .context(ks_err!("GetKeyCallback failed: {:?}", error)),
-            )
-            .unwrap();
-        }
-        Ok(())
+    pub fn new_native_binder(
+        key_tx: oneshot::Sender<Result<RemotelyProvisionedKey>>,
+    ) -> Strong<dyn IGetKeyCallback> {
+        let result: Self = GetKeyCallback { key_tx: SafeSender::new(key_tx) };
+        BnGetKeyCallback::new_binder(result, BinderFeatures::default())
     }
 }
 
@@ -160,15 +151,46 @@
 impl IGetKeyCallback for GetKeyCallback {
     fn onSuccess(&self, key: &RemotelyProvisionedKey) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetKeyCallback::onSuccess", 500);
-        map_or_log_err(self.on_success(key), Ok)
+        self.key_tx.send(Ok(RemotelyProvisionedKey {
+            keyBlob: key.keyBlob.clone(),
+            encodedCertChain: key.encodedCertChain.clone(),
+        }));
+        Ok(())
     }
     fn onCancel(&self) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetKeyCallback::onCancel", 500);
-        map_or_log_err(self.on_cancel(), Ok)
+        log::warn!("IGetKeyCallback cancelled");
+        self.key_tx.send(
+            Err(Error::Rc(ResponseCode::OUT_OF_KEYS)).context(ks_err!("GetKeyCallback cancelled.")),
+        );
+        Ok(())
     }
     fn onError(&self, error: &str) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetKeyCallback::onError", 500);
-        map_or_log_err(self.on_error(error), Ok)
+        log::error!("IGetKeyCallback failed: {error}");
+        self.key_tx.send(
+            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(),
     }
 }
 
@@ -179,16 +201,59 @@
     let registration = get_rkpd_registration(security_level)
         .await
         .context(ks_err!("Trying to get to IRegistration service."))?;
+    get_rkpd_attestation_key_from_registration_async(&registration, caller_uid).await
+}
 
+struct StoreUpgradedKeyCallback {
+    completer: SafeSender<Result<()>>,
+}
+
+impl StoreUpgradedKeyCallback {
+    pub fn new_native_binder(
+        completer: oneshot::Sender<Result<()>>,
+    ) -> Strong<dyn IStoreUpgradedKeyCallback> {
+        let result: Self = StoreUpgradedKeyCallback { completer: SafeSender::new(completer) };
+        BnStoreUpgradedKeyCallback::new_binder(result, BinderFeatures::default())
+    }
+}
+
+impl Interface for StoreUpgradedKeyCallback {}
+
+impl IStoreUpgradedKeyCallback for StoreUpgradedKeyCallback {
+    fn onSuccess(&self) -> binder::Result<()> {
+        let _wp = wd::watch_millis("IGetRegistrationCallback::onSuccess", 500);
+        self.completer.send(Ok(()));
+        Ok(())
+    }
+
+    fn onError(&self, error: &str) -> binder::Result<()> {
+        let _wp = wd::watch_millis("IGetRegistrationCallback::onError", 500);
+        log::error!("IGetRegistrationCallback failed: {error}");
+        self.completer.send(
+            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 = GetKeyCallback::new_native_binder(tx)
-        .context(ks_err!("Trying to create a IGetKeyCallback."))?;
+    let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
 
     registration
-        .getKey(caller_uid.try_into().unwrap(), &cb)
-        .context(ks_err!("Trying to get key."))?;
+        .storeUpgradedKeyAsync(key_blob, upgraded_blob, &cb)
+        .context(ks_err!("Failed to store upgraded blob with RKPD."))?;
 
-    rx.await.unwrap()
+    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(
@@ -199,11 +264,7 @@
     let registration = get_rkpd_registration(security_level)
         .await
         .context(ks_err!("Trying to get to IRegistration service."))?;
-
-    registration
-        .storeUpgradedKey(key_blob, upgraded_blob)
-        .context(ks_err!("Failed to store upgraded blob with RKPD."))?;
-    Ok(())
+    store_rkpd_attestation_key_with_registration_async(&registration, key_blob, upgraded_blob).await
 }
 
 /// Get attestation key from RKPD.
@@ -212,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.
@@ -222,26 +283,33 @@
     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 android_security_rkp_aidl::aidl::android::security::rkp::IRegistration::BnRegistration;
-    use std::sync::Arc;
+    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())
         }
     }
@@ -249,57 +317,98 @@
     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<()> {
             todo!()
         }
 
-        fn storeUpgradedKey(&self, _: &[u8], _: &[u8]) -> binder::Result<()> {
-            todo!()
+        fn storeUpgradedKeyAsync(
+            &self,
+            _: &[u8],
+            _: &[u8],
+            cb: &Strong<dyn IStoreUpgradedKeyCallback>,
+        ) -> binder::Result<()> {
+            // 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).unwrap();
-        let mock_registration = MockRegistration::new_native_binder();
+        let cb = GetRegistrationCallback::new_native_binder(tx);
+        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());
     }
 
     #[test]
     fn test_get_registration_cb_cancel() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx).unwrap();
+        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)
         );
     }
 
     #[test]
     fn test_get_registration_cb_error() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx).unwrap();
+        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)
         );
     }
 
@@ -308,80 +417,171 @@
         let mock_key =
             RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
         let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx).unwrap();
+        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);
     }
 
     #[test]
     fn test_get_key_cb_cancel() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx).unwrap();
+        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)
         );
     }
 
     #[test]
     fn test_get_key_cb_error() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx).unwrap();
+        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)
         );
     }
 
     #[test]
-    #[ignore]
+    fn test_store_upgraded_cb_success() {
+        let (tx, rx) = oneshot::channel();
+        let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+        assert!(cb.onSuccess().is_ok());
+
+        tokio_rt().block_on(rx).unwrap().unwrap();
+    }
+
+    #[test]
+    fn test_store_upgraded_key_cb_error() {
+        let (tx, rx) = oneshot::channel();
+        let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+        assert!(cb.onError("oh no! it failed").is_ok());
+
+        let result = tokio_rt().block_on(rx).unwrap();
+        assert_eq!(
+            result.unwrap_err().downcast::<Error>().unwrap(),
+            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() {
-        let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, 0).unwrap();
+        binder::ProcessState::start_thread_pool();
+        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());
     }
 
     #[test]
-    #[ignore]
     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);
     }
 
     #[test]
-    #[ignore]
     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]
-    #[ignore]
+    // 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);
     }
 }
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index 75d98e2..1bae75e 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -276,7 +276,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/test_utils/authorizations.rs b/keystore2/test_utils/authorizations.rs
index 7dcee83..4608bc5 100644
--- a/keystore2/test_utils/authorizations.rs
+++ b/keystore2/test_utils/authorizations.rs
@@ -71,15 +71,6 @@
         self
     }
 
-    /// Add Attestation-ID.
-    pub fn attestation_app_id(mut self, b: Vec<u8>) -> Self {
-        self.0.push(KeyParameter {
-            tag: Tag::ATTESTATION_APPLICATION_ID,
-            value: KeyParameterValue::Blob(b),
-        });
-        self
-    }
-
     /// Add No_auth_required.
     pub fn no_auth_required(mut self) -> Self {
         self.0.push(KeyParameter {
diff --git a/keystore2/test_utils/key_generations.rs b/keystore2/test_utils/key_generations.rs
index 53597af..f9aaabb 100644
--- a/keystore2/test_utils/key_generations.rs
+++ b/keystore2/test_utils/key_generations.rs
@@ -58,8 +58,6 @@
     pub block_mode: Option<BlockMode>,
     /// Attestation challenge.
     pub att_challenge: Option<Vec<u8>>,
-    /// Attestation app id.
-    pub att_app_id: Option<Vec<u8>>,
 }
 
 /// DER-encoded PKCS#8 format RSA key. Generated using:
@@ -338,7 +336,6 @@
     nspace: i64,
     alias: Option<String>,
     att_challenge: Option<&[u8]>,
-    att_app_id: Option<&[u8]>,
 ) -> binder::Result<KeyMetadata> {
     let mut key_attest = false;
     let mut gen_params = AuthSetBuilder::new()
@@ -354,11 +351,6 @@
         gen_params = gen_params.clone().attestation_challenge(challenge.to_vec());
     }
 
-    if let Some(app_id) = att_app_id {
-        key_attest = true;
-        gen_params = gen_params.clone().attestation_app_id(app_id.to_vec());
-    }
-
     match sec_level.generateKey(
         &KeyDescriptor { domain, nspace, alias, blob: None },
         None,
@@ -453,9 +445,6 @@
     if let Some(value) = &key_params.att_challenge {
         gen_params = gen_params.attestation_challenge(value.to_vec())
     }
-    if let Some(value) = &key_params.att_app_id {
-        gen_params = gen_params.attestation_app_id(value.to_vec())
-    }
 
     let key_metadata = sec_level.generateKey(
         &KeyDescriptor { domain, nspace, alias, blob: None },
@@ -468,8 +457,7 @@
     // Must have a public key.
     assert!(key_metadata.certificate.is_some());
 
-    if attest_key.is_none() && key_params.att_challenge.is_some() && key_params.att_app_id.is_some()
-    {
+    if attest_key.is_none() && key_params.att_challenge.is_some() {
         // Should have an attestation record.
         assert!(key_metadata.certificateChain.is_some());
     } else {
@@ -578,7 +566,6 @@
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
     algorithm: Algorithm,
     att_challenge: &[u8],
-    att_app_id: &[u8],
 ) -> binder::Result<KeyMetadata> {
     assert!(algorithm == Algorithm::RSA || algorithm == Algorithm::EC);
 
@@ -597,7 +584,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: Some(att_challenge.to_vec()),
-                att_app_id: Some(att_app_id.to_vec()),
             },
             None,
         )
@@ -607,7 +593,6 @@
         let metadata = generate_ec_attestation_key(
             sec_level,
             att_challenge,
-            att_app_id,
             Digest::SHA_2_256,
             EcCurve::P_256,
         )
@@ -622,7 +607,6 @@
 pub fn generate_ec_attestation_key(
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
     att_challenge: &[u8],
-    att_app_id: &[u8],
     digest: Digest,
     ec_curve: EcCurve,
 ) -> binder::Result<KeyMetadata> {
@@ -633,8 +617,7 @@
         .purpose(KeyPurpose::ATTEST_KEY)
         .ec_curve(ec_curve)
         .digest(digest)
-        .attestation_challenge(att_challenge.to_vec())
-        .attestation_app_id(att_app_id.to_vec());
+        .attestation_challenge(att_challenge.to_vec());
 
     let attestation_key_metadata = sec_level.generateKey(
         &KeyDescriptor {
@@ -662,7 +645,6 @@
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
     alias: Option<String>,
     att_challenge: &[u8],
-    att_app_id: &[u8],
     attest_key: &KeyDescriptor,
 ) -> binder::Result<KeyMetadata> {
     let ec_gen_params = AuthSetBuilder::new()
@@ -672,8 +654,7 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(att_challenge.to_vec())
-        .attestation_app_id(att_app_id.to_vec());
+        .attestation_challenge(att_challenge.to_vec());
 
     let ec_key_metadata = sec_level
         .generateKey(
diff --git a/keystore2/tests/ffi_test_utils.cpp b/keystore2/tests/ffi_test_utils.cpp
index 45ce02c..de20d83 100644
--- a/keystore2/tests/ffi_test_utils.cpp
+++ b/keystore2/tests/ffi_test_utils.cpp
@@ -237,7 +237,7 @@
     }
 
     // Perform ASN.1 DER encoding of KeyDescription.
-    size_t asn1_data_len = i2d_TEST_KEY_DESCRIPTION(key_description.get(), nullptr);
+    int asn1_data_len = i2d_TEST_KEY_DESCRIPTION(key_description.get(), nullptr);
     if (asn1_data_len < 0) {
         cxx_result.error = keymaster::TranslateLastOpenSslError();
         return cxx_result;
@@ -341,7 +341,7 @@
     }
 
     // ASN.1 DER-encoding of secure key wrapper.
-    size_t asn1_data_len = i2d_TEST_SECURE_KEY_WRAPPER(sec_key_wrapper.get(), nullptr);
+    int asn1_data_len = i2d_TEST_SECURE_KEY_WRAPPER(sec_key_wrapper.get(), nullptr);
     if (asn1_data_len < 0) {
         cxx_result.error = keymaster::TranslateLastOpenSslError();
         return cxx_result;
diff --git a/keystore2/tests/keystore2_client_aes_key_tests.rs b/keystore2/tests/keystore2_client_aes_key_tests.rs
index 885cbf5..313f596 100644
--- a/keystore2/tests/keystore2_client_aes_key_tests.rs
+++ b/keystore2/tests/keystore2_client_aes_key_tests.rs
@@ -26,8 +26,7 @@
 };
 
 use crate::keystore2_client_test_utils::{
-    has_trusty_keymint, perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op,
-    SAMPLE_PLAIN_TEXT,
+    perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op, SAMPLE_PLAIN_TEXT,
 };
 
 /// Generate a AES key. Create encrypt and decrypt operations using the generated key.
@@ -393,11 +392,11 @@
     ));
     assert!(result.is_err());
 
-    if has_trusty_keymint() {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::MISSING_MAC_LENGTH));
-    } else {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::UNSUPPORTED_MAC_LENGTH));
-    }
+    let e = result.unwrap_err();
+    assert!(
+        e == Error::Km(ErrorCode::MISSING_MAC_LENGTH)
+            || e == Error::Km(ErrorCode::UNSUPPORTED_MAC_LENGTH)
+    );
 }
 
 /// Generate a AES-GCM key with `MIN_MAC_LENGTH`. Try to create an operation using this
diff --git a/keystore2/tests/keystore2_client_attest_key_tests.rs b/keystore2/tests/keystore2_client_attest_key_tests.rs
index b286b2f..4febd9b 100644
--- a/keystore2/tests/keystore2_client_attest_key_tests.rs
+++ b/keystore2/tests/keystore2_client_attest_key_tests.rs
@@ -43,13 +43,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
         let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sec_level, algo, att_challenge, att_app_id)
-                .unwrap();
+            key_generations::generate_attestation_key(&sec_level, algo, att_challenge).unwrap();
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -71,7 +69,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: Some(att_challenge.to_vec()),
-                att_app_id: Some(att_app_id.to_vec()),
             },
             Some(&attestation_key_metadata.key),
         )
@@ -94,13 +91,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
         let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sec_level, algo, att_challenge, att_app_id)
-                .unwrap();
+            key_generations::generate_attestation_key(&sec_level, algo, att_challenge).unwrap();
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -122,7 +117,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: Some(att_challenge.to_vec()),
-                att_app_id: Some(att_app_id.to_vec()),
             },
             Some(&attestation_key_metadata.key),
         )
@@ -146,20 +140,16 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
         let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sec_level, algo, att_challenge, att_app_id)
-                .unwrap();
+            key_generations::generate_attestation_key(&sec_level, algo, att_challenge).unwrap();
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
         cert_chain.extend(attestation_key_metadata.certificateChain.as_ref().unwrap());
-        // The server seems to be issuing test certs with invalid subject names.
-        // Re-enable when b/263254184 is fixed
-        // validate_certchain(&cert_chain).expect("Error while validating cert chain.");
+        validate_certchain(&cert_chain).expect("Error while validating cert chain.");
 
         // Create EC key and use attestation key to sign it.
         let ec_key_alias = format!("ks_ec_attested_test_key_{}", getuid());
@@ -167,7 +157,6 @@
             &sec_level,
             Some(ec_key_alias),
             att_challenge,
-            att_app_id,
             &attestation_key_metadata.key,
         )
         .unwrap();
@@ -177,9 +166,7 @@
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
         cert_chain.extend(attestation_key_metadata.certificateChain.as_ref().unwrap());
 
-        // The server seems to be issuing test certs with invalid subject names.
-        // Re-enable when b/263254184 is fixed
-        // validate_certchain(&cert_chain).expect("Error while validating cert chain.");
+        validate_certchain(&cert_chain).expect("Error while validating cert chain.");
     }
 }
 
@@ -193,13 +180,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     // Create EcCurve::CURVE_25519 attestation key.
     let attestation_key_metadata = key_generations::generate_ec_attestation_key(
         &sec_level,
         att_challenge,
-        att_app_id,
         Digest::NONE,
         EcCurve::CURVE_25519,
     )
@@ -225,7 +210,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: Some(att_challenge.to_vec()),
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&attestation_key_metadata.key),
     )
@@ -327,16 +311,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     // Create RSA attestation key.
-    let attestation_key_metadata = key_generations::generate_attestation_key(
-        &sec_level,
-        Algorithm::RSA,
-        att_challenge,
-        att_app_id,
-    )
-    .unwrap();
+    let attestation_key_metadata =
+        key_generations::generate_attestation_key(&sec_level, Algorithm::RSA, att_challenge)
+            .unwrap();
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -358,7 +337,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&attestation_key_metadata.key),
     ));
@@ -376,7 +354,6 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     let alias = format!("non_attest_key_{}", getuid());
     let non_attest_key_metadata = key_generations::generate_ec_p256_signing_key(
@@ -385,7 +362,6 @@
         -1,
         Some(alias),
         None,
-        None,
     )
     .unwrap();
 
@@ -404,7 +380,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: Some(att_challenge.to_vec()),
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&non_attest_key_metadata.key),
     ));
@@ -421,7 +396,6 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     let alias = "aes_attest_key";
     let sym_key_metadata = key_generations::generate_sym_key(
@@ -450,7 +424,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: Some(att_challenge.to_vec()),
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&sym_key_metadata.key),
     ));
@@ -468,16 +441,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     // Create attestation key.
-    let attestation_key_metadata = key_generations::generate_attestation_key(
-        &sec_level,
-        Algorithm::RSA,
-        att_challenge,
-        att_app_id,
-    )
-    .unwrap();
+    let attestation_key_metadata =
+        key_generations::generate_attestation_key(&sec_level, Algorithm::RSA, att_challenge)
+            .unwrap();
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -493,8 +461,7 @@
         .key_size(128)
         .padding_mode(PaddingMode::NONE)
         .block_mode(BlockMode::ECB)
-        .attestation_challenge(att_challenge.to_vec())
-        .attestation_app_id(att_app_id.to_vec());
+        .attestation_challenge(att_challenge.to_vec());
 
     let alias = format!("ks_test_sym_key_attest_{}", getuid());
     let aes_key_metadata = sec_level
diff --git a/keystore2/tests/keystore2_client_ec_key_tests.rs b/keystore2/tests/keystore2_client_ec_key_tests.rs
index 726d61c..c2034de 100644
--- a/keystore2/tests/keystore2_client_ec_key_tests.rs
+++ b/keystore2/tests/keystore2_client_ec_key_tests.rs
@@ -209,7 +209,6 @@
         key_generations::SELINUX_SHELL_NAMESPACE,
         None,
         None,
-        None,
     )
     .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_grant_key_tests.rs b/keystore2/tests/keystore2_client_grant_key_tests.rs
index 827a0de..7c75734 100644
--- a/keystore2/tests/keystore2_client_grant_key_tests.rs
+++ b/keystore2/tests/keystore2_client_grant_key_tests.rs
@@ -44,7 +44,6 @@
         key_generations::SELINUX_SHELL_NAMESPACE,
         Some(alias),
         None,
-        None,
     )
     .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_import_keys_tests.rs b/keystore2/tests/keystore2_client_import_keys_tests.rs
index c8f94b6..ecba402 100644
--- a/keystore2/tests/keystore2_client_import_keys_tests.rs
+++ b/keystore2/tests/keystore2_client_import_keys_tests.rs
@@ -35,7 +35,7 @@
 use crate::ffi_test_utils::{create_wrapped_key, create_wrapped_key_additional_auth_data};
 
 use crate::keystore2_client_test_utils::{
-    encrypt_secure_key, encrypt_transport_key, has_trusty_keymint,
+    encrypt_secure_key, encrypt_transport_key, has_default_keymint,
     perform_sample_asym_sign_verify_op, perform_sample_hmac_sign_verify_op,
     perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op, SAMPLE_PLAIN_TEXT,
 };
@@ -286,7 +286,7 @@
         key_generations::RSA_2048_KEY,
     ));
 
-    if has_trusty_keymint() {
+    if has_default_keymint() {
         assert!(result.is_err());
         assert_eq!(Error::Km(ErrorCode::INCOMPATIBLE_PURPOSE), result.unwrap_err());
     } else {
diff --git a/keystore2/tests/keystore2_client_list_entries_tests.rs b/keystore2/tests/keystore2_client_list_entries_tests.rs
index def9d94..62e3dd0 100644
--- a/keystore2/tests/keystore2_client_list_entries_tests.rs
+++ b/keystore2/tests/keystore2_client_list_entries_tests.rs
@@ -89,7 +89,6 @@
                 key_generations::SELINUX_SHELL_NAMESPACE,
                 Some(alias.to_string()),
                 None,
-                None,
             )
             .unwrap();
 
@@ -128,7 +127,6 @@
                     -1,
                     Some(alias.to_string()),
                     None,
-                    None,
                 )
                 .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_operation_tests.rs b/keystore2/tests/keystore2_client_operation_tests.rs
index e1102dd..9714900 100644
--- a/keystore2/tests/keystore2_client_operation_tests.rs
+++ b/keystore2/tests/keystore2_client_operation_tests.rs
@@ -307,7 +307,6 @@
         key_generations::SELINUX_SHELL_NAMESPACE,
         Some(alias),
         None,
-        None,
     )
     .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_rsa_key_tests.rs b/keystore2/tests/keystore2_client_rsa_key_tests.rs
index 3139c2b..ad176a4 100644
--- a/keystore2/tests/keystore2_client_rsa_key_tests.rs
+++ b/keystore2/tests/keystore2_client_rsa_key_tests.rs
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    BlockMode::BlockMode, Digest::Digest, ErrorCode::ErrorCode, KeyPurpose::KeyPurpose,
-    PaddingMode::PaddingMode, SecurityLevel::SecurityLevel,
+    Digest::Digest, ErrorCode::ErrorCode, KeyPurpose::KeyPurpose, PaddingMode::PaddingMode,
+    SecurityLevel::SecurityLevel,
 };
 use android_system_keystore2::aidl::android::system::keystore2::{
     CreateOperationResponse::CreateOperationResponse, Domain::Domain,
@@ -25,9 +25,7 @@
     authorizations, get_keystore_service, key_generations, key_generations::Error,
 };
 
-use crate::keystore2_client_test_utils::{
-    delete_app_key, has_trusty_keymint, perform_sample_sign_operation, ForcedOp,
-};
+use crate::keystore2_client_test_utils::{delete_app_key, perform_sample_sign_operation, ForcedOp};
 
 /// This macro is used for creating signing key operation tests using digests and paddings
 /// for various key sizes.
@@ -59,7 +57,6 @@
                 stringify!($test_name),
                 $padding,
                 None,
-                None,
             );
         }
     };
@@ -73,7 +70,6 @@
                 stringify!($test_name),
                 $padding,
                 $mgf_digest,
-                Some(BlockMode::ECB),
             );
         }
     };
@@ -133,7 +129,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
@@ -170,18 +165,16 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
     ));
     assert!(result.is_err());
 
-    if has_trusty_keymint() {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::UNKNOWN_ERROR));
-    } else {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::INCOMPATIBLE_DIGEST));
-    }
+    let e = result.unwrap_err();
+    assert!(
+        e == Error::Km(ErrorCode::UNKNOWN_ERROR) || e == Error::Km(ErrorCode::INCOMPATIBLE_DIGEST)
+    );
 
     delete_app_key(&keystore2, alias).unwrap();
 }
@@ -193,7 +186,6 @@
     alias: &str,
     padding: PaddingMode,
     mgf_digest: Option<Digest>,
-    block_mode: Option<BlockMode>,
 ) {
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
@@ -209,9 +201,8 @@
             padding: Some(padding),
             digest,
             mgf_digest,
-            block_mode,
+            block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1559,7 +1550,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
@@ -1592,7 +1582,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1624,7 +1613,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1654,7 +1642,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
@@ -1684,7 +1671,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1714,7 +1700,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::AGREE_KEY,
         ForcedOp(false),
@@ -1747,7 +1732,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: None,
-                att_app_id: None,
             },
             KeyPurpose::DECRYPT,
             ForcedOp(false),
@@ -1781,7 +1765,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: None,
-                att_app_id: None,
             },
             KeyPurpose::SIGN,
             ForcedOp(false),
@@ -1813,7 +1796,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::ENCRYPT,
         ForcedOp(false),
@@ -1845,7 +1827,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1876,7 +1857,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1906,7 +1886,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         None,
     ));
diff --git a/keystore2/tests/keystore2_client_test_utils.rs b/keystore2/tests/keystore2_client_test_utils.rs
index 59819df..56995e4 100644
--- a/keystore2/tests/keystore2_client_test_utils.rs
+++ b/keystore2/tests/keystore2_client_test_utils.rs
@@ -84,7 +84,8 @@
     };
 }
 
-pub fn has_trusty_keymint() -> bool {
+/// Indicate whether the default device is KeyMint (rather than Keymaster).
+pub fn has_default_keymint() -> bool {
     binder::is_declared("android.hardware.security.keymint.IKeyMintDevice/default")
         .expect("Could not check for declared keymint interface")
 }
@@ -102,10 +103,9 @@
     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, nspace, alias, None, None,
-    )
-    .unwrap();
+    let key_metadata =
+        key_generations::generate_ec_p256_signing_key(&sec_level, domain, nspace, alias, None)
+            .unwrap();
 
     sec_level.createOperation(
         &key_metadata.key,
diff --git a/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs b/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs
index 32ecd03..63122fe 100644
--- a/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs
+++ b/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs
@@ -165,14 +165,12 @@
                 .unwrap();
             // Generate Key BLOB and prepare legacy keystore blob files.
             let att_challenge: &[u8] = b"foo";
-            let att_app_id: &[u8] = b"bar";
             let key_metadata = key_generations::generate_ec_p256_signing_key(
                 &sec_level,
                 Domain::BLOB,
                 SELINUX_SHELL_NAMESPACE,
                 None,
                 Some(att_challenge),
-                Some(att_app_id),
             )
             .expect("Failed to generate key blob");
 
@@ -424,14 +422,12 @@
                 .unwrap();
             // Generate Key BLOB and prepare legacy keystore blob files.
             let att_challenge: &[u8] = b"foo";
-            let att_app_id: &[u8] = b"bar";
             let key_metadata = key_generations::generate_ec_p256_signing_key(
                 &sec_level,
                 Domain::BLOB,
                 SELINUX_SHELL_NAMESPACE,
                 None,
                 Some(att_challenge),
-                Some(att_app_id),
             )
             .expect("Failed to generate key blob");