Merge "Change requests for Rust update v1.66.0"
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/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..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/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/rkpd_client.rs b/keystore2/src/rkpd_client.rs
index d611678..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, 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,7 +283,7 @@
     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)]
@@ -254,19 +291,25 @@
     use super::*;
     use android_security_rkp_aidl::aidl::android::security::rkp::IRegistration::BnRegistration;
     use std::sync::atomic::{AtomicU32, Ordering};
-    use std::sync::Arc;
 
     #[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())
         }
     }
@@ -274,8 +317,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<()> {
@@ -286,19 +343,33 @@
             &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
@@ -310,7 +381,8 @@
 
     #[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());
     }
 
@@ -320,10 +392,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)
         );
     }
 
@@ -333,10 +405,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)
         );
     }
 
@@ -348,7 +420,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);
     }
 
@@ -358,10 +430,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)
         );
     }
 
@@ -371,10 +443,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)
         );
     }
 
@@ -384,7 +456,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]
@@ -393,17 +465,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());
     }
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index 4481172..7bc548e 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -277,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/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();
+    }
+}