diff --git a/OWNERS b/OWNERS
index 6fdb550..b1c1fea 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,11 +1,7 @@
-alanstokes@google.com
 drysdale@google.com
-eranm@google.com
 hasinitg@google.com
 jbires@google.com
 jeffv@google.com
-kroot@google.com
 sethmo@google.com
 swillden@google.com
-trong@google.com
 zeuthen@google.com
diff --git a/keystore2/OWNERS b/keystore2/OWNERS
index aeb8390..1fcc785 100644
--- a/keystore2/OWNERS
+++ b/keystore2/OWNERS
@@ -1,7 +1,9 @@
 set noparent
 # Bug component: 1084732
+cvlasov@google.com
 drysdale@google.com
 hasinitg@google.com
 jbires@google.com
+kwadhera@google.com
 sethmo@google.com
 swillden@google.com
diff --git a/keystore2/aidl/Android.bp b/keystore2/aidl/Android.bp
index afc2743..13bf455 100644
--- a/keystore2/aidl/Android.bp
+++ b/keystore2/aidl/Android.bp
@@ -24,8 +24,8 @@
 aidl_interface {
     name: "android.security.authorization",
     srcs: ["android/security/authorization/*.aidl"],
+    defaults: ["android.hardware.security.keymint-latest-defaults"],
     imports: [
-        "android.hardware.security.keymint-V3",
         "android.hardware.security.secureclock-V1",
     ],
     unstable: true,
@@ -63,8 +63,8 @@
 aidl_interface {
     name: "android.security.compat",
     srcs: ["android/security/compat/*.aidl"],
+    defaults: ["android.hardware.security.keymint-latest-defaults"],
     imports: [
-        "android.hardware.security.keymint-V3",
         "android.hardware.security.secureclock-V1",
         "android.hardware.security.sharedsecret-V1",
     ],
@@ -86,8 +86,8 @@
 aidl_interface {
     name: "android.security.maintenance",
     srcs: ["android/security/maintenance/*.aidl"],
-    imports: [
-        "android.system.keystore2-V4",
+    defaults: [
+        "android.system.keystore2-latest-defaults",
     ],
     unstable: true,
     backend: {
@@ -145,8 +145,8 @@
 aidl_interface {
     name: "android.security.metrics",
     srcs: ["android/security/metrics/*.aidl"],
-    imports: [
-        "android.system.keystore2-V4",
+    defaults: [
+        "android.system.keystore2-latest-defaults",
     ],
     unstable: true,
     backend: {
@@ -169,21 +169,21 @@
 java_defaults {
     name: "keystore2_use_latest_aidl_java_static",
     static_libs: [
-        "android.system.keystore2-V4-java-source",
+        "android.system.keystore2-V5-java-source",
     ],
 }
 
 java_defaults {
     name: "keystore2_use_latest_aidl_java_shared",
     libs: [
-        "android.system.keystore2-V4-java-source",
+        "android.system.keystore2-V5-java-source",
     ],
 }
 
 java_defaults {
     name: "keystore2_use_latest_aidl_java",
     libs: [
-        "android.system.keystore2-V4-java",
+        "android.system.keystore2-V5-java",
     ],
 }
 
@@ -193,28 +193,28 @@
 cc_defaults {
     name: "keystore2_use_latest_aidl_ndk_static",
     static_libs: [
-        "android.system.keystore2-V4-ndk",
+        "android.system.keystore2-V5-ndk",
     ],
 }
 
 cc_defaults {
     name: "keystore2_use_latest_aidl_ndk_shared",
     shared_libs: [
-        "android.system.keystore2-V4-ndk",
+        "android.system.keystore2-V5-ndk",
     ],
 }
 
 cc_defaults {
     name: "keystore2_use_latest_aidl_cpp_shared",
     shared_libs: [
-        "android.system.keystore2-V4-cpp",
+        "android.system.keystore2-V5-cpp",
     ],
 }
 
 cc_defaults {
     name: "keystore2_use_latest_aidl_cpp_static",
     static_libs: [
-        "android.system.keystore2-V4-cpp",
+        "android.system.keystore2-V5-cpp",
     ],
 }
 
@@ -224,6 +224,6 @@
 rust_defaults {
     name: "keystore2_use_latest_aidl_rust",
     rustlibs: [
-        "android.system.keystore2-V4-rust",
+        "android.system.keystore2-V5-rust",
     ],
 }
diff --git a/keystore2/android.system.keystore2-service.xml b/keystore2/android.system.keystore2-service.xml
index 4d8a756..35b9cc8 100644
--- a/keystore2/android.system.keystore2-service.xml
+++ b/keystore2/android.system.keystore2-service.xml
@@ -1,7 +1,7 @@
 <manifest version="1.0" type="framework">
     <hal format="aidl">
         <name>android.system.keystore2</name>
-        <version>4</version>
+        <version>5</version>
         <interface>
             <name>IKeystoreService</name>
             <instance>default</instance>
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 8451a33..c815622 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -125,6 +125,14 @@
     }
 }
 
+/// Access information for a key.
+#[derive(Debug)]
+struct KeyAccessInfo {
+    key_id: i64,
+    descriptor: KeyDescriptor,
+    vector: Option<KeyPermSet>,
+}
+
 /// If the database returns a busy error code, retry after this interval.
 const DB_BUSY_RETRY_INTERVAL: Duration = Duration::from_micros(500);
 
@@ -2071,7 +2079,7 @@
         key: &KeyDescriptor,
         key_type: KeyType,
         caller_uid: u32,
-    ) -> Result<(i64, KeyDescriptor, Option<KeyPermSet>)> {
+    ) -> Result<KeyAccessInfo> {
         match key.domain {
             // Domain App or SELinux. In this case we load the key_id from
             // the keyentry database for further loading of key components.
@@ -2087,7 +2095,7 @@
                 let key_id = Self::load_key_entry_id(tx, &access_key, key_type)
                     .with_context(|| format!("With key.domain = {:?}.", access_key.domain))?;
 
-                Ok((key_id, access_key, None))
+                Ok(KeyAccessInfo { key_id, descriptor: access_key, vector: None })
             }
 
             // Domain::GRANT. In this case we load the key_id and the access_vector
@@ -2113,7 +2121,11 @@
                         ))
                     })
                     .context("Domain::GRANT.")?;
-                Ok((key_id, key.clone(), Some(access_vector.into())))
+                Ok(KeyAccessInfo {
+                    key_id,
+                    descriptor: key.clone(),
+                    vector: Some(access_vector.into()),
+                })
             }
 
             // Domain::KEY_ID. In this case we load the domain and namespace from the
@@ -2169,7 +2181,7 @@
                 access_key.domain = domain;
                 access_key.nspace = namespace;
 
-                Ok((key_id, access_key, access_vector))
+                Ok(KeyAccessInfo { key_id, descriptor: access_key, vector: access_vector })
             }
             _ => Err(anyhow!(KsError::Rc(ResponseCode::INVALID_ARGUMENT))),
         }
@@ -2356,12 +2368,11 @@
             .context(ks_err!("Failed to initialize transaction."))?;
 
         // Load the key_id and complete the access control tuple.
-        let (key_id, access_key_descriptor, access_vector) =
-            Self::load_access_tuple(&tx, key, key_type, caller_uid).context(ks_err!())?;
+        let access = Self::load_access_tuple(&tx, key, key_type, caller_uid).context(ks_err!())?;
 
         // Perform access control. It is vital that we return here if the permission is denied.
         // So do not touch that '?' at the end.
-        check_permission(&access_key_descriptor, access_vector).context(ks_err!())?;
+        check_permission(&access.descriptor, access.vector).context(ks_err!())?;
 
         // KEY ID LOCK 2/2
         // If we did not get a key id lock by now, it was because we got a key descriptor
