Tests updated to handle RKP-ONLY property on GSI builds.

GSI replaces the values for remote_prov_prop properties (since they’re
system_internal_prop properties), so on GSI the properties are not
reliable indicators of whether StrongBox/TEE are RKP-only or not.

Tests are updated to stop running on GSI build if generateKey API
returns `ATTESTATION_KEYS_NOT_PROVISIONED` error code when RKP_ONLY
property is undetermined.

Impacted tests: test that generates a key with attestation but doesn't provide an ATTEST_KEY,
i.e. any test that requires batch keys.

Bug: 359743897
Test: atest keystore2_client_tests
Change-Id: I8028c20a47cf46aa807270e948e37ed6a08bfa02
diff --git a/keystore2/test_utils/key_generations.rs b/keystore2/test_utils/key_generations.rs
index c845332..258c68f 100644
--- a/keystore2/test_utils/key_generations.rs
+++ b/keystore2/test_utils/key_generations.rs
@@ -38,6 +38,7 @@
 use nix::unistd::getuid;
 use std::collections::HashSet;
 use std::fmt::Write;
+use std::path::PathBuf;
 
 /// Shell namespace.
 pub const SELINUX_SHELL_NAMESPACE: i64 = 1;
@@ -50,6 +51,10 @@
 /// Vold context
 pub const TARGET_VOLD_CTX: &str = "u:r:vold:s0";
 
+const TEE_KEYMINT_RKP_ONLY: &str = "remote_provisioning.tee.rkp_only";
+
+const STRONGBOX_KEYMINT_RKP_ONLY: &str = "remote_provisioning.strongbox.rkp_only";
+
 /// Allowed tags in generated/imported key authorizations.
 /// See hardware/interfaces/security/keymint/aidl/android/hardware/security/keymint/Tag.aidl for the
 /// list feature tags.
@@ -387,6 +392,33 @@
     })
 }
 
+/// Get the value of the given system property, if the given system property doesn't exist
+/// then returns an empty byte vector.
+pub fn get_system_prop(name: &str) -> Vec<u8> {
+    match rustutils::system_properties::read(name) {
+        Ok(Some(value)) => value.as_bytes().to_vec(),
+        _ => vec![],
+    }
+}
+
+/// Determines whether test is running on GSI.
+pub fn is_gsi() -> bool {
+    // This file is only present on GSI builds.
+    PathBuf::from("/system/system_ext/etc/init/init.gsi.rc").as_path().is_file()
+}
+
+/// Determines whether the test is on a GSI build where the rkp-only status of the device is
+/// unknown. GSI replaces the values for remote_prov_prop properties (since they’re
+/// system_internal_prop properties), so on GSI the properties are not reliable indicators of
+/// whether StrongBox/TEE is RKP-only or not.
+pub fn is_rkp_only_unknown_on_gsi(sec_level: SecurityLevel) -> bool {
+    if sec_level == SecurityLevel::TRUSTED_ENVIRONMENT {
+        is_gsi() && get_system_prop(TEE_KEYMINT_RKP_ONLY).is_empty()
+    } else {
+        is_gsi() && get_system_prop(STRONGBOX_KEYMINT_RKP_ONLY).is_empty()
+    }
+}
+
 /// Verify that given key param is listed in given authorizations list.
 pub fn check_key_param(authorizations: &[Authorization], key_param: &KeyParameter) -> bool {
     authorizations.iter().any(|auth| &auth.keyParameter == key_param)
@@ -634,7 +666,7 @@
     alias: Option<String>,
     key_params: &KeyParams,
     attest_key: Option<&KeyDescriptor>,
-) -> binder::Result<KeyMetadata> {
+) -> binder::Result<Option<KeyMetadata>> {
     let mut gen_params = AuthSetBuilder::new()
         .no_auth_required()
         .algorithm(Algorithm::RSA)
@@ -660,13 +692,29 @@
         gen_params = gen_params.attestation_challenge(value.to_vec())
     }
 
-    let key_metadata = sl.binder.generateKey(
+    let key_metadata = match sl.binder.generateKey(
         &KeyDescriptor { domain, nspace, alias, blob: None },
         attest_key,
         &gen_params,
         0,
         b"entropy",
-    )?;
+    ) {
+        Ok(metadata) => metadata,
+        Err(e) => {
+            return if is_rkp_only_unknown_on_gsi(sl.level)
+                && e.service_specific_error() == ErrorCode::ATTESTATION_KEYS_NOT_PROVISIONED.0
+            {
+                // GSI replaces the values for remote_prov_prop properties (since they’re
+                // system_internal_prop properties), so on GSI the properties are not
+                // reliable indicators of whether StrongBox/TEE are RKP-only or not.
+                // Test can be skipped if it generates a key with attestation but doesn't provide
+                // an ATTEST_KEY and rkp-only property is undetermined.
+                Ok(None)
+            } else {
+                Err(e)
+            };
+        }
+    };
 
     // Must have a public key.
     assert!(key_metadata.certificate.is_some());
@@ -697,7 +745,7 @@
             }
         ));
     }
-    Ok(key_metadata)
+    Ok(Some(key_metadata))
 }
 
 /// Generate AES/3DES key.
@@ -795,12 +843,12 @@
     sl: &SecLevel,
     algorithm: Algorithm,
     att_challenge: &[u8],
-) -> binder::Result<KeyMetadata> {
+) -> binder::Result<Option<KeyMetadata>> {
     assert!(algorithm == Algorithm::RSA || algorithm == Algorithm::EC);
 
     if algorithm == Algorithm::RSA {
         let alias = "ks_rsa_attest_test_key";
-        let metadata = generate_rsa_key(
+        generate_rsa_key(
             sl,
             Domain::APP,
             -1,
@@ -816,14 +864,8 @@
             },
             None,
         )
-        .unwrap();
-        Ok(metadata)
     } else {
-        let metadata =
-            generate_ec_attestation_key(sl, att_challenge, Digest::SHA_2_256, EcCurve::P_256)
-                .unwrap();
-
-        Ok(metadata)
+        generate_ec_attestation_key(sl, att_challenge, Digest::SHA_2_256, EcCurve::P_256)
     }
 }
 
@@ -834,7 +876,7 @@
     att_challenge: &[u8],
     digest: Digest,
     ec_curve: EcCurve,
-) -> binder::Result<KeyMetadata> {
+) -> binder::Result<Option<KeyMetadata>> {
     let alias = "ks_attest_ec_test_key";
     let gen_params = AuthSetBuilder::new()
         .no_auth_required()
@@ -844,7 +886,7 @@
         .digest(digest)
         .attestation_challenge(att_challenge.to_vec());
 
-    let attestation_key_metadata = sl.binder.generateKey(
+    let attestation_key_metadata = match sl.binder.generateKey(
         &KeyDescriptor {
             domain: Domain::APP,
             nspace: -1,
@@ -855,7 +897,23 @@
         &gen_params,
         0,
         b"entropy",
-    )?;
+    ) {
+        Ok(metadata) => metadata,
+        Err(e) => {
+            return if is_rkp_only_unknown_on_gsi(sl.level)
+                && e.service_specific_error() == ErrorCode::ATTESTATION_KEYS_NOT_PROVISIONED.0
+            {
+                // GSI replaces the values for remote_prov_prop properties (since they’re
+                // system_internal_prop properties), so on GSI the properties are not
+                // reliable indicators of whether StrongBox/TEE are RKP-only or not.
+                // Test can be skipped if it generates a key with attestation but doesn't provide
+                // an ATTEST_KEY and rkp-only property is undetermined.
+                Ok(None)
+            } else {
+                Err(e)
+            };
+        }
+    };
 
     // Should have public certificate.
     assert!(attestation_key_metadata.certificate.is_some());
@@ -868,7 +926,7 @@
         &gen_params,
         KeyOrigin::GENERATED,
     );
-    Ok(attestation_key_metadata)
+    Ok(Some(attestation_key_metadata))
 }
 
 /// Generate EC-P-256 key and attest it with given attestation key.
@@ -1432,8 +1490,8 @@
     sl: &SecLevel,
     gen_params: &AuthSetBuilder,
     alias: &str,