@@ -2375,13 +2386,13 @@
         // that the caller had access to the given key. But we need to make sure that the
         // key id still exists. So we have to load the key entry by key id this time.
         let (key_id_guard, tx) = match key_id_guard {
-            None => match KEY_ID_LOCK.try_get(key_id) {
+            None => match KEY_ID_LOCK.try_get(access.key_id) {
                 None => {
                     // Roll back the transaction.
                     tx.rollback().context(ks_err!("Failed to roll back transaction."))?;
 
                     // Block until we have a key id lock.
-                    let key_id_guard = KEY_ID_LOCK.get(key_id);
+                    let key_id_guard = KEY_ID_LOCK.get(access.key_id);
 
                     // Create a new transaction.
                     let tx = self
@@ -2395,7 +2406,7 @@
                         // alias may have been rebound after we rolled back the transaction.
                         &KeyDescriptor {
                             domain: Domain::KEY_ID,
-                            nspace: key_id,
+                            nspace: access.key_id,
                             ..Default::default()
                         },
                         key_type,
@@ -2451,16 +2462,15 @@
         let _wp = wd::watch("KeystoreDB::unbind_key");
 
         self.with_transaction(Immediate("TX_unbind_key"), |tx| {
-            let (key_id, access_key_descriptor, access_vector) =
-                Self::load_access_tuple(tx, key, key_type, caller_uid)
-                    .context("Trying to get access tuple.")?;
+            let access = Self::load_access_tuple(tx, key, key_type, caller_uid)
+                .context("Trying to get access tuple.")?;
 
             // Perform access control. It is vital that we return here if the permission is denied.
             // So do not touch that '?' at the end.
-            check_permission(&access_key_descriptor, access_vector)
+            check_permission(&access.descriptor, access.vector)
                 .context("While checking permission.")?;
 
-            Self::mark_unreferenced(tx, key_id)
+            Self::mark_unreferenced(tx, access.key_id)
                 .map(|need_gc| (need_gc, ()))
                 .context("Trying to mark the key unreferenced.")
         })
@@ -2830,7 +2840,7 @@
             // We could check key.domain == Domain::GRANT and fail early.
             // But even if we load the access tuple by grant here, the permission
             // check denies the attempt to create a grant by grant descriptor.
-            let (key_id, access_key_descriptor, _) =
+            let access =
                 Self::load_access_tuple(tx, key, KeyType::Client, caller_uid).context(ks_err!())?;
 
             // Perform access control. It is vital that we return here if the permission
@@ -2838,14 +2848,14 @@
             // This permission check checks if the caller has the grant permission
             // for the given key and in addition to all of the permissions
             // expressed in `access_vector`.
-            check_permission(&access_key_descriptor, &access_vector)
+            check_permission(&access.descriptor, &access_vector)
                 .context(ks_err!("check_permission failed"))?;
 
             let grant_id = if let Some(grant_id) = tx
                 .query_row(
                     "SELECT id FROM persistent.grant
                 WHERE keyentryid = ? AND grantee = ?;",
-                    params![key_id, grantee_uid],
+                    params![access.key_id, grantee_uid],
                     |row| row.get(0),
                 )
                 .optional()
@@ -2864,7 +2874,7 @@
                     tx.execute(
                         "INSERT INTO persistent.grant (id, grantee, keyentryid, access_vector)
                         VALUES (?, ?, ?, ?);",
-                        params![id, grantee_uid, key_id, i32::from(access_vector)],
+                        params![id, grantee_uid, access.key_id, i32::from(access_vector)],
                     )
                 })
                 .context(ks_err!())?
@@ -2889,18 +2899,17 @@
         self.with_transaction(Immediate("TX_ungrant"), |tx| {
             // Load the key_id and complete the access control tuple.
             // We ignore the access vector here because grants cannot be granted.
-            let (key_id, access_key_descriptor, _) =
+            let access =
                 Self::load_access_tuple(tx, key, KeyType::Client, caller_uid).context(ks_err!())?;
 
             // Perform access control. We must return here if the permission
             // was denied. So do not touch the '?' at the end of this line.
-            check_permission(&access_key_descriptor)
-                .context(ks_err!("check_permission failed."))?;
+            check_permission(&access.descriptor).context(ks_err!("check_permission failed."))?;
 
             tx.execute(
                 "DELETE FROM persistent.grant
                 WHERE keyentryid = ? AND grantee = ?;",
-                params![key_id, grantee_uid],
+                params![access.key_id, grantee_uid],
             )
             .context("Failed to delete grant.")?;
 
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index 39d6f9c..0e8892b 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -271,17 +271,8 @@
     // If the KeyMint device is back-level, use a wrapper that intercepts and
     // emulates things that are not supported by the hardware.
     let keymint = match hal_version {
-        Some(300) => {
-            // Current KeyMint version: use as-is as v3 Keymint is current version
-            log::info!(
-                "KeyMint device is current version ({:?}) for security level: {:?}",
-                hal_version,
-                security_level
-            );
-            keymint
-        }
-        Some(200) => {
-            // Previous KeyMint version: use as-is as we don't have any software emulation of v3-specific KeyMint features.
+        Some(400) | Some(300) | Some(200) => {
+            // KeyMint v2+: use as-is (we don't have any software emulation of v3 or v4-specific KeyMint features).
             log::info!(
                 "KeyMint device is current version ({:?}) for security level: {:?}",
                 hal_version,
diff --git a/keystore2/src/km_compat.rs b/keystore2/src/km_compat.rs
index 5e3bdfa..95e9294 100644
--- a/keystore2/src/km_compat.rs
+++ b/keystore2/src/km_compat.rs
@@ -214,6 +214,12 @@
     fn sendRootOfTrust(&self, root_of_trust: &[u8]) -> binder::Result<()> {
         self.real.sendRootOfTrust(root_of_trust)
     }
+    fn setAdditionalAttestationInfo(
+        &self,
+        additional_attestation_info: &[KeyParameter],
+    ) -> binder::Result<()> {
+        self.real.setAdditionalAttestationInfo(additional_attestation_info)
+    }
 
     // For methods that emit keyblobs, check whether the underlying real device
     // supports the relevant parameters, and forward to the appropriate device.
diff --git a/keystore2/src/km_compat/km_compat.cpp b/keystore2/src/km_compat/km_compat.cpp
index e9ff1ff..7a6ef4a 100644
--- a/keystore2/src/km_compat/km_compat.cpp
+++ b/keystore2/src/km_compat/km_compat.cpp
@@ -839,6 +839,11 @@
     return convertErrorCode(KMV1::ErrorCode::UNIMPLEMENTED);
 }
 
+ScopedAStatus KeyMintDevice::setAdditionalAttestationInfo(
+    const std::vector<KeyParameter>& /* additionalAttestationInfo */) {
+    return convertErrorCode(KMV1::ErrorCode::UNIMPLEMENTED);
+}
+
 ScopedAStatus KeyMintOperation::updateAad(const std::vector<uint8_t>& input,
                                           const std::optional<HardwareAuthToken>& optAuthToken,
                                           const std::optional<TimeStampToken>& optTimeStampToken) {
diff --git a/keystore2/src/km_compat/km_compat.h b/keystore2/src/km_compat/km_compat.h
index c4bcdaa..71f7fbe 100644
--- a/keystore2/src/km_compat/km_compat.h
+++ b/keystore2/src/km_compat/km_compat.h
@@ -147,6 +147,9 @@
                                  std::vector<uint8_t>* rootOfTrust);
     ScopedAStatus sendRootOfTrust(const std::vector<uint8_t>& rootOfTrust);
 
+    ScopedAStatus
+    setAdditionalAttestationInfo(const std::vector<KeyParameter>& additionalAttestationInfo);
+
     // These are public to allow testing code to use them directly.
     // This class should not be used publicly anyway.
     std::variant<std::vector<Certificate>, KMV1_ErrorCode>
diff --git a/keystore2/src/km_compat/km_compat_type_conversion.h b/keystore2/src/km_compat/km_compat_type_conversion.h
index 5db7e3d..d6a2dcc 100644
--- a/keystore2/src/km_compat/km_compat_type_conversion.h
+++ b/keystore2/src/km_compat/km_compat_type_conversion.h
@@ -750,8 +750,12 @@
     case KMV1::Tag::CERTIFICATE_SUBJECT:
     case KMV1::Tag::CERTIFICATE_NOT_BEFORE:
     case KMV1::Tag::CERTIFICATE_NOT_AFTER:
+        // These tags do not exist in KM < KeyMint 1.
+        break;
     case KMV1::Tag::ATTESTATION_ID_SECOND_IMEI:
-        // These tags do not exist in KM < KeyMint 1.0.
+        // This tag doesn't exist in KM < KeyMint 3.
+    case KMV1::Tag::MODULE_HASH:
+        // This tag doesn't exist in KM < KeyMint 4.
         break;
     case KMV1::Tag::MAX_BOOT_LEVEL:
         // Does not exist in API level 30 or below.
diff --git a/keystore2/src/permission.rs b/keystore2/src/permission.rs
index 7bf17b5..023774f 100644
--- a/keystore2/src/permission.rs
+++ b/keystore2/src/permission.rs
@@ -282,12 +282,19 @@
 ///                      SELinux keystore key backend, and the result is used
 ///                      as target context.
 pub fn check_grant_permission(
+    caller_uid: u32,
     caller_ctx: &CStr,
     access_vec: KeyPermSet,
     key: &KeyDescriptor,
 ) -> anyhow::Result<()> {
     let target_context = match key.domain {
-        Domain::APP => getcon().context("check_grant_permission: getcon failed.")?,
+        Domain::APP => {
+            if caller_uid as i64 != key.nspace {
+                return Err(selinux::Error::perm())
+                    .context("Trying to access key without ownership.");
+            }
+            getcon().context("check_grant_permission: getcon failed.")?
+        }
         Domain::SELINUX => lookup_keystore2_key_context(key.nspace)
             .context("check_grant_permission: Domain::SELINUX: Failed to lookup namespace.")?,
         _ => return Err(KsError::sys()).context(format!("Cannot grant {:?}.", key.domain)),
diff --git a/keystore2/src/permission/tests.rs b/keystore2/src/permission/tests.rs
index c9aebfe..68c9b74 100644
--- a/keystore2/src/permission/tests.rs
+++ b/keystore2/src/permission/tests.rs
@@ -135,11 +135,11 @@
 fn check_grant_permission_app() -> Result<()> {
     let system_server_ctx = Context::new("u:r:system_server:s0")?;
     let key = KeyDescriptor { domain: Domain::APP, nspace: 0, alias: None, blob: None };
-    check_grant_permission(&system_server_ctx, SYSTEM_SERVER_PERMISSIONS_NO_GRANT, &key)
+    check_grant_permission(0, &system_server_ctx, SYSTEM_SERVER_PERMISSIONS_NO_GRANT, &key)
         .expect("Grant permission check failed.");
 
     // attempts to grant the grant permission must always fail even when privileged.
-    assert_perm_failed!(check_grant_permission(&system_server_ctx, KeyPerm::Grant.into(), &key));
+    assert_perm_failed!(check_grant_permission(0, &system_server_ctx, KeyPerm::Grant.into(), &key));
     Ok(())
 }
 
@@ -153,12 +153,12 @@
         blob: None,
     };
     if is_su {
-        assert!(check_grant_permission(&sctx, NOT_GRANT_PERMS, &key).is_ok());
+        assert!(check_grant_permission(0, &sctx, NOT_GRANT_PERMS, &key).is_ok());
         // attempts to grant the grant permission must always fail even when privileged.
-        assert_perm_failed!(check_grant_permission(&sctx, KeyPerm::Grant.into(), &key));
+        assert_perm_failed!(check_grant_permission(0, &sctx, KeyPerm::Grant.into(), &key));
     } else {
         // unprivileged grant attempts always fail. shell does not have the grant permission.
-        assert_perm_failed!(check_grant_permission(&sctx, UNPRIV_PERMS, &key));
+        assert_perm_failed!(check_grant_permission(0, &sctx, UNPRIV_PERMS, &key));
     }
     Ok(())
 }
diff --git a/keystore2/src/utils.rs b/keystore2/src/utils.rs
index 2b69d1e..c6dc11e 100644
--- a/keystore2/src/utils.rs
+++ b/keystore2/src/utils.rs
@@ -80,6 +80,7 @@
 pub fn check_grant_permission(access_vec: KeyPermSet, key: &KeyDescriptor) -> anyhow::Result<()> {
     ThreadState::with_calling_sid(|calling_sid| {
         permission::check_grant_permission(
+            ThreadState::get_calling_uid(),
             calling_sid
                 .ok_or_else(Error::sys)
                 .context(ks_err!("Cannot check permission without calling_sid."))?,
diff --git a/keystore2/test_utils/Android.bp b/keystore2/test_utils/Android.bp
index d0b5540..57da27f 100644
--- a/keystore2/test_utils/Android.bp
+++ b/keystore2/test_utils/Android.bp
@@ -62,8 +62,8 @@
     static_libs: [
         // Also include static_libs for the NDK variants so that they are available
         // for dependencies.
-        "android.system.keystore2-V4-ndk",
-        "android.hardware.security.keymint-V3-ndk",
+        "android.system.keystore2-V5-ndk",
+        "android.hardware.security.keymint-V4-ndk",
     ],
 }
 
diff --git a/keystore2/test_utils/run_as.rs b/keystore2/test_utils/run_as.rs
index d4fd06c..7a9acb7 100644
--- a/keystore2/test_utils/run_as.rs
+++ b/keystore2/test_utils/run_as.rs
@@ -352,6 +352,40 @@
     }
 }
 
+/// Run the given closure in a new process running as an untrusted app with the given `uid` and
+/// `gid`. Parent process will run without waiting for child status.
+///
+/// # Safety
+/// run_as_child runs the given closure in the client branch of fork. And it uses non
+/// async signal safe API. This means that calling this function in a multi threaded program
+/// yields undefined behavior in the child. As of this writing, it is safe to call this function
+/// from a Rust device test, because every test itself is spawned as a separate process.
+///
+/// # Safety Binder
+/// It is okay for the closure to use binder services, however, this does not work
+/// if the parent initialized libbinder already. So do not use binder outside of the closure
+/// in your test.
+pub unsafe fn run_as_child_app<F, R, M>(
+    uid: u32,
+    gid: u32,
+    f: F,
+) -> Result<ChildHandle<R, M>, nix::Error>
+where
+    R: Serialize + DeserializeOwned,
+    M: Serialize + DeserializeOwned,
+    F: 'static + Send + FnOnce(&mut ChannelReader<M>, &mut ChannelWriter<M>) -> R,
+{
+    // Safety: Caller guarantees that the process only has a single thread.
+    unsafe {
+        run_as_child(
+            "u:r:untrusted_app:s0:c91,c256,c10,c20",
+            Uid::from_raw(uid),
+            Gid::from_raw(gid),
+            f,
+        )
+    }
+}
+
 /// Run the given closure in a new process running with the new identity given as
 /// `uid`, `gid`, and `se_context`. Parent process will run without waiting for child status.
 ///
@@ -383,7 +417,7 @@
     let (response_reader, mut response_writer) =
         pipe_channel().expect("Failed to create cmd pipe.");
 
-    // SAFETY: Our caller guarantees that the process only has a single thread, so calling
+    // Safety: Our caller guarantees that the process only has a single thread, so calling
     // non-async-signal-safe functions in the child is in fact safe.
     match unsafe { fork() } {
         Ok(ForkResult::Parent { child, .. }) => {
diff --git a/keystore2/tests/Android.bp b/keystore2/tests/Android.bp
index 0406a71..f18824b 100644
--- a/keystore2/tests/Android.bp
+++ b/keystore2/tests/Android.bp
@@ -31,8 +31,8 @@
     static_libs: [
         // Also include static_libs for the NDK variants so that they are available
         // for dependencies.
-        "android.system.keystore2-V4-ndk",
-        "android.hardware.security.keymint-V3-ndk",
+        "android.system.keystore2-V5-ndk",
+        "android.hardware.security.keymint-V4-ndk",
     ],
     srcs: ["keystore2_client_tests.rs"],
     test_suites: [
diff --git a/keystore2/tests/keystore2_client_grant_key_tests.rs b/keystore2/tests/keystore2_client_grant_key_tests.rs
index 5391d20..c171ab1 100644
--- a/keystore2/tests/keystore2_client_grant_key_tests.rs
+++ b/keystore2/tests/keystore2_client_grant_key_tests.rs
@@ -19,20 +19,35 @@
     Digest::Digest, KeyPurpose::KeyPurpose,
 };
 use android_system_keystore2::aidl::android::system::keystore2::{
-    Domain::Domain, KeyDescriptor::KeyDescriptor, KeyPermission::KeyPermission,
-    ResponseCode::ResponseCode,
+    Domain::Domain, IKeystoreService::IKeystoreService, KeyDescriptor::KeyDescriptor,
+    KeyEntryResponse::KeyEntryResponse, KeyPermission::KeyPermission, ResponseCode::ResponseCode,
 };
 use keystore2_test_utils::{
-    authorizations, get_keystore_service, key_generations, key_generations::Error, run_as, SecLevel,
+    authorizations, get_keystore_service, key_generations,
+    key_generations::{map_ks_error, Error},
+    run_as, SecLevel,
 };
 use nix::unistd::getuid;
 use rustutils::users::AID_USER_OFFSET;
 
-/// Generate an EC signing key and grant it to the user with given access vector.
-fn generate_ec_key_and_grant_to_user(
-    grantee_uid: i32,
+/// Produce a [`KeyDescriptor`] for a granted key.
+fn granted_key_descriptor(nspace: i64) -> KeyDescriptor {
+    KeyDescriptor { domain: Domain::GRANT, nspace, alias: None, blob: None }
+}
+
+fn get_granted_key(
+    ks2: &binder::Strong<dyn IKeystoreService>,
+    nspace: i64,
+) -> Result<KeyEntryResponse, Error> {
+    map_ks_error(ks2.getKeyEntry(&granted_key_descriptor(nspace)))
+}
+
+/// Generate an EC signing key in the SELINUX domain and grant it to the user with given access
+/// vector.
+fn generate_and_grant_selinux_key(
+    grantee_uid: u32,
     access_vector: i32,
-) -> binder::Result<KeyDescriptor> {
+) -> Result<KeyDescriptor, Error> {
     let sl = SecLevel::tee();
     let alias = format!("{}{}", "ks_grant_test_key_1", getuid());
 
@@ -45,61 +60,54 @@
     )
     .unwrap();
 
-    sl.keystore2.grant(&key_metadata.key, grantee_uid, access_vector)
+    map_ks_error(sl.keystore2.grant(
+        &key_metadata.key,
+        grantee_uid.try_into().unwrap(),
+        access_vector,
+    ))
 }
 
-fn load_grant_key_and_perform_sign_operation(
-    sl: &SecLevel,
-    grant_key_nspace: i64,
-) -> Result<(), binder::Status> {
-    let key_entry_response = sl.keystore2.getKeyEntry(&KeyDescriptor {
-        domain: Domain::GRANT,
-        nspace: grant_key_nspace,
-        alias: None,
-        blob: None,
-    })?;
+/// Use a granted key to perform a signing operation.
+fn sign_with_granted_key(grant_key_nspace: i64) -> Result<(), Error> {
+    let sl = SecLevel::tee();
+    let key_entry_response = get_granted_key(&sl.keystore2, grant_key_nspace)?;
 
     // Perform sample crypto operation using granted key.
-    let op_response = sl.binder.createOperation(
+    let op_response = map_ks_error(sl.binder.createOperation(
         &key_entry_response.metadata.key,
         &authorizations::AuthSetBuilder::new().purpose(KeyPurpose::SIGN).digest(Digest::SHA_2_256),
         false,
-    )?;
+    ))?;
 
     assert!(op_response.iOperation.is_some());
     assert_eq!(
         Ok(()),
-        key_generations::map_ks_error(perform_sample_sign_operation(
-            &op_response.iOperation.unwrap()
-        ))
+        map_ks_error(perform_sample_sign_operation(&op_response.iOperation.unwrap()))
     );
 
     Ok(())
 }
 
-/// Try to grant a key with permission that does not map to any of the `KeyPermission` values.
-/// An error is expected with values that does not map to set of permissions listed in
+/// Try to grant an SELINUX key with permission that does not map to any of the `KeyPermission`
+/// values.  An error is expected with values that does not map to set of permissions listed in
 /// `KeyPermission`.
 #[test]
-fn keystore2_grant_key_with_invalid_perm_expecting_syserror() {
+fn grant_selinux_key_with_invalid_perm() {
     const USER_ID: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     let grantee_uid = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
     let invalid_access_vector = KeyPermission::CONVERT_STORAGE_KEY_TO_EPHEMERAL.0 << 19;
 
-    let result = key_generations::map_ks_error(generate_ec_key_and_grant_to_user(
-        grantee_uid.try_into().unwrap(),
-        invalid_access_vector,
-    ));
+    let result = generate_and_grant_selinux_key(grantee_uid, invalid_access_vector);
     assert!(result.is_err());
     assert_eq!(Error::Rc(ResponseCode::SYSTEM_ERROR), result.unwrap_err());
 }
 
-/// Try to grant a key with empty access vector `KeyPermission::NONE`, should be able to grant a
-/// key with empty access vector successfully. In grantee context try to use the granted key, it
-/// should fail to load the key with permission denied error.
+/// Try to grant an SELINUX key with empty access vector `KeyPermission::NONE`, should be able to
+/// grant a key with empty access vector successfully. In grantee context try to use the granted
+/// key, it should fail to load the key with permission denied error.
 #[test]
-fn keystore2_grant_key_with_perm_none() {
+fn grant_selinux_key_with_perm_none() {
     const USER_ID: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
@@ -108,11 +116,7 @@
     let grantor_fn = || {
         let empty_access_vector = KeyPermission::NONE.0;
 
-        let grant_key = key_generations::map_ks_error(generate_ec_key_and_grant_to_user(
-            GRANTEE_UID.try_into().unwrap(),
-            empty_access_vector,
-        ))
-        .unwrap();
+        let grant_key = generate_and_grant_selinux_key(GRANTEE_UID, empty_access_vector).unwrap();
 
         assert_eq!(grant_key.domain, Domain::GRANT);
 
@@ -128,12 +132,7 @@
     let grantee_fn = move || {
         let keystore2 = get_keystore_service();
 
-        let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
-            domain: Domain::GRANT,
-            nspace: grant_key_nspace,
-            alias: None,
-            blob: None,
-        }));
+        let result = get_granted_key(&keystore2, grant_key_nspace);
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
     };
@@ -143,13 +142,13 @@
     unsafe { run_as::run_as_app(GRANTEE_UID, GRANTEE_GID, grantee_fn) };
 }
 
-/// Grant a key to the user (grantee) with `GET_INFO|USE` key permissions. Verify whether grantee
-/// can succeed in loading the granted key and try to perform simple operation using this granted
-/// key. Grantee should be able to load the key and use the key to perform crypto operation
+/// Grant an SELINUX key to the user (grantee) with `GET_INFO|USE` key permissions. Verify whether
+/// grantee can succeed in loading the granted key and try to perform simple operation using this
+/// granted key. Grantee should be able to load the key and use the key to perform crypto operation
 /// successfully. Try to delete the granted key in grantee context where it is expected to fail to
 /// delete it as `DELETE` permission is not granted.
 #[test]
-fn keystore2_grant_get_info_use_key_perm() {
+fn grant_selinux_key_get_info_use_perms() {
     const USER_ID: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
@@ -158,11 +157,7 @@
     // Generate a key and grant it to a user with GET_INFO|USE key permissions.
     let grantor_fn = || {
         let access_vector = KeyPermission::GET_INFO.0 | KeyPermission::USE.0;
-        let grant_key = key_generations::map_ks_error(generate_ec_key_and_grant_to_user(
-            GRANTEE_UID.try_into().unwrap(),
-            access_vector,
-        ))
-        .unwrap();
+        let grant_key = generate_and_grant_selinux_key(GRANTEE_UID, access_vector).unwrap();
 
         assert_eq!(grant_key.domain, Domain::GRANT);
 
@@ -178,15 +173,7 @@
         let sl = SecLevel::tee();
 
         // Load the granted key.
-        let key_entry_response = sl
-            .keystore2
-            .getKeyEntry(&KeyDescriptor {
-                domain: Domain::GRANT,
-                nspace: grant_key_nspace,
-                alias: None,
-                blob: None,
-            })
-            .unwrap();
+        let key_entry_response = get_granted_key(&sl.keystore2, grant_key_nspace).unwrap();
 
         // Perform sample crypto operation using granted key.
         let op_response = sl
@@ -202,18 +189,12 @@
         assert!(op_response.iOperation.is_some());
         assert_eq!(
             Ok(()),
-            key_generations::map_ks_error(perform_sample_sign_operation(
-                &op_response.iOperation.unwrap()
-            ))
+            map_ks_error(perform_sample_sign_operation(&op_response.iOperation.unwrap()))
         );
 
         // Try to delete the key, it is expected to be fail with permission denied error.
-        let result = key_generations::map_ks_error(sl.keystore2.deleteKey(&KeyDescriptor {
-            domain: Domain::GRANT,
-            nspace: grant_key_nspace,
-            alias: None,
-            blob: None,
-        }));
+        let result =
+            map_ks_error(sl.keystore2.deleteKey(&granted_key_descriptor(grant_key_nspace)));
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
     };
@@ -223,11 +204,143 @@
     unsafe { run_as::run_as_app(GRANTEE_UID, GRANTEE_GID, grantee_fn) };
 }
 
-/// Grant a key to the user with DELETE access. In grantee context load the key and delete it.
+/// Grant an SELINUX key to the user (grantee) with just `GET_INFO` key permissions. Verify whether
+/// grantee can succeed in loading the granted key and try to perform simple operation using this
+/// granted key.
+#[test]
+fn grant_selinux_key_get_info_only() {
+    const USER_ID: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+
+    // Generate a key and grant it to a user with (just) GET_INFO key permissions.
+    let grantor_fn = || {
+        let access_vector = KeyPermission::GET_INFO.0;
+        let grant_key = generate_and_grant_selinux_key(GRANTEE_UID, access_vector).unwrap();
+
+        assert_eq!(grant_key.domain, Domain::GRANT);
+
+        grant_key.nspace
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    let grant_key_nspace = unsafe { run_as::run_as_root(grantor_fn) };
+
+    // In grantee context load the key and try to perform crypto operation.
+    let grantee_fn = move || {
+        let sl = SecLevel::tee();
+
+        // Load the granted key.
+        let key_entry_response = get_granted_key(&sl.keystore2, grant_key_nspace)
+            .expect("failed to get info for granted key");
+
+        // Attempt to perform sample crypto operation using granted key, now identified by <KEY_ID,
+        // key_id>.
+        let result = map_ks_error(
+            sl.binder.createOperation(
+                &key_entry_response.metadata.key,
+                &authorizations::AuthSetBuilder::new()
+                    .purpose(KeyPurpose::SIGN)
+                    .digest(Digest::SHA_2_256),
+                false,
+            ),
+        );
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+
+        // Try to delete the key using a <GRANT, grant_id> descriptor.
+        let result =
+            map_ks_error(sl.keystore2.deleteKey(&granted_key_descriptor(grant_key_nspace)));
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+
+        // Try to delete the key using a <KEY_ID, key_id> descriptor.
+        let result = map_ks_error(sl.keystore2.deleteKey(&key_entry_response.metadata.key));
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    unsafe { run_as::run_as_app(GRANTEE_UID, GRANTEE_GID, grantee_fn) };
+}
+
+/// Grant an APP key to the user (grantee) with just `GET_INFO` key permissions. Verify whether
+/// grantee can succeed in loading the granted key and try to perform simple operation using this
+/// granted key.
+#[test]
+fn grant_app_key_get_info_only() {
+    const USER_ID: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+    static ALIAS: &str = "ks_grant_key_info_only";
+
+    // Generate a key and grant it to a user with (just) GET_INFO key permissions.
+    let grantor_fn = || {
+        let sl = SecLevel::tee();
+        let access_vector = KeyPermission::GET_INFO.0;
+        let mut grant_keys = generate_ec_key_and_grant_to_users(
+            &sl,
+            Some(ALIAS.to_string()),
+            vec![GRANTEE_UID.try_into().unwrap()],
+            access_vector,
+        )
+        .unwrap();
+
+        grant_keys.remove(0)
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    let grant_key_nspace = unsafe { run_as::run_as_root(grantor_fn) };
+
+    // In grantee context load the key and try to perform crypto operation.
+    let grantee_fn = move || {
+        let sl = SecLevel::tee();
+
+        // Load the granted key.
+        let key_entry_response = get_granted_key(&sl.keystore2, grant_key_nspace)
+            .expect("failed to get info for granted key");
+
+        // Attempt to perform sample crypto operation using granted key, now identified by <KEY_ID,
+        // key_id>.
+        let result = map_ks_error(
+            sl.binder.createOperation(
+                &key_entry_response.metadata.key,
+                &authorizations::AuthSetBuilder::new()
+                    .purpose(KeyPurpose::SIGN)
+                    .digest(Digest::SHA_2_256),
+                false,
+            ),
+        );
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+
+        // Try to delete the key using a <GRANT, grant_id> descriptor.
+        let result =
+            map_ks_error(sl.keystore2.deleteKey(&granted_key_descriptor(grant_key_nspace)));
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+
+        // Try to delete the key using a <KEY_ID, key_id> descriptor.
+        let result = map_ks_error(sl.keystore2.deleteKey(&key_entry_response.metadata.key));
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    unsafe { run_as::run_as_app(GRANTEE_UID, GRANTEE_GID, grantee_fn) };
+}
+
+/// Grant an APP key to the user with DELETE access. In grantee context load the key and delete it.
 /// Verify that grantee should succeed in deleting the granted key and in grantor context test
 /// should fail to find the key with error response `KEY_NOT_FOUND`.
 #[test]
-fn keystore2_grant_delete_key_success() {
+fn grant_app_key_delete_success() {
     const USER_ID: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
@@ -256,14 +369,7 @@
     // Grantee context, delete the key.
     let grantee_fn = move || {
         let keystore2 = get_keystore_service();
-        keystore2
-            .deleteKey(&KeyDescriptor {
-                domain: Domain::GRANT,
-                nspace: grant_key_nspace,
-                alias: None,
-                blob: None,
-            })
-            .unwrap();
+        keystore2.deleteKey(&granted_key_descriptor(grant_key_nspace)).unwrap();
     };
 
     // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
@@ -273,7 +379,7 @@
     // Verify whether key got deleted in grantor's context.
     let grantor_fn = move || {
         let keystore2_inst = get_keystore_service();
-        let result = key_generations::map_ks_error(keystore2_inst.getKeyEntry(&KeyDescriptor {
+        let result = map_ks_error(keystore2_inst.getKeyEntry(&KeyDescriptor {
             domain: Domain::APP,
             nspace: -1,
             alias: Some(ALIAS.to_string()),
@@ -288,13 +394,17 @@
     unsafe { run_as::run_as_root(grantor_fn) };
 }
 
-/// Grant a key to the user. In grantee context load the granted key and try to grant it to second
-/// user. Test should fail with a response code `PERMISSION_DENIED` to grant a key to second user
-/// from grantee context. Test should make sure second grantee should not have a access to granted
-/// key.
+/// Grant an APP key to the user. In grantee context load the granted key and try to grant it to
+/// second user. Test should fail with a response code `PERMISSION_DENIED` to grant a key to second
+/// user from grantee context. Test should make sure second grantee should not have a access to
+/// granted key.
 #[test]
-#[ignore]
-fn keystore2_grant_key_fails_with_permission_denied() {
+fn grant_granted_app_key_fails() {
+    const GRANTOR_USER_ID: u32 = 97;
+    const GRANTOR_APPLICATION_ID: u32 = 10003;
+    static GRANTOR_UID: u32 = GRANTOR_USER_ID * AID_USER_OFFSET + GRANTOR_APPLICATION_ID;
+    static GRANTOR_GID: u32 = GRANTOR_UID;
+
     const USER_ID: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
@@ -322,23 +432,25 @@
     };
     // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
     // `--test-threads=1`), and nothing yet done with binder.
-    let grant_key_nspace = unsafe { run_as::run_as_root(grantor_fn) };
+    let grant_key_nspace = unsafe { run_as::run_as_app(GRANTOR_UID, GRANTOR_GID, grantor_fn) };
 
     // Grantee context, load the granted key and try to grant it to `SEC_GRANTEE_UID` grantee.
     let grantee_fn = move || {
         let keystore2 = get_keystore_service();
         let access_vector = KeyPermission::GET_INFO.0;
 
-        let key_entry_response = keystore2
-            .getKeyEntry(&KeyDescriptor {
-                domain: Domain::GRANT,
-                nspace: grant_key_nspace,
-                alias: None,
-                blob: None,
-            })
-            .unwrap();
+        // Try to grant when identifying the key with <GRANT, grant_nspace>.
+        let result = map_ks_error(keystore2.grant(
+            &granted_key_descriptor(grant_key_nspace),
+            SEC_GRANTEE_UID.try_into().unwrap(),
+            access_vector,
+        ));
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::SYSTEM_ERROR), result.unwrap_err());
 
-        let result = key_generations::map_ks_error(keystore2.grant(
+        // Load the key info and try to grant when identifying the key with <KEY_ID, keyid>.
+        let key_entry_response = get_granted_key(&keystore2, grant_key_nspace).unwrap();
+        let result = map_ks_error(keystore2.grant(
             &key_entry_response.metadata.key,
             SEC_GRANTEE_UID.try_into().unwrap(),
             access_vector,
@@ -354,14 +466,7 @@
     // Make sure second grantee shouldn't have access to the above granted key.
     let grantee2_fn = move || {
         let keystore2 = get_keystore_service();
-
-        let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
-            domain: Domain::GRANT,
-            nspace: grant_key_nspace,
-            alias: None,
-            blob: None,
-        }));
-
+        let result = get_granted_key(&keystore2, grant_key_nspace);
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
     };
@@ -371,10 +476,88 @@
     unsafe { run_as::run_as_app(SEC_GRANTEE_UID, SEC_GRANTEE_GID, grantee2_fn) };
 }
 
-/// Try to grant a key with `GRANT` access. Keystore2 system shouldn't allow to grant a key with
-/// `GRANT` access. Test should fail to grant a key with `PERMISSION_DENIED` error response code.
+/// Grant an APP key to one user, from a normal user. Check that grantee context can load the
+/// granted key, but that a second unrelated context cannot.
 #[test]
-fn keystore2_grant_key_fails_with_grant_perm_expect_perm_denied() {
+fn grant_app_key_only_to_grantee() {
+    const GRANTOR_USER_ID: u32 = 97;
+    const GRANTOR_APPLICATION_ID: u32 = 10003;
+    static GRANTOR_UID: u32 = GRANTOR_USER_ID * AID_USER_OFFSET + GRANTOR_APPLICATION_ID;
+    static GRANTOR_GID: u32 = GRANTOR_UID;
+
+    const USER_ID: u32 = 99;
+    const APPLICATION_ID: u32 = 10001;
+    static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
+    static GRANTEE_GID: u32 = GRANTEE_UID;
+
+    const SEC_USER_ID: u32 = 98;
+    const SEC_APPLICATION_ID: u32 = 10001;
+    static SEC_GRANTEE_UID: u32 = SEC_USER_ID * AID_USER_OFFSET + SEC_APPLICATION_ID;
+    static SEC_GRANTEE_GID: u32 = SEC_GRANTEE_UID;
+
+    // Child function to generate a key and grant it to a user with `GET_INFO` permission.
+    let grantor_fn = || {
+        let sl = SecLevel::tee();
+        let access_vector = KeyPermission::GET_INFO.0;
+        let alias = format!("ks_grant_single_{}", getuid());
+        let mut grant_keys = generate_ec_key_and_grant_to_users(
+            &sl,
+            Some(alias),
+            vec![GRANTEE_UID.try_into().unwrap()],
+            access_vector,
+        )
+        .unwrap();
+
+        grant_keys.remove(0)
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    let grant_key_nspace = unsafe { run_as::run_as_app(GRANTOR_UID, GRANTOR_GID, grantor_fn) };
+
+    // Child function for the grantee context: can load the granted key.
+    let grantee_fn = move || {
+        let keystore2 = get_keystore_service();
+        let rsp = get_granted_key(&keystore2, grant_key_nspace).expect("failed to get granted key");
+
+        // Return the underlying key ID to simulate an ID leak.
+        assert_eq!(rsp.metadata.key.domain, Domain::KEY_ID);
+        rsp.metadata.key.nspace
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    let key_id = unsafe { run_as::run_as_app(GRANTEE_UID, GRANTEE_GID, grantee_fn) };
+
+    // Second context does not have access to the above granted key, because it's identified
+    // by <uid, grant_nspace> and the implicit uid value is different.  Also, even if the
+    // second context gets hold of the key ID somehow, that also doesn't work.
+    let non_grantee_fn = move || {
+        let keystore2 = get_keystore_service();
+        let result = get_granted_key(&keystore2, grant_key_nspace);
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
+
+        let result = map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
+            domain: Domain::KEY_ID,
+            nspace: key_id,
+            alias: None,
+            blob: None,
+        }));
+        assert!(result.is_err());
+        assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
+    };
+
+    // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
+    // `--test-threads=1`), and nothing yet done with binder on the main thread.
+    unsafe { run_as::run_as_app(SEC_GRANTEE_UID, SEC_GRANTEE_GID, non_grantee_fn) };
+}
+
+/// Try to grant an APP key with `GRANT` access. Keystore2 system shouldn't allow to grant a key
+/// with `GRANT` access. Test should fail to grant a key with `PERMISSION_DENIED` error response
+/// code.
+#[test]
+fn grant_app_key_with_grant_perm_fails() {
     let sl = SecLevel::tee();
     let access_vector = KeyPermission::GRANT.0;
     let alias = format!("ks_grant_access_vec_key_{}", getuid());
@@ -382,7 +565,7 @@
     let application_id = 10001;
     let grantee_uid = user_id * AID_USER_OFFSET + application_id;
 
-    let result = key_generations::map_ks_error(generate_ec_key_and_grant_to_users(
+    let result = map_ks_error(generate_ec_key_and_grant_to_users(
         &sl,
         Some(alias),
         vec![grantee_uid.try_into().unwrap()],
@@ -392,10 +575,10 @@
     assert_eq!(Error::Rc(ResponseCode::PERMISSION_DENIED), result.unwrap_err());
 }
 
-/// Try to grant a non-existing key to the user. Test should fail with `KEY_NOT_FOUND` error
+/// Try to grant a non-existing SELINUX key to the user. Test should fail with `KEY_NOT_FOUND` error
 /// response.
 #[test]
-fn keystore2_grant_fails_with_non_existing_key_expect_key_not_found_err() {
+fn grant_fails_with_non_existing_selinux_key() {
     let keystore2 = get_keystore_service();
     let alias = format!("ks_grant_test_non_existing_key_5_{}", getuid());
     let user_id = 98;
@@ -403,7 +586,7 @@
     let grantee_uid = user_id * AID_USER_OFFSET + application_id;
     let access_vector = KeyPermission::GET_INFO.0;
 
-    let result = key_generations::map_ks_error(keystore2.grant(
+    let result = map_ks_error(keystore2.grant(
         &KeyDescriptor {
             domain: Domain::SELINUX,
             nspace: key_generations::SELINUX_SHELL_NAMESPACE,
@@ -417,10 +600,10 @@
     assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
 }
 
-/// Grant a key to the user and immediately ungrant the granted key. In grantee context try to load
+/// Grant an APP key to the user and immediately ungrant the granted key. In grantee context try to load
 /// the key. Grantee should fail to load the ungranted key with `KEY_NOT_FOUND` error response.
 #[test]
-fn keystore2_ungrant_key_success() {
+fn ungrant_app_key_success() {
     const USER_ID: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
@@ -459,12 +642,7 @@
     // Grantee context, try to load the ungranted key.
     let grantee_fn = move || {
         let keystore2 = get_keystore_service();
-        let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
-            domain: Domain::GRANT,
-            nspace: grant_key_nspace,
-            alias: None,
-            blob: None,
-        }));
+        let result = get_granted_key(&keystore2, grant_key_nspace);
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
     };
@@ -480,7 +658,7 @@
 /// key in grantee context. Test should fail to load the granted key in grantee context as the
 /// associated key is deleted from grantor context.
 #[test]
-fn keystore2_ungrant_fails_with_non_existing_key_expect_key_not_found_error() {
+fn ungrant_deleted_app_key_fails() {
     const APPLICATION_ID: u32 = 10001;
     const USER_ID: u32 = 99;
     static GRANTEE_UID: u32 = USER_ID * AID_USER_OFFSET + APPLICATION_ID;
@@ -510,9 +688,8 @@
         sl.keystore2.deleteKey(&key_metadata.key).unwrap();
 
         // Try to ungrant above granted key.
-        let result = key_generations::map_ks_error(
-            sl.keystore2.ungrant(&key_metadata.key, GRANTEE_UID.try_into().unwrap()),
-        );
+        let result =
+            map_ks_error(sl.keystore2.ungrant(&key_metadata.key, GRANTEE_UID.try_into().unwrap()));
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
 
@@ -539,13 +716,7 @@
     // grantor context.
     let grantee_fn = move || {
         let keystore2 = get_keystore_service();
-
-        let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
-            domain: Domain::GRANT,
-            nspace: grant_key_nspace,
-            alias: None,
-            blob: None,
-        }));
+        let result = get_granted_key(&keystore2, grant_key_nspace);
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
     };
@@ -558,7 +729,7 @@
 /// Grant a key to multiple users. Verify that all grantees should succeed in loading the key and
 /// use it for performing an operation successfully.
 #[test]
-fn keystore2_grant_key_to_multi_users_success() {
+fn grant_app_key_to_multi_users_success() {
     const APPLICATION_ID: u32 = 10001;
     const USER_ID_1: u32 = 99;
     static GRANTEE_1_UID: u32 = USER_ID_1 * AID_USER_OFFSET + APPLICATION_ID;
@@ -592,15 +763,7 @@
     {
         let grant_key_nspace = grant_keys.remove(0);
         let grantee_fn = move || {
-            let sl = SecLevel::tee();
-
-            assert_eq!(
-                Ok(()),
-                key_generations::map_ks_error(load_grant_key_and_perform_sign_operation(
-                    &sl,
-                    grant_key_nspace
-                ))
-            );
+            assert_eq!(Ok(()), sign_with_granted_key(grant_key_nspace));
         };
         // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
         // `--test-threads=1`), and nothing yet done with binder.
@@ -612,7 +775,7 @@
 /// use the key and delete it. Try to load the granted key in another grantee context. Test should
 /// fail to load the granted key with `KEY_NOT_FOUND` error response.
 #[test]
-fn keystore2_grant_key_to_multi_users_delete_fails_with_key_not_found_error() {
+fn grant_app_key_to_multi_users_delete_then_key_not_found() {
     const USER_ID_1: u32 = 99;
     const APPLICATION_ID: u32 = 10001;
     static GRANTEE_1_UID: u32 = USER_ID_1 * AID_USER_OFFSET + APPLICATION_ID;
@@ -645,25 +808,10 @@
     // Grantee #1 context
     let grant_key1_nspace = grant_keys.remove(0);
     let grantee1_fn = move || {
-        let sl = SecLevel::tee();
-
-        assert_eq!(
-            Ok(()),
-            key_generations::map_ks_error(load_grant_key_and_perform_sign_operation(
-                &sl,
-                grant_key1_nspace
-            ))
-        );
+        assert_eq!(Ok(()), sign_with_granted_key(grant_key1_nspace));
 
         // Delete the granted key.
-        sl.keystore2
-            .deleteKey(&KeyDescriptor {
-                domain: Domain::GRANT,
-                nspace: grant_key1_nspace,
-                alias: None,
-                blob: None,
-            })
-            .unwrap();
+        get_keystore_service().deleteKey(&granted_key_descriptor(grant_key1_nspace)).unwrap();
     };
 
     // Safety: only one thread at this point (enforced by `AndroidTest.xml` setting
@@ -675,12 +823,8 @@
     let grantee2_fn = move || {
         let keystore2 = get_keystore_service();
 
-        let result = key_generations::map_ks_error(keystore2.getKeyEntry(&KeyDescriptor {
-            domain: Domain::GRANT,
-            nspace: grant_key2_nspace,
-            alias: None,
-            blob: None,
-        }));
+        let result =
+            map_ks_error(keystore2.getKeyEntry(&granted_key_descriptor(grant_key2_nspace)));
         assert!(result.is_err());
         assert_eq!(Error::Rc(ResponseCode::KEY_NOT_FOUND), result.unwrap_err());
     };
diff --git a/keystore2/tests/user_auth.rs b/keystore2/tests/user_auth.rs
index 336af4f..187256b 100644
--- a/keystore2/tests/user_auth.rs
+++ b/keystore2/tests/user_auth.rs
@@ -43,12 +43,9 @@
     run_as::{ChannelReader, ChannelWriter}, expect_km_error,
 };
 use log::{warn, info};
-use nix::unistd::{Gid, Uid};
 use rustutils::users::AID_USER_OFFSET;
 use std::{time::Duration, thread::sleep};
 
-/// SELinux context.
-const CTX: &str = "u:r:untrusted_app:s0:c91,c256,c10,c20";
 /// Test user ID.
 const TEST_USER_ID: i32 = 100;
 /// Corresponding uid value.
@@ -263,7 +260,7 @@
     // `--test-threads=1`), and nothing yet done with binder.
     let mut child_handle = unsafe {
         // Perform keystore actions while running as the test user.
-        run_as::run_as_child(CTX, Uid::from_raw(UID), Gid::from_raw(UID), child_fn)
+        run_as::run_as_child_app(UID, UID, child_fn)
     }
     .unwrap();
 
@@ -385,7 +382,7 @@
     // `--test-threads=1`), and nothing yet done with binder.
     let mut child_handle = unsafe {
         // Perform keystore actions while running as the test user.
-        run_as::run_as_child(CTX, Uid::from_raw(UID), Gid::from_raw(UID), child_fn)
+        run_as::run_as_child_app(UID, UID, child_fn)
     }
     .unwrap();
 
@@ -514,7 +511,7 @@
     // `--test-threads=1`), and nothing yet done with binder.
     let mut child_handle = unsafe {
         // Perform keystore actions while running as the test user.
-        run_as::run_as_child(CTX, Uid::from_raw(UID), Gid::from_raw(UID), child_fn)
+        run_as::run_as_child_app(UID, UID, child_fn)
     }
     .unwrap();
 
@@ -659,7 +656,7 @@
     // `--test-threads=1`), and nothing yet done with binder.
     let mut child_handle = unsafe {
         // Perform keystore actions while running as the test user.
-        run_as::run_as_child(CTX, Uid::from_raw(UID), Gid::from_raw(UID), child_fn)
+        run_as::run_as_child_app(UID, UID, child_fn)
     }
     .unwrap();
 
@@ -788,7 +785,7 @@
     // `--test-threads=1`), and nothing yet done with binder.
     let mut child_handle = unsafe {
         // Perform keystore actions while running as the test user.
-        run_as::run_as_child(CTX, Uid::from_raw(UID), Gid::from_raw(UID), child_fn)
+        run_as::run_as_child_app(UID, UID, child_fn)
     }
     .unwrap();
 
diff --git a/provisioner/rkp_factory_extraction_lib.cpp b/provisioner/rkp_factory_extraction_lib.cpp
index 37b81fd..d4aed45 100644
--- a/provisioner/rkp_factory_extraction_lib.cpp
+++ b/provisioner/rkp_factory_extraction_lib.cpp
@@ -25,7 +25,6 @@
 #include <cstring>
 #include <iterator>
 #include <keymaster/cppcose/cppcose.h>
-#include <openssl/base64.h>
 #include <remote_prov/remote_prov_utils.h>
 #include <sys/random.h>
 
@@ -33,6 +32,7 @@
 #include <optional>
 #include <string>
 #include <string_view>
+#include <unordered_set>
 #include <vector>
 
 #include "cppbor_parse.h"
@@ -42,6 +42,7 @@
 using aidl::android::hardware::security::keymint::MacedPublicKey;
 using aidl::android::hardware::security::keymint::ProtectedData;
 using aidl::android::hardware::security::keymint::RpcHardwareInfo;
+using aidl::android::hardware::security::keymint::remote_prov::BccEntryData;
 using aidl::android::hardware::security::keymint::remote_prov::EekChain;
 using aidl::android::hardware::security::keymint::remote_prov::generateEekChain;
 using aidl::android::hardware::security::keymint::remote_prov::getProdEekChain;
@@ -50,35 +51,13 @@
 using aidl::android::hardware::security::keymint::remote_prov::verifyFactoryCsr;
 using aidl::android::hardware::security::keymint::remote_prov::verifyFactoryProtectedData;
 
-using namespace cppbor;
-using namespace cppcose;
+using cppbor::Array;
+using cppbor::Map;
+using cppbor::Null;
+template <class T> using ErrMsgOr = cppcose::ErrMsgOr<T>;
 
 constexpr size_t kVersionWithoutSuperencryption = 3;
 
-std::string toBase64(const std::vector<uint8_t>& buffer) {
-    size_t base64Length;
-    int rc = EVP_EncodedLength(&base64Length, buffer.size());
-    if (!rc) {
-        std::cerr << "Error getting base64 length. Size overflow?" << std::endl;
-        exit(-1);
-    }
-
-    std::string base64(base64Length, ' ');
-    rc = EVP_EncodeBlock(reinterpret_cast<uint8_t*>(base64.data()), buffer.data(), buffer.size());
-    ++rc;  // Account for NUL, which BoringSSL does not for some reason.
-    if (rc != base64Length) {
-        std::cerr << "Error writing base64. Expected " << base64Length
-                  << " bytes to be written, but " << rc << " bytes were actually written."
-                  << std::endl;
-        exit(-1);
-    }
-
-    // BoringSSL automatically adds a NUL -- remove it from the string data
-    base64.pop_back();
-
-    return base64;
-}
-
 std::vector<uint8_t> generateChallenge() {
     std::vector<uint8_t> challenge(kChallengeSize);
 
@@ -90,7 +69,8 @@
             if (errno == EINTR) {
                 continue;
             } else {
-                std::cerr << errno << ": " << strerror(errno) << std::endl;
+                std::cerr << "generateChallenge: getrandom returned an error with errno " << errno
+                          << ": " << strerror(errno) << ". Exiting..." << std::endl;
                 exit(-1);
             }
         }
@@ -118,7 +98,7 @@
         return {nullptr, parsedVerifiedDeviceInfo.moveMessage()};
     }
 
-    auto [parsedProtectedData, ignore2, errMsg] = parse(protectedData.protectedData);
+    auto [parsedProtectedData, ignore2, errMsg] = cppbor::parse(protectedData.protectedData);
     if (!parsedProtectedData) {
         std::cerr << "Error parsing protected data: '" << errMsg << "'" << std::endl;
         return {nullptr, errMsg};
@@ -145,7 +125,7 @@
     if (!status.isOk()) {
         std::cerr << "Failed to get hardware info for '" << componentName
                   << "'. Description: " << status.getDescription() << "." << std::endl;
-        exit(-1);
+        return {nullptr, status.getDescription()};
     }
 
     const std::vector<uint8_t> eek = getProdEekChain(hwInfo.supportedEekCurve);
@@ -156,13 +136,14 @@
     if (!status.isOk()) {
         std::cerr << "Bundle extraction failed for '" << componentName
                   << "'. Description: " << status.getDescription() << "." << std::endl;
-        exit(-1);
+        return {nullptr, status.getDescription()};
     }
     return composeCertificateRequestV1(protectedData, verifiedDeviceInfo, challenge, keysToSignMac,
                                        irpc);
 }
 
-void selfTestGetCsrV1(std::string_view componentName, IRemotelyProvisionedComponent* irpc) {
+std::optional<std::string> selfTestGetCsrV1(std::string_view componentName,
+                                            IRemotelyProvisionedComponent* irpc) {
     std::vector<uint8_t> keysToSignMac;
     std::vector<MacedPublicKey> emptyKeys;
     DeviceInfo verifiedDeviceInfo;
@@ -172,14 +153,14 @@
     if (!status.isOk()) {
         std::cerr << "Failed to get hardware info for '" << componentName
                   << "'. Description: " << status.getDescription() << "." << std::endl;
-        exit(-1);
+        return status.getDescription();
     }
 
     const std::vector<uint8_t> eekId = {0, 1, 2, 3, 4, 5, 6, 7};
     ErrMsgOr<EekChain> eekChain = generateEekChain(hwInfo.supportedEekCurve, /*length=*/3, eekId);
     if (!eekChain) {
         std::cerr << "Error generating test EEK certificate chain: " << eekChain.message();
-        exit(-1);
+        return eekChain.message();
     }
     const std::vector<uint8_t> challenge = generateChallenge();
     status = irpc->generateCertificateRequest(
@@ -188,7 +169,7 @@
     if (!status.isOk()) {
         std::cerr << "Error generating test cert chain for '" << componentName
                   << "'. Description: " << status.getDescription() << "." << std::endl;
-        exit(-1);
+        return status.getDescription();
     }
 
     auto result = verifyFactoryProtectedData(
@@ -198,8 +179,9 @@
     if (!result) {
         std::cerr << "Self test failed for IRemotelyProvisionedComponent '" << componentName
                   << "'. Error message: '" << result.message() << "'." << std::endl;
-        exit(-1);
+        return result.message();
     }
+    return std::nullopt;
 }
 
 CborResult<Array> composeCertificateRequestV3(const std::vector<uint8_t>& csr) {
@@ -223,9 +205,8 @@
     return {std::unique_ptr<Array>(parsedCsr.release()->asArray()), ""};
 }
 
-CborResult<cppbor::Array> getCsrV3(std::string_view componentName,
-                                   IRemotelyProvisionedComponent* irpc, bool selfTest,
-                                   bool allowDegenerate) {
+CborResult<Array> getCsrV3(std::string_view componentName, IRemotelyProvisionedComponent* irpc,
+                           bool selfTest, bool allowDegenerate, bool requireUdsCerts) {
     std::vector<uint8_t> csr;
     std::vector<MacedPublicKey> emptyKeys;
     const std::vector<uint8_t> challenge = generateChallenge();
@@ -234,16 +215,17 @@
     if (!status.isOk()) {
         std::cerr << "Bundle extraction failed for '" << componentName
                   << "'. Description: " << status.getDescription() << "." << std::endl;
-        exit(-1);
+        return {nullptr, status.getDescription()};
     }
 
     if (selfTest) {
-        auto result = verifyFactoryCsr(/*keysToSign=*/cppbor::Array(), csr, irpc,
-                                       std::string(componentName), challenge, allowDegenerate);
+        auto result =
+            verifyFactoryCsr(/*keysToSign=*/cppbor::Array(), csr, irpc, std::string(componentName),
+                             challenge, allowDegenerate, requireUdsCerts);
         if (!result) {
             std::cerr << "Self test failed for IRemotelyProvisionedComponent '" << componentName
                       << "'. Error message: '" << result.message() << "'." << std::endl;
-            exit(-1);
+            return {nullptr, result.message()};
         }
     }
 
@@ -251,35 +233,35 @@
 }
 
 CborResult<Array> getCsr(std::string_view componentName, IRemotelyProvisionedComponent* irpc,
-                         bool selfTest, bool allowDegenerate) {
+                         bool selfTest, bool allowDegenerate, bool requireUdsCerts) {
     RpcHardwareInfo hwInfo;
     auto status = irpc->getHardwareInfo(&hwInfo);
     if (!status.isOk()) {
         std::cerr << "Failed to get hardware info for '" << componentName
                   << "'. Description: " << status.getDescription() << "." << std::endl;
-        exit(-1);
+        return {nullptr, status.getDescription()};
     }
 
     if (hwInfo.versionNumber < kVersionWithoutSuperencryption) {
         if (selfTest) {
-            selfTestGetCsrV1(componentName, irpc);
+            auto errMsg = selfTestGetCsrV1(componentName, irpc);
+            if (errMsg) {
+                return {nullptr, *errMsg};
+            }
         }
         return getCsrV1(componentName, irpc);
     } else {
-        return getCsrV3(componentName, irpc, selfTest, allowDegenerate);
+        return getCsrV3(componentName, irpc, selfTest, allowDegenerate, requireUdsCerts);
     }
 }
 
-bool isRemoteProvisioningSupported(IRemotelyProvisionedComponent* irpc) {
-    RpcHardwareInfo hwInfo;
-    auto status = irpc->getHardwareInfo(&hwInfo);
-    if (status.isOk()) {
-        return true;
+std::unordered_set<std::string> parseCommaDelimited(const std::string& input) {
+    std::stringstream ss(input);
+    std::unordered_set<std::string> result;
+    while (ss.good()) {
+        std::string name;
+        std::getline(ss, name, ',');
+        result.insert(name);
     }
-    if (status.getExceptionCode() == EX_UNSUPPORTED_OPERATION) {
-        return false;
-    }
-    std::cerr << "Unexpected error when getting hardware info. Description: "
-              << status.getDescription() << "." << std::endl;
-    exit(-1);
-}
+    return result;
+}
\ No newline at end of file
diff --git a/provisioner/rkp_factory_extraction_lib.h b/provisioner/rkp_factory_extraction_lib.h
index 94bd751..2c1e2ff 100644
--- a/provisioner/rkp_factory_extraction_lib.h
+++ b/provisioner/rkp_factory_extraction_lib.h
@@ -23,8 +23,12 @@
 #include <memory>
 #include <string>
 #include <string_view>
+#include <unordered_set>
 #include <vector>
 
+// Parse a comma-delimited string.
+std::unordered_set<std::string> parseCommaDelimited(const std::string& input);
+
 // Challenge size must be between 32 and 64 bytes inclusive.
 constexpr size_t kChallengeSize = 64;
 
@@ -35,9 +39,6 @@
     std::string errMsg;
 };
 
-// Return `buffer` encoded as a base64 string.
-std::string toBase64(const std::vector<uint8_t>& buffer);
-
 // Generate a random challenge containing `kChallengeSize` bytes.
 std::vector<uint8_t> generateChallenge();
 
@@ -47,13 +48,4 @@
 CborResult<cppbor::Array>
 getCsr(std::string_view componentName,
        aidl::android::hardware::security::keymint::IRemotelyProvisionedComponent* irpc,
-       bool selfTest, bool allowDegenerate);
-
-// Generates a test certificate chain and validates it, exiting the process on error.
-void selfTestGetCsr(
-    std::string_view componentName,
-    aidl::android::hardware::security::keymint::IRemotelyProvisionedComponent* irpc);
-
-// Returns true if the given IRemotelyProvisionedComponent supports remote provisioning.
-bool isRemoteProvisioningSupported(
-    aidl::android::hardware::security::keymint::IRemotelyProvisionedComponent* irpc);
+       bool selfTest, bool allowDegenerate, bool requireUdsCerts);
\ No newline at end of file
diff --git a/provisioner/rkp_factory_extraction_lib_test.cpp b/provisioner/rkp_factory_extraction_lib_test.cpp
index 247c508..746ce41 100644
--- a/provisioner/rkp_factory_extraction_lib_test.cpp
+++ b/provisioner/rkp_factory_extraction_lib_test.cpp
@@ -25,6 +25,8 @@
 #include <android-base/properties.h>
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include <openssl/base64.h>
+#include <remote_prov/MockIRemotelyProvisionedComponent.h>
 
 #include <cstdint>
 #include <memory>
@@ -60,26 +62,29 @@
 
 }  // namespace cppbor
 
-class MockIRemotelyProvisionedComponent : public IRemotelyProvisionedComponentDefault {
-  public:
-    MOCK_METHOD(ScopedAStatus, getHardwareInfo, (RpcHardwareInfo * _aidl_return), (override));
-    MOCK_METHOD(ScopedAStatus, generateEcdsaP256KeyPair,
-                (bool in_testMode, MacedPublicKey* out_macedPublicKey,
-                 std::vector<uint8_t>* _aidl_return),
-                (override));
-    MOCK_METHOD(ScopedAStatus, generateCertificateRequest,
-                (bool in_testMode, const std::vector<MacedPublicKey>& in_keysToSign,
-                 const std::vector<uint8_t>& in_endpointEncryptionCertChain,
-                 const std::vector<uint8_t>& in_challenge, DeviceInfo* out_deviceInfo,
-                 ProtectedData* out_protectedData, std::vector<uint8_t>* _aidl_return),
-                (override));
-    MOCK_METHOD(ScopedAStatus, generateCertificateRequestV2,
-                (const std::vector<MacedPublicKey>& in_keysToSign,
-                 const std::vector<uint8_t>& in_challenge, std::vector<uint8_t>* _aidl_return),
-                (override));
-    MOCK_METHOD(ScopedAStatus, getInterfaceVersion, (int32_t * _aidl_return), (override));
-    MOCK_METHOD(ScopedAStatus, getInterfaceHash, (std::string * _aidl_return), (override));
-};
+std::string toBase64(const std::vector<uint8_t>& buffer) {
+    size_t base64Length;
+    int rc = EVP_EncodedLength(&base64Length, buffer.size());
+    if (!rc) {
+        std::cerr << "Error getting base64 length. Size overflow?" << std::endl;
+        exit(-1);
+    }
+
+    std::string base64(base64Length, ' ');
+    rc = EVP_EncodeBlock(reinterpret_cast<uint8_t*>(base64.data()), buffer.data(), buffer.size());
+    ++rc;  // Account for NUL, which BoringSSL does not for some reason.
+    if (rc != base64Length) {
+        std::cerr << "Error writing base64. Expected " << base64Length
+                  << " bytes to be written, but " << rc << " bytes were actually written."
+                  << std::endl;
+        exit(-1);
+    }
+
+    // BoringSSL automatically adds a NUL -- remove it from the string data
+    base64.pop_back();
+
+    return base64;
+}
 
 TEST(LibRkpFactoryExtractionTests, ToBase64) {
     std::vector<uint8_t> input(UINT8_MAX + 1);
@@ -87,7 +92,7 @@
         input[i] = i;
     }
 
-    // Test three lengths so we get all the different paddding options
+    // Test three lengths so we get all the different padding options
     EXPECT_EQ("AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8gISIjJCUmJygpKissLS4"
               "vMDEyMzQ1Njc4OTo7PD0+P0BBQkNERUZHSElKS0xNTk9QUVJTVFVWV1hZWltcXV"
               "5fYGFiY2RlZmdoaWprbG1ub3BxcnN0dXZ3eHl6e3x9fn+AgYKDhIWGh4iJiouMj"
@@ -160,7 +165,7 @@
     std::vector<uint8_t> challenge;
 
     // Set up mock, then call getSCsr
-    auto mockRpc = SharedRefBase::make<MockIRemotelyProvisionedComponent>();
+    auto mockRpc = SharedRefBase::make<remote_prov::MockIRemotelyProvisionedComponent>();
     EXPECT_CALL(*mockRpc, getHardwareInfo(NotNull())).WillRepeatedly([](RpcHardwareInfo* hwInfo) {
         hwInfo->versionNumber = 2;
         return ScopedAStatus::ok();
@@ -180,8 +185,9 @@
                         SetArgPointee<6>(kFakeMac),             //
                         Return(ByMove(ScopedAStatus::ok()))));  //
 
-    auto [csr, csrErrMsg] = getCsr("mock component name", mockRpc.get(),
-                                   /*selfTest=*/false, /*allowDegenerate=*/true);
+    auto [csr, csrErrMsg] =
+        getCsr("mock component name", mockRpc.get(),
+               /*selfTest=*/false, /*allowDegenerate=*/true, /*requireUdsCerts=*/false);
     ASSERT_THAT(csr, NotNull()) << csrErrMsg;
     ASSERT_THAT(csr->asArray(), Pointee(Property(&Array::size, Eq(4))));
 
@@ -230,7 +236,7 @@
 
 TEST(LibRkpFactoryExtractionTests, GetCsrWithV3Hal) {
     const std::vector<uint8_t> kCsr = Array()
-                                          .add(3 /* version */)
+                                          .add(1 /* version */)
                                           .add(Map() /* UdsCerts */)
                                           .add(Array() /* DiceCertChain */)
                                           .add(Array() /* SignedData */)
@@ -238,7 +244,7 @@
     std::vector<uint8_t> challenge;
 
     // Set up mock, then call getCsr
-    auto mockRpc = SharedRefBase::make<MockIRemotelyProvisionedComponent>();
+    auto mockRpc = SharedRefBase::make<remote_prov::MockIRemotelyProvisionedComponent>();
     EXPECT_CALL(*mockRpc, getHardwareInfo(NotNull())).WillRepeatedly([](RpcHardwareInfo* hwInfo) {
         hwInfo->versionNumber = 3;
         return ScopedAStatus::ok();
@@ -250,12 +256,13 @@
         .WillOnce(DoAll(SaveArg<1>(&challenge), SetArgPointee<2>(kCsr),
                         Return(ByMove(ScopedAStatus::ok()))));
 
-    auto [csr, csrErrMsg] = getCsr("mock component name", mockRpc.get(),
-                                   /*selfTest=*/false, /*allowDegenerate=*/true);
+    auto [csr, csrErrMsg] =
+        getCsr("mock component name", mockRpc.get(),
+               /*selfTest=*/false, /*allowDegenerate=*/true, /*requireUdsCerts=*/false);
     ASSERT_THAT(csr, NotNull()) << csrErrMsg;
     ASSERT_THAT(csr, Pointee(Property(&Array::size, Eq(5))));
 
-    EXPECT_THAT(csr->get(0 /* version */), Pointee(Eq(Uint(3))));
+    EXPECT_THAT(csr->get(0 /* version */), Pointee(Eq(Uint(1))));
     EXPECT_THAT(csr->get(1)->asMap(), NotNull());
     EXPECT_THAT(csr->get(2)->asArray(), NotNull());
     EXPECT_THAT(csr->get(3)->asArray(), NotNull());
@@ -266,3 +273,72 @@
     const Tstr fingerprint(android::base::GetProperty("ro.build.fingerprint", ""));
     EXPECT_THAT(*unverifedDeviceInfo->get("fingerprint")->asTstr(), Eq(fingerprint));
 }
+
+TEST(LibRkpFactoryExtractionTests, requireUdsCerts) {
+    const std::vector<uint8_t> kCsr = Array()
+                                          .add(1 /* version */)
+                                          .add(Map() /* UdsCerts */)
+                                          .add(Array() /* DiceCertChain */)
+                                          .add(Array() /* SignedData */)
+                                          .encode();
+    std::vector<uint8_t> challenge;
+
+    // Set up mock, then call getCsr
+    auto mockRpc = SharedRefBase::make<remote_prov::MockIRemotelyProvisionedComponent>();
+    EXPECT_CALL(*mockRpc, getHardwareInfo(NotNull())).WillRepeatedly([](RpcHardwareInfo* hwInfo) {
+        hwInfo->versionNumber = 3;
+        return ScopedAStatus::ok();
+    });
+    EXPECT_CALL(*mockRpc,
+                generateCertificateRequestV2(IsEmpty(),   // keysToSign
+                                             _,           // challenge
+                                             NotNull()))  // _aidl_return
+        .WillOnce(DoAll(SaveArg<1>(&challenge), SetArgPointee<2>(kCsr),
+                        Return(ByMove(ScopedAStatus::ok()))));
+
+    auto [csr, csrErrMsg] =
+        getCsr("mock component name", mockRpc.get(),
+               /*selfTest=*/true, /*allowDegenerate=*/false, /*requireUdsCerts=*/true);
+    ASSERT_EQ(csr, nullptr);
+    ASSERT_THAT(csrErrMsg, testing::HasSubstr("UdsCerts must not be empty"));
+}
+
+TEST(LibRkpFactoryExtractionTests, dontRequireUdsCerts) {
+    const std::vector<uint8_t> kCsr = Array()
+                                          .add(1 /* version */)
+                                          .add(Map() /* UdsCerts */)
+                                          .add(Array() /* DiceCertChain */)
+                                          .add(Array() /* SignedData */)
+                                          .encode();
+    std::vector<uint8_t> challenge;
+
+    // Set up mock, then call getCsr
+    auto mockRpc = SharedRefBase::make<remote_prov::MockIRemotelyProvisionedComponent>();
+    EXPECT_CALL(*mockRpc, getHardwareInfo(NotNull())).WillRepeatedly([](RpcHardwareInfo* hwInfo) {
+        hwInfo->versionNumber = 3;
+        return ScopedAStatus::ok();
+    });
+    EXPECT_CALL(*mockRpc,
+                generateCertificateRequestV2(IsEmpty(),   // keysToSign
+                                             _,           // challenge
+                                             NotNull()))  // _aidl_return
+        .WillOnce(DoAll(SaveArg<1>(&challenge), SetArgPointee<2>(kCsr),
+                        Return(ByMove(ScopedAStatus::ok()))));
+
+    auto [csr, csrErrMsg] =
+        getCsr("mock component name", mockRpc.get(),
+               /*selfTest=*/true, /*allowDegenerate=*/false, /*requireUdsCerts=*/false);
+    ASSERT_EQ(csr, nullptr);
+    ASSERT_THAT(csrErrMsg, testing::Not(testing::HasSubstr("UdsCerts must not be empty")));
+}
+
+TEST(LibRkpFactoryExtractionTests, parseCommaDelimitedString) {
+    const auto& rpcNames = "default,avf,,default,Strongbox,strongbox,,";
+    const auto& rpcSet = parseCommaDelimited(rpcNames);
+
+    ASSERT_EQ(rpcSet.size(), 5);
+    ASSERT_TRUE(rpcSet.count("default") == 1);
+    ASSERT_TRUE(rpcSet.count("avf") == 1);
+    ASSERT_TRUE(rpcSet.count("strongbox") == 1);
+    ASSERT_TRUE(rpcSet.count("Strongbox") == 1);
+}
\ No newline at end of file
diff --git a/provisioner/rkp_factory_extraction_tool.cpp b/provisioner/rkp_factory_extraction_tool.cpp
index c0f6beb..eaa0acc 100644
--- a/provisioner/rkp_factory_extraction_tool.cpp
+++ b/provisioner/rkp_factory_extraction_tool.cpp
@@ -26,6 +26,7 @@
 
 #include <future>
 #include <string>
+#include <unordered_set>
 #include <vector>
 
 #include "DrmRkpAdapter.h"
@@ -33,10 +34,10 @@
 
 using aidl::android::hardware::drm::IDrmFactory;
 using aidl::android::hardware::security::keymint::IRemotelyProvisionedComponent;
+using aidl::android::hardware::security::keymint::RpcHardwareInfo;
+using aidl::android::hardware::security::keymint::remote_prov::deviceSuffix;
 using aidl::android::hardware::security::keymint::remote_prov::jsonEncodeCsrWithBuild;
-
-using namespace cppbor;
-using namespace cppcose;
+using aidl::android::hardware::security::keymint::remote_prov::RKPVM_INSTANCE_NAME;
 
 DEFINE_string(output_format, "build+csr", "How to format the output. Defaults to 'build+csr'.");
 DEFINE_bool(self_test, true,
@@ -47,6 +48,10 @@
             "If true, self_test validation will allow degenerate DICE chains in the CSR.");
 DEFINE_string(serialno_prop, "ro.serialno",
               "The property of getting serial number. Defaults to 'ro.serialno'.");
+DEFINE_string(require_uds_certs, "",
+              "The comma-delimited names of remotely provisioned "
+              "components whose UDS certificate chains are required to be present in the CSR. "
+              "Example: avf,default,strongbox");
 
 namespace {
 
@@ -59,15 +64,15 @@
     return  std::string(descriptor) + "/" + name;
 }
 
-void writeOutput(const std::string instance_name, const Array& csr) {
+void writeOutput(const std::string instance_name, const cppbor::Array& csr) {
     if (FLAGS_output_format == kBinaryCsrOutput) {
         auto bytes = csr.encode();
         std::copy(bytes.begin(), bytes.end(), std::ostream_iterator<char>(std::cout));
     } else if (FLAGS_output_format == kBuildPlusCsr) {
         auto [json, error] = jsonEncodeCsrWithBuild(instance_name, csr, FLAGS_serialno_prop);
         if (!error.empty()) {
-            std::cerr << "Error JSON encoding the output: " << error;
-            exit(1);
+            std::cerr << "Error JSON encoding the output: " << error << std::endl;
+            exit(-1);
         }
         std::cout << json << std::endl;
     } else {
@@ -75,20 +80,28 @@
         std::cerr << "Valid formats:" << std::endl;
         std::cerr << "  " << kBinaryCsrOutput << std::endl;
         std::cerr << "  " << kBuildPlusCsr << std::endl;
-        exit(1);
+        exit(-1);
     }
 }
 
-void getCsrForIRpc(const char* descriptor, const char* name, IRemotelyProvisionedComponent* irpc) {
+void getCsrForIRpc(const char* descriptor, const char* name, IRemotelyProvisionedComponent* irpc,
+                   bool requireUdsCerts) {
     // AVF RKP HAL is not always supported, so we need to check if it is supported before
     // generating the CSR.
-    if (std::string(name) == "avf" && !isRemoteProvisioningSupported(irpc)) {
-        return;
+    if (std::string(name) == deviceSuffix(RKPVM_INSTANCE_NAME)) {
+        RpcHardwareInfo hwInfo;
+        auto status = irpc->getHardwareInfo(&hwInfo);
+        if (!status.isOk()) {
+            return;
+        }
     }
-    auto [request, errMsg] = getCsr(name, irpc, FLAGS_self_test, FLAGS_allow_degenerate);
-    auto fullName = getFullServiceName(descriptor, name);
+
+    auto [request, errMsg] =
+        getCsr(name, irpc, FLAGS_self_test, FLAGS_allow_degenerate, requireUdsCerts);
     if (!request) {
-        std::cerr << "Unable to build CSR for '" << fullName << ": " << errMsg << std::endl;
+        auto fullName = getFullServiceName(descriptor, name);
+        std::cerr << "Unable to build CSR for '" << fullName << "': " << errMsg << ", exiting."
+                  << std::endl;
         exit(-1);
     }
 
@@ -97,23 +110,33 @@
 
 // Callback for AServiceManager_forEachDeclaredInstance that writes out a CSR
 // for every IRemotelyProvisionedComponent.
-void getCsrForInstance(const char* name, void* /*context*/) {
+void getCsrForInstance(const char* name, void* context) {
     auto fullName = getFullServiceName(IRemotelyProvisionedComponent::descriptor, name);
-    std::future<AIBinder*> wait_for_service_func =
+    std::future<AIBinder*> waitForServiceFunc =
         std::async(std::launch::async, AServiceManager_waitForService, fullName.c_str());
-    if (wait_for_service_func.wait_for(std::chrono::seconds(10)) == std::future_status::timeout) {
-        std::cerr << "Wait for service timed out after 10 seconds: " << fullName;
+    if (waitForServiceFunc.wait_for(std::chrono::seconds(10)) == std::future_status::timeout) {
+        std::cerr << "Wait for service timed out after 10 seconds: '" << fullName << "', exiting."
+                  << std::endl;
         exit(-1);
     }
-    AIBinder* rkpAiBinder = wait_for_service_func.get();
+    AIBinder* rkpAiBinder = waitForServiceFunc.get();
     ::ndk::SpAIBinder rkp_binder(rkpAiBinder);
-    auto rkp_service = IRemotelyProvisionedComponent::fromBinder(rkp_binder);
-    if (!rkp_service) {
-        std::cerr << "Unable to get binder object for '" << fullName << "', skipping.";
+    auto rkpService = IRemotelyProvisionedComponent::fromBinder(rkp_binder);
+    if (!rkpService) {
+        std::cerr << "Unable to get binder object for '" << fullName << "', exiting." << std::endl;
         exit(-1);
     }
 
-    getCsrForIRpc(IRemotelyProvisionedComponent::descriptor, name, rkp_service.get());
+    if (context == nullptr) {
+        std::cerr << "Unable to get context for '" << fullName << "', exiting." << std::endl;
+        exit(-1);
+    }
+
+    auto requireUdsCertsRpcNames = static_cast<std::unordered_set<std::string>*>(context);
+    auto requireUdsCerts = requireUdsCertsRpcNames->count(name) != 0;
+    requireUdsCertsRpcNames->erase(name);
+    getCsrForIRpc(IRemotelyProvisionedComponent::descriptor, name, rkpService.get(),
+                  requireUdsCerts);
 }
 
 }  // namespace
@@ -121,12 +144,21 @@
 int main(int argc, char** argv) {
     gflags::ParseCommandLineFlags(&argc, &argv, /*remove_flags=*/true);
 
+    auto requireUdsCertsRpcNames = parseCommaDelimited(FLAGS_require_uds_certs);
+
     AServiceManager_forEachDeclaredInstance(IRemotelyProvisionedComponent::descriptor,
-                                            /*context=*/nullptr, getCsrForInstance);
+                                            &requireUdsCertsRpcNames, getCsrForInstance);
 
     // Append drm csr's
-    for (auto const& e : android::mediadrm::getDrmRemotelyProvisionedComponents()) {
-        getCsrForIRpc(IDrmFactory::descriptor, e.first.c_str(), e.second.get());
+    for (auto const& [name, irpc] : android::mediadrm::getDrmRemotelyProvisionedComponents()) {
+        auto requireUdsCerts = requireUdsCertsRpcNames.count(name) != 0;
+        requireUdsCertsRpcNames.erase(name);
+        getCsrForIRpc(IDrmFactory::descriptor, name.c_str(), irpc.get(), requireUdsCerts);
+    }
+
+    for (auto const& rpcName : requireUdsCertsRpcNames) {
+        std::cerr << "WARNING: You requested to enforce the presence of UDS Certs for '" << rpcName
+                  << "', but no Remotely Provisioned Component had that name." << std::endl;
     }
 
     return 0;