-) -> binder::Result<KeyMetadata> {
-    let key_metadata = sl.binder.generateKey(
+) -> binder::Result<Option<KeyMetadata>> {
+    let key_metadata = match sl.binder.generateKey(
         &KeyDescriptor {
             domain: Domain::APP,
             nspace: -1,
@@ -1444,7 +1502,23 @@
         gen_params,
         0,
         b"entropy",
-    )?;
+    ) {
+        Ok(metadata) => metadata,
+        Err(e) => {
+            return if is_rkp_only_unknown_on_gsi(sl.level)
+                && e.service_specific_error() == ErrorCode::ATTESTATION_KEYS_NOT_PROVISIONED.0
+            {
+                // GSI replaces the values for remote_prov_prop properties (since they’re
+                // system_internal_prop properties), so on GSI the properties are not
+                // reliable indicators of whether StrongBox/TEE are RKP-only or not.
+                // Test can be skipped if it generates a key with attestation but doesn't provide
+                // an ATTEST_KEY and rkp-only property is undetermined.
+                Ok(None)
+            } else {
+                Err(e)
+            };
+        }
+    };
 
     if gen_params.iter().any(|kp| {
         matches!(
@@ -1489,7 +1563,7 @@
     }
     check_key_authorizations(sl, &key_metadata.authorizations, gen_params, KeyOrigin::GENERATED);
 
-    Ok(key_metadata)
+    Ok(Some(key_metadata))
 }
 
 /// Generate a key using given authorizations and create an operation using the generated key.
@@ -1498,8 +1572,10 @@
     gen_params: &AuthSetBuilder,
     op_params: &AuthSetBuilder,
     alias: &str,
-) -> binder::Result<CreateOperationResponse> {
-    let key_metadata = generate_key(sl, gen_params, alias)?;
+) -> binder::Result<Option<CreateOperationResponse>> {
+    let Some(key_metadata) = generate_key(sl, gen_params, alias)? else {
+        return Ok(None);
+    };
 
-    sl.binder.createOperation(&key_metadata.key, op_params, false)
+    sl.binder.createOperation(&key_metadata.key, op_params, false).map(Some)
 }
diff --git a/keystore2/tests/keystore2_client_attest_key_tests.rs b/keystore2/tests/keystore2_client_attest_key_tests.rs
index 266f2f2..7597ba1 100644
--- a/keystore2/tests/keystore2_client_attest_key_tests.rs
+++ b/keystore2/tests/keystore2_client_attest_key_tests.rs
@@ -47,8 +47,12 @@
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
-        let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sl, algo, att_challenge).unwrap();
+        let Some(attestation_key_metadata) = key_generations::map_ks_error(
+            key_generations::generate_attestation_key(&sl, algo, att_challenge),
+        )
+        .unwrap() else {
+            return;
+        };
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -57,7 +61,7 @@
 
         // Create RSA signing key and use attestation key to sign it.
         let sign_key_alias = format!("ks_attest_rsa_signing_key_{}", getuid());
-        let sign_key_metadata = key_generations::generate_rsa_key(
+        let Some(sign_key_metadata) = key_generations::generate_rsa_key(
             &sl,
             Domain::APP,
             -1,
@@ -73,7 +77,9 @@
             },
             Some(&attestation_key_metadata.key),
         )
-        .unwrap();
+        .unwrap() else {
+            return;
+        };
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(sign_key_metadata.certificate.as_ref().unwrap());
@@ -94,8 +100,12 @@
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
-        let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sl, algo, att_challenge).unwrap();
+        let Some(attestation_key_metadata) = key_generations::map_ks_error(
+            key_generations::generate_attestation_key(&sl, algo, att_challenge),
+        )
+        .unwrap() else {
+            return;
+        };
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -104,7 +114,7 @@
 
         // Create RSA encrypt/decrypt key and use attestation key to sign it.
         let decrypt_key_alias = format!("ks_attest_rsa_encrypt_key_{}", getuid());
-        let decrypt_key_metadata = key_generations::generate_rsa_key(
+        let Some(decrypt_key_metadata) = key_generations::generate_rsa_key(
             &sl,
             Domain::APP,
             -1,
@@ -120,7 +130,9 @@
             },
             Some(&attestation_key_metadata.key),
         )
-        .unwrap();
+        .unwrap() else {
+            return;
+        };
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(decrypt_key_metadata.certificate.as_ref().unwrap());
@@ -142,8 +154,12 @@
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
-        let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sl, algo, att_challenge).unwrap();
+        let Some(attestation_key_metadata) = key_generations::map_ks_error(
+            key_generations::generate_attestation_key(&sl, algo, att_challenge),
+        )
+        .unwrap() else {
+            return;
+        };
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -187,13 +203,17 @@
     let att_challenge: &[u8] = b"foo";
 
     // Create EcCurve::CURVE_25519 attestation key.
-    let attestation_key_metadata = key_generations::generate_ec_attestation_key(
-        &sl,
-        att_challenge,
-        Digest::NONE,
-        EcCurve::CURVE_25519,
-    )
-    .unwrap();
+    let Some(attestation_key_metadata) =
+        key_generations::map_ks_error(key_generations::generate_ec_attestation_key(
+            &sl,
+            att_challenge,
+            Digest::NONE,
+            EcCurve::CURVE_25519,
+        ))
+        .unwrap()
+    else {
+        return;
+    };
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -202,7 +222,7 @@
 
     // Create RSA signing key and use attestation key to sign it.
     let sign_key_alias = format!("ksrsa_attested_sign_test_key_{}", getuid());
-    let sign_key_metadata = key_generations::generate_rsa_key(
+    let Some(sign_key_metadata) = key_generations::generate_rsa_key(
         &sl,
         Domain::APP,
         -1,
@@ -218,7 +238,9 @@
         },
         Some(&attestation_key_metadata.key),
     )
-    .unwrap();
+    .unwrap() else {
+        return;
+    };
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(sign_key_metadata.certificate.as_ref().unwrap());
@@ -327,8 +349,12 @@
     let att_challenge: &[u8] = b"foo";
 
     // Create RSA attestation key.
-    let attestation_key_metadata =
-        key_generations::generate_attestation_key(&sl, Algorithm::RSA, att_challenge).unwrap();
+    let Some(attestation_key_metadata) = key_generations::map_ks_error(
+        key_generations::generate_attestation_key(&sl, Algorithm::RSA, att_challenge),
+    )
+    .unwrap() else {
+        return;
+    };
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -448,8 +474,12 @@
     let att_challenge: &[u8] = b"foo";
 
     // Create attestation key.
-    let attestation_key_metadata =
-        key_generations::generate_attestation_key(&sl, Algorithm::RSA, att_challenge).unwrap();
+    let Some(attestation_key_metadata) = key_generations::map_ks_error(
+        key_generations::generate_attestation_key(&sl, Algorithm::RSA, att_challenge),
+    )
+    .unwrap() else {
+        return;
+    };
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -527,9 +557,12 @@
     let sl = SecLevel::tee();
 
     let att_challenge: &[u8] = b"foo";
-
-    let attest_key_metadata =
-        key_generations::generate_attestation_key(&sl, algorithm, att_challenge).unwrap();
+    let Some(attest_key_metadata) = key_generations::map_ks_error(
+        key_generations::generate_attestation_key(&sl, algorithm, att_challenge),
+    )
+    .unwrap() else {
+        return;
+    };
 
     let attest_id_params = get_attestation_ids(&sl.keystore2);
 
@@ -589,9 +622,12 @@
     let att_challenge: &[u8] = b"foo";
 
     // Create EC-Attestation key.
-    let attest_key_metadata =
-        key_generations::generate_ec_attestation_key(&sl, att_challenge, digest, EcCurve::P_256)
-            .unwrap();
+    let Some(attest_key_metadata) = key_generations::map_ks_error(
+        key_generations::generate_ec_attestation_key(&sl, att_challenge, digest, EcCurve::P_256),
+    )
+    .unwrap() else {
+        return;
+    };
 
     let attest_id_params = vec![
         (Tag::ATTESTATION_ID_BRAND, b"invalid-brand".to_vec()),
@@ -634,8 +670,12 @@
     let sl = SecLevel::tee();
 
     let att_challenge: &[u8] = b"foo";
-    let attest_key_metadata =
-        key_generations::generate_attestation_key(&sl, Algorithm::RSA, att_challenge).unwrap();
+    let Some(attest_key_metadata) = key_generations::map_ks_error(
+        key_generations::generate_attestation_key(&sl, Algorithm::RSA, att_challenge),
+    )
+    .unwrap() else {
+        return;
+    };
 
     let attest_id_params = get_attestation_ids(&sl.keystore2);
     for (attest_id, value) in attest_id_params {
diff --git a/keystore2/tests/keystore2_client_authorizations_tests.rs b/keystore2/tests/keystore2_client_authorizations_tests.rs
index 0981783..6732f5c 100644
--- a/keystore2/tests/keystore2_client_authorizations_tests.rs
+++ b/keystore2/tests/keystore2_client_authorizations_tests.rs
@@ -13,10 +13,10 @@
 // limitations under the License.
 
 use crate::keystore2_client_test_utils::{
-    app_attest_key_feature_exists, delete_app_key, perform_sample_asym_sign_verify_op,
-    perform_sample_hmac_sign_verify_op, perform_sample_sym_key_decrypt_op,
-    perform_sample_sym_key_encrypt_op, verify_certificate_serial_num,
-    verify_certificate_subject_name, SAMPLE_PLAIN_TEXT,
+    app_attest_key_feature_exists, delete_app_key,
+    perform_sample_asym_sign_verify_op, perform_sample_hmac_sign_verify_op,
+    perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op,
+    verify_certificate_serial_num, verify_certificate_subject_name, SAMPLE_PLAIN_TEXT,
 };
 use crate::{require_keymint, skip_test_if_no_app_attest_key_feature};
 use aconfig_android_hardware_biometrics_rust;
@@ -44,7 +44,7 @@
 use openssl::x509::X509NameBuilder;
 use std::time::SystemTime;
 
-fn gen_key_including_unique_id(sl: &SecLevel, alias: &str) -> Vec<u8> {
+fn gen_key_including_unique_id(sl: &SecLevel, alias: &str) -> Option<Vec<u8>> {
     let gen_params = authorizations::AuthSetBuilder::new()
         .no_auth_required()
         .algorithm(Algorithm::EC)
@@ -55,7 +55,9 @@
         .attestation_challenge(b"foo".to_vec())
         .include_unique_id();
 
-    let key_metadata = key_generations::generate_key(sl, &gen_params, alias).unwrap();
+    let key_metadata =
+        key_generations::map_ks_error(key_generations::generate_key(sl, &gen_params, alias))
+            .unwrap()?;
 
     let unique_id = get_value_from_attest_record(
         key_metadata.certificate.as_ref().unwrap(),
@@ -64,7 +66,7 @@
     )
     .expect("Unique id not found.");
     assert!(!unique_id.is_empty());
-    unique_id
+    Some(unique_id)
 }
 
 fn generate_key_and_perform_sign_verify_op_max_times(
@@ -72,8 +74,10 @@
     gen_params: &authorizations::AuthSetBuilder,
     alias: &str,
     max_usage_count: i32,
-) -> binder::Result<KeyMetadata> {
-    let key_metadata = key_generations::generate_key(sl, gen_params, alias)?;
+) -> binder::Result<Option<KeyMetadata>> {
+    let Some(key_metadata) = key_generations::generate_key(sl, gen_params, alias)? else {
+        return Ok(None);
+    };
 
     // Use above generated key `max_usage_count` times.
     for _ in 0..max_usage_count {
@@ -85,7 +89,7 @@
         );
     }
 
-    Ok(key_metadata)
+    Ok(Some(key_metadata))
 }
 
 /// Generate a key with `USAGE_COUNT_LIMIT` and verify the key characteristics. Test should be able
@@ -100,9 +104,12 @@
     check_attestation: bool,
 ) {
     // Generate a key and use the key for `max_usage_count` times.
-    let key_metadata =
+    let Some(key_metadata) =
         generate_key_and_perform_sign_verify_op_max_times(sl, gen_params, alias, max_usage_count)
-            .unwrap();
+            .unwrap()
+    else {
+        return;
+    };
 
     let auth = key_generations::get_key_auth(&key_metadata.authorizations, Tag::USAGE_COUNT_LIMIT)
         .unwrap();
@@ -158,7 +165,6 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .active_date_time(active_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
@@ -189,7 +195,6 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .active_date_time(future_active_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
@@ -220,7 +225,6 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .origination_expire_date_time(origination_expire_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
@@ -251,7 +255,6 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .origination_expire_date_time(origination_expire_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
@@ -286,7 +289,9 @@
         .usage_expire_date_time(usage_expire_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_hmac_verify_success";
-    let key_metadata = key_generations::generate_key(&sl, &gen_params, alias).unwrap();
+    let Some(key_metadata) = key_generations::generate_key(&sl, &gen_params, alias).unwrap() else {
+        return;
+    };
 
     perform_sample_hmac_sign_verify_op(&sl.binder, &key_metadata.key);
     delete_app_key(&sl.keystore2, alias).unwrap();
@@ -313,7 +318,9 @@
         .usage_expire_date_time(usage_expire_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_hamc_verify_fail";
-    let key_metadata = key_generations::generate_key(&sl, &gen_params, alias).unwrap();
+    let Some(key_metadata) = key_generations::generate_key(&sl, &gen_params, alias).unwrap() else {
+        return;
+    };
 
     let result = key_generations::map_ks_error(
         sl.binder.createOperation(
@@ -349,7 +356,9 @@
         .usage_expire_date_time(usage_expire_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
-    let key_metadata = key_generations::generate_key(&sl, &gen_params, alias).unwrap();
+    let Some(key_metadata) = key_generations::generate_key(&sl, &gen_params, alias).unwrap() else {
+        return;
+    };
     let cipher_text = perform_sample_sym_key_encrypt_op(
         &sl.binder,
         PaddingMode::PKCS7,
@@ -398,7 +407,9 @@
         .usage_expire_date_time(usage_expire_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
-    let key_metadata = key_generations::generate_key(&sl, &gen_params, alias).unwrap();
+    let Some(key_metadata) = key_generations::generate_key(&sl, &gen_params, alias).unwrap() else {
+        return;
+    };
     let cipher_text = perform_sample_sym_key_encrypt_op(
         &sl.binder,
         PaddingMode::PKCS7,
@@ -440,7 +451,6 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .early_boot_only();
 
     let alias = "ks_test_auth_tags_test";
@@ -471,14 +481,16 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .max_uses_per_boot(MAX_USES_COUNT);
 
     let alias = "ks_test_auth_tags_test";
     // Generate a key and use the key for `MAX_USES_COUNT` times.
-    let key_metadata =
+    let Some(key_metadata) =
         generate_key_and_perform_sign_verify_op_max_times(&sl, &gen_params, alias, MAX_USES_COUNT)
-            .unwrap();
+            .unwrap()
+    else {
+        return;
+    };
 
     // Try to use the key one more time.
     let result = key_generations::map_ks_error(sl.binder.createOperation(
@@ -583,7 +595,6 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .creation_date_time(creation_datetime.try_into().unwrap());
 
     let alias = "ks_test_auth_tags_test";
@@ -611,15 +622,15 @@
     let sl = SecLevel::tee();
 
     let alias_first = "ks_test_auth_tags_test_1";
-    let unique_id_first = gen_key_including_unique_id(&sl, alias_first);
+    if let Some(unique_id_first) = gen_key_including_unique_id(&sl, alias_first) {
+        let alias_second = "ks_test_auth_tags_test_2";
+        let unique_id_second = gen_key_including_unique_id(&sl, alias_second).unwrap();
 
-    let alias_second = "ks_test_auth_tags_test_2";
-    let unique_id_second = gen_key_including_unique_id(&sl, alias_second);
+        assert_eq!(unique_id_first, unique_id_second);
 
-    assert_eq!(unique_id_first, unique_id_second);
-
-    delete_app_key(&sl.keystore2, alias_first).unwrap();
-    delete_app_key(&sl.keystore2, alias_second).unwrap();
+        delete_app_key(&sl.keystore2, alias_first).unwrap();
+        delete_app_key(&sl.keystore2, alias_second).unwrap();
+    }
 }
 
 /// Generate a key with `APPLICATION_DATA`. Test should create an operation using the
@@ -759,8 +770,11 @@
         .ec_curve(EcCurve::P_256)
         .attestation_challenge(b"foo".to_vec());
     let attest_alias = "ks_test_auth_tags_attest_key";
-    let attest_key_metadata =
-        key_generations::generate_key(&sl, &attest_gen_params, attest_alias).unwrap();
+    let Some(attest_key_metadata) =
+        key_generations::generate_key(&sl, &attest_gen_params, attest_alias).unwrap()
+    else {
+        return;
+    };
 
     // Generate attested key.
     let alias = "ks_test_auth_tags_attested_key";
@@ -815,8 +829,11 @@
         .app_id(b"app-id".to_vec())
         .app_data(b"app-data".to_vec());
     let attest_alias = "ks_test_auth_tags_attest_key";
-    let attest_key_metadata =
-        key_generations::generate_key(&sl, &attest_gen_params, attest_alias).unwrap();
+    let Some(attest_key_metadata) =
+        key_generations::generate_key(&sl, &attest_gen_params, attest_alias).unwrap()
+    else {
+        return;
+    };
 
     // Generate new key using above generated attestation key without providing app-id and app-data.
     let alias = "ks_test_auth_tags_attested_key";
@@ -923,12 +940,13 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(b"foo".to_vec())
         .cert_subject_name(x509_name)
         .cert_serial(serial.to_vec());
 
     let alias = "ks_test_auth_tags_test";
-    let key_metadata = key_generations::generate_key(&sl, &gen_params, alias).unwrap();
+    let Some(key_metadata) = key_generations::generate_key(&sl, &gen_params, alias).unwrap() else {
+        return;
+    };
     verify_certificate_subject_name(
         key_metadata.certificate.as_ref().unwrap(),
         cert_subject.as_bytes(),
diff --git a/keystore2/tests/keystore2_client_device_unique_attestation_tests.rs b/keystore2/tests/keystore2_client_device_unique_attestation_tests.rs
index cc5d85c..91370c7 100644
--- a/keystore2/tests/keystore2_client_device_unique_attestation_tests.rs
+++ b/keystore2/tests/keystore2_client_device_unique_attestation_tests.rs
@@ -130,7 +130,7 @@
         let alias = "ks_test_device_unique_attest_id_test";
         match key_generations::map_ks_error(key_generations::generate_key(&sl, &gen_params, alias))
         {
-            Ok(key_metadata) => {
+            Ok(Some(key_metadata)) => {
                 let attest_id_value = get_value_from_attest_record(
                     key_metadata.certificate.as_ref().unwrap(),
                     attest_id,
@@ -140,6 +140,7 @@
                 assert_eq!(attest_id_value, value);
                 delete_app_key(&sl.keystore2, alias).unwrap();
             }
+            Ok(None) => {}
             Err(e) => {
                 assert_eq!(e, Error::Km(ErrorCode::CANNOT_ATTEST_IDS));
             }
@@ -196,7 +197,7 @@
 
     let alias = "ks_device_unique_ec_key_attest_test";
     match key_generations::map_ks_error(key_generations::generate_key(&sl, &gen_params, alias)) {
-        Ok(key_metadata) => {
+        Ok(Some(key_metadata)) => {
             perform_sample_asym_sign_verify_op(
                 &sl.binder,
                 &key_metadata,
@@ -205,6 +206,7 @@
             );
             delete_app_key(&sl.keystore2, alias).unwrap();
         }
+        Ok(None) => {}
         Err(e) => {
             assert_eq!(e, Error::Km(ErrorCode::CANNOT_ATTEST_IDS));
         }
@@ -235,7 +237,7 @@
 
     let alias = "ks_device_unique_rsa_key_attest_test";
     match key_generations::map_ks_error(key_generations::generate_key(&sl, &gen_params, alias)) {
-        Ok(key_metadata) => {
+        Ok(Some(key_metadata)) => {
             perform_sample_asym_sign_verify_op(
                 &sl.binder,
                 &key_metadata,
@@ -244,6 +246,7 @@
             );
             delete_app_key(&sl.keystore2, alias).unwrap();
         }
+        Ok(None) => {}
         Err(e) => {
             assert_eq!(e, Error::Km(ErrorCode::CANNOT_ATTEST_IDS));
         }
@@ -285,7 +288,7 @@
         let result =
             key_generations::map_ks_error(key_generations::generate_key(&sl, &gen_params, alias));
         assert!(result.is_err());
-        assert!(matches!(result.unwrap_err(), Error::Km(ErrorCode::CANNOT_ATTEST_IDS)));
+        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::CANNOT_ATTEST_IDS));
     }
 }
 
diff --git a/keystore2/tests/keystore2_client_rsa_key_tests.rs b/keystore2/tests/keystore2_client_rsa_key_tests.rs
index 1559008..cb8729f 100644
--- a/keystore2/tests/keystore2_client_rsa_key_tests.rs
+++ b/keystore2/tests/keystore2_client_rsa_key_tests.rs
@@ -78,9 +78,12 @@
     key_params: &key_generations::KeyParams,
     op_purpose: KeyPurpose,
     forced_op: ForcedOp,
-) -> binder::Result<CreateOperationResponse> {
-    let key_metadata =
-        key_generations::generate_rsa_key(sl, domain, nspace, alias, key_params, None)?;
+) -> binder::Result<Option<CreateOperationResponse>> {
+    let Some(key_metadata) =
+        key_generations::generate_rsa_key(sl, domain, nspace, alias, key_params, None)?
+    else {
+        return Ok(None);
+    };
 
     let mut op_params = authorizations::AuthSetBuilder::new().purpose(op_purpose);
 
@@ -97,7 +100,7 @@
         op_params = op_params.block_mode(value)
     }
 
-    sl.binder.createOperation(&key_metadata.key, &op_params, forced_op.0)
+    sl.binder.createOperation(&key_metadata.key, &op_params, forced_op.0).map(Some)
 }
 
 /// Generate RSA signing key with given parameters and perform signing operation.
@@ -109,7 +112,7 @@
 ) {
     let sl = SecLevel::tee();
 
-    let op_response = create_rsa_key_and_operation(
+    let Some(op_response) = create_rsa_key_and_operation(
         &sl,
         Domain::APP,
         -1,
@@ -126,7 +129,9 @@
         KeyPurpose::SIGN,
         ForcedOp(false),
     )
-    .expect("Failed to create an operation.");
+    .expect("Failed to create an operation.") else {
+        return;
+    };
 
     assert!(op_response.iOperation.is_some());
     assert_eq!(
diff --git a/keystore2/tests/keystore2_client_test_utils.rs b/keystore2/tests/keystore2_client_test_utils.rs
index d4ca2ae..f028a65 100644
--- a/keystore2/tests/keystore2_client_test_utils.rs
+++ b/keystore2/tests/keystore2_client_test_utils.rs
@@ -40,7 +40,6 @@
 use openssl::x509::X509;
 use packagemanager_aidl::aidl::android::content::pm::IPackageManagerNative::IPackageManagerNative;
 use serde::{Deserialize, Serialize};
-use std::path::PathBuf;
 use std::process::{Command, Output};
 
 /// This enum is used to communicate between parent and child processes.
@@ -100,10 +99,7 @@
     // (ro.product.*_for_attestation) reading logic would not be available for such devices
     // hence skipping this test for such scenario.
 
-    // This file is only present on GSI builds.
-    let gsi_marker = PathBuf::from("/system/system_ext/etc/init/init.gsi.rc");
-
-    get_vsr_api_level() < 34 && gsi_marker.as_path().is_file()
+    get_vsr_api_level() < 34 && key_generations::is_gsi()
 }
 
 #[macro_export]