Merge "Replaced size_t by int to handle the errors."
diff --git a/keystore2/src/rkpd_client.rs b/keystore2/src/rkpd_client.rs
index 2d1b23b..b426440 100644
--- a/keystore2/src/rkpd_client.rs
+++ b/keystore2/src/rkpd_client.rs
@@ -14,7 +14,7 @@
 
 //! Helper wrapper around RKPD interface.
 
-use crate::error::{map_binder_status_code, map_or_log_err, Error, ErrorCode};
+use crate::error::{map_binder_status_code, Error, ErrorCode};
 use crate::globals::get_remotely_provisioned_component_name;
 use crate::ks_err;
 use crate::utils::watchdog as wd;
@@ -23,7 +23,10 @@
     IGetKeyCallback::BnGetKeyCallback, IGetKeyCallback::IGetKeyCallback,
     IGetRegistrationCallback::BnGetRegistrationCallback,
     IGetRegistrationCallback::IGetRegistrationCallback, IRegistration::IRegistration,
-    IRemoteProvisioning::IRemoteProvisioning, RemotelyProvisionedKey::RemotelyProvisionedKey,
+    IRemoteProvisioning::IRemoteProvisioning,
+    IStoreUpgradedKeyCallback::BnStoreUpgradedKeyCallback,
+    IStoreUpgradedKeyCallback::IStoreUpgradedKeyCallback,
+    RemotelyProvisionedKey::RemotelyProvisionedKey,
 };
 use android_security_rkp_aidl::binder::{BinderFeatures, Interface, Strong};
 use anyhow::{Context, Result};
@@ -31,45 +34,36 @@
 use futures::executor::block_on;
 use std::sync::Mutex;
 
-type RegistrationSender = oneshot::Sender<Result<binder::Strong<dyn IRegistration>>>;
+/// Thread-safe channel for sending a value once and only once. If a value has
+/// already been send, subsequent calls to send will noop.
+struct SafeSender<T> {
+    inner: Mutex<Option<oneshot::Sender<T>>>,
+}
+
+impl<T> SafeSender<T> {
+    fn new(sender: oneshot::Sender<T>) -> Self {
+        Self { inner: Mutex::new(Some(sender)) }
+    }
+
+    fn send(&self, value: T) {
+        if let Some(inner) = self.inner.lock().unwrap().take() {
+            // assert instead of unwrap, because on failure send returns Err(value)
+            assert!(inner.send(value).is_ok(), "thread state is terminally broken");
+        }
+    }
+}
 
 struct GetRegistrationCallback {
-    registration_tx: Mutex<Option<RegistrationSender>>,
+    registration_tx: SafeSender<Result<binder::Strong<dyn IRegistration>>>,
 }
 
 impl GetRegistrationCallback {
     pub fn new_native_binder(
-        registration_tx: RegistrationSender,
-    ) -> Result<Strong<dyn IGetRegistrationCallback>> {
+        registration_tx: oneshot::Sender<Result<binder::Strong<dyn IRegistration>>>,
+    ) -> Strong<dyn IGetRegistrationCallback> {
         let result: Self =
-            GetRegistrationCallback { registration_tx: Mutex::new(Some(registration_tx)) };
-        Ok(BnGetRegistrationCallback::new_binder(result, BinderFeatures::default()))
-    }
-    fn on_success(&self, registration: &binder::Strong<dyn IRegistration>) -> Result<()> {
-        if let Some(tx) = self.registration_tx.lock().unwrap().take() {
-            tx.send(Ok(registration.clone())).unwrap();
-        }
-        Ok(())
-    }
-    fn on_cancel(&self) -> Result<()> {
-        if let Some(tx) = self.registration_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
-                    .context(ks_err!("GetRegistrationCallback cancelled.")),
-            )
-            .unwrap();
-        }
-        Ok(())
-    }
-    fn on_error(&self, error: &str) -> Result<()> {
-        if let Some(tx) = self.registration_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
-                    .context(ks_err!("GetRegistrationCallback failed: {:?}", error)),
-            )
-            .unwrap();
-        }
-        Ok(())
+            GetRegistrationCallback { registration_tx: SafeSender::new(registration_tx) };
+        BnGetRegistrationCallback::new_binder(result, BinderFeatures::default())
     }
 }
 
@@ -78,15 +72,26 @@
 impl IGetRegistrationCallback for GetRegistrationCallback {
     fn onSuccess(&self, registration: &Strong<dyn IRegistration>) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetRegistrationCallback::onSuccess", 500);
-        map_or_log_err(self.on_success(registration), Ok)
+        self.registration_tx.send(Ok(registration.clone()));
+        Ok(())
     }
     fn onCancel(&self) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetRegistrationCallback::onCancel", 500);
-        map_or_log_err(self.on_cancel(), Ok)
+        log::warn!("IGetRegistrationCallback cancelled");
+        self.registration_tx.send(
+            Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
+                .context(ks_err!("GetRegistrationCallback cancelled.")),
+        );
+        Ok(())
     }
     fn onError(&self, error: &str) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetRegistrationCallback::onError", 500);
-        map_or_log_err(self.on_error(error), Ok)
+        log::error!("IGetRegistrationCallback failed: '{error}'");
+        self.registration_tx.send(
+            Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
+                .context(ks_err!("GetRegistrationCallback failed: {:?}", error)),
+        );
+        Ok(())
     }
 }
 
@@ -102,8 +107,7 @@
         .context(ks_err!("Trying to get IRPC name."))?;
 
     let (tx, rx) = oneshot::channel();
-    let cb = GetRegistrationCallback::new_native_binder(tx)
-        .context(ks_err!("Trying to create a IGetRegistrationCallback."))?;
+    let cb = GetRegistrationCallback::new_native_binder(tx);
 
     remote_provisioning
         .getRegistration(&rpc_name, &cb)
@@ -112,46 +116,16 @@
     rx.await.unwrap()
 }
 
-type KeySender = oneshot::Sender<Result<RemotelyProvisionedKey>>;
-
 struct GetKeyCallback {
-    key_tx: Mutex<Option<KeySender>>,
+    key_tx: SafeSender<Result<RemotelyProvisionedKey>>,
 }
 
 impl GetKeyCallback {
-    pub fn new_native_binder(key_tx: KeySender) -> Result<Strong<dyn IGetKeyCallback>> {
-        let result: Self = GetKeyCallback { key_tx: Mutex::new(Some(key_tx)) };
-        Ok(BnGetKeyCallback::new_binder(result, BinderFeatures::default()))
-    }
-    fn on_success(&self, key: &RemotelyProvisionedKey) -> Result<()> {
-        if let Some(tx) = self.key_tx.lock().unwrap().take() {
-            tx.send(Ok(RemotelyProvisionedKey {
-                keyBlob: key.keyBlob.clone(),
-                encodedCertChain: key.encodedCertChain.clone(),
-            }))
-            .unwrap();
-        }
-        Ok(())
-    }
-    fn on_cancel(&self) -> Result<()> {
-        if let Some(tx) = self.key_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
-                    .context(ks_err!("GetKeyCallback cancelled.")),
-            )
-            .unwrap();
-        }
-        Ok(())
-    }
-    fn on_error(&self, error: &str) -> Result<()> {
-        if let Some(tx) = self.key_tx.lock().unwrap().take() {
-            tx.send(
-                Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
-                    .context(ks_err!("GetKeyCallback failed: {:?}", error)),
-            )
-            .unwrap();
-        }
-        Ok(())
+    pub fn new_native_binder(
+        key_tx: oneshot::Sender<Result<RemotelyProvisionedKey>>,
+    ) -> Strong<dyn IGetKeyCallback> {
+        let result: Self = GetKeyCallback { key_tx: SafeSender::new(key_tx) };
+        BnGetKeyCallback::new_binder(result, BinderFeatures::default())
     }
 }
 
@@ -160,15 +134,29 @@
 impl IGetKeyCallback for GetKeyCallback {
     fn onSuccess(&self, key: &RemotelyProvisionedKey) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetKeyCallback::onSuccess", 500);
-        map_or_log_err(self.on_success(key), Ok)
+        self.key_tx.send(Ok(RemotelyProvisionedKey {
+            keyBlob: key.keyBlob.clone(),
+            encodedCertChain: key.encodedCertChain.clone(),
+        }));
+        Ok(())
     }
     fn onCancel(&self) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetKeyCallback::onCancel", 500);
-        map_or_log_err(self.on_cancel(), Ok)
+        log::warn!("IGetKeyCallback cancelled");
+        self.key_tx.send(
+            Err(Error::Km(ErrorCode::OPERATION_CANCELLED))
+                .context(ks_err!("GetKeyCallback cancelled.")),
+        );
+        Ok(())
     }
     fn onError(&self, error: &str) -> binder::Result<()> {
         let _wp = wd::watch_millis("IGetKeyCallback::onError", 500);
-        map_or_log_err(self.on_error(error), Ok)
+        log::error!("IGetKeyCallback failed: {error}");
+        self.key_tx.send(
+            Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
+                .context(ks_err!("GetKeyCallback failed: {:?}", error)),
+        );
+        Ok(())
     }
 }
 
@@ -181,8 +169,7 @@
         .context(ks_err!("Trying to get to IRegistration service."))?;
 
     let (tx, rx) = oneshot::channel();
-    let cb = GetKeyCallback::new_native_binder(tx)
-        .context(ks_err!("Trying to create a IGetKeyCallback."))?;
+    let cb = GetKeyCallback::new_native_binder(tx);
 
     registration
         .getKey(caller_uid.try_into().unwrap(), &cb)
@@ -191,6 +178,39 @@
     rx.await.unwrap()
 }
 
+struct StoreUpgradedKeyCallback {
+    completer: SafeSender<Result<()>>,
+}
+
+impl StoreUpgradedKeyCallback {
+    pub fn new_native_binder(
+        completer: oneshot::Sender<Result<()>>,
+    ) -> Strong<dyn IStoreUpgradedKeyCallback> {
+        let result: Self = StoreUpgradedKeyCallback { completer: SafeSender::new(completer) };
+        BnStoreUpgradedKeyCallback::new_binder(result, BinderFeatures::default())
+    }
+}
+
+impl Interface for StoreUpgradedKeyCallback {}
+
+impl IStoreUpgradedKeyCallback for StoreUpgradedKeyCallback {
+    fn onSuccess(&self) -> binder::Result<()> {
+        let _wp = wd::watch_millis("IGetRegistrationCallback::onSuccess", 500);
+        self.completer.send(Ok(()));
+        Ok(())
+    }
+
+    fn onError(&self, error: &str) -> binder::Result<()> {
+        let _wp = wd::watch_millis("IGetRegistrationCallback::onError", 500);
+        log::error!("IGetRegistrationCallback failed: {error}");
+        self.completer.send(
+            Err(Error::Km(ErrorCode::UNKNOWN_ERROR))
+                .context(ks_err!("Failed to store upgraded key: {:?}", error)),
+        );
+        Ok(())
+    }
+}
+
 async fn store_rkpd_attestation_key_async(
     security_level: &SecurityLevel,
     key_blob: &[u8],
@@ -200,10 +220,14 @@
         .await
         .context(ks_err!("Trying to get to IRegistration service."))?;
 
+    let (tx, rx) = oneshot::channel();
+    let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+
     registration
-        .storeUpgradedKey(key_blob, upgraded_blob)
+        .storeUpgradedKeyAsync(key_blob, upgraded_blob, &cb)
         .context(ks_err!("Failed to store upgraded blob with RKPD."))?;
-    Ok(())
+
+    rx.await.unwrap()
 }
 
 /// Get attestation key from RKPD.
@@ -257,14 +281,19 @@
             todo!()
         }
 
-        fn storeUpgradedKey(&self, _: &[u8], _: &[u8]) -> binder::Result<()> {
+        fn storeUpgradedKeyAsync(
+            &self,
+            _: &[u8],
+            _: &[u8],
+            _: &Strong<dyn IStoreUpgradedKeyCallback>,
+        ) -> binder::Result<()> {
             todo!()
         }
     }
 
     fn get_mock_registration() -> Result<binder::Strong<dyn IRegistration>> {
         let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx).unwrap();
+        let cb = GetRegistrationCallback::new_native_binder(tx);
         let mock_registration = MockRegistration::new_native_binder();
 
         assert!(cb.onSuccess(&mock_registration).is_ok());
@@ -280,7 +309,7 @@
     #[test]
     fn test_get_registration_cb_cancel() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx).unwrap();
+        let cb = GetRegistrationCallback::new_native_binder(tx);
         assert!(cb.onCancel().is_ok());
 
         let result = block_on(rx).unwrap();
@@ -293,7 +322,7 @@
     #[test]
     fn test_get_registration_cb_error() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetRegistrationCallback::new_native_binder(tx).unwrap();
+        let cb = GetRegistrationCallback::new_native_binder(tx);
         assert!(cb.onError("error").is_ok());
 
         let result = block_on(rx).unwrap();
@@ -308,7 +337,7 @@
         let mock_key =
             RemotelyProvisionedKey { keyBlob: vec![1, 2, 3], encodedCertChain: vec![4, 5, 6] };
         let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx).unwrap();
+        let cb = GetKeyCallback::new_native_binder(tx);
         assert!(cb.onSuccess(&mock_key).is_ok());
 
         let key = block_on(rx).unwrap().unwrap();
@@ -318,7 +347,7 @@
     #[test]
     fn test_get_key_cb_cancel() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx).unwrap();
+        let cb = GetKeyCallback::new_native_binder(tx);
         assert!(cb.onCancel().is_ok());
 
         let result = block_on(rx).unwrap();
@@ -331,7 +360,7 @@
     #[test]
     fn test_get_key_cb_error() {
         let (tx, rx) = oneshot::channel();
-        let cb = GetKeyCallback::new_native_binder(tx).unwrap();
+        let cb = GetKeyCallback::new_native_binder(tx);
         assert!(cb.onError("error").is_ok());
 
         let result = block_on(rx).unwrap();
@@ -342,16 +371,38 @@
     }
 
     #[test]
-    #[ignore]
+    fn test_store_upgraded_cb_success() {
+        let (tx, rx) = oneshot::channel();
+        let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+        assert!(cb.onSuccess().is_ok());
+
+        block_on(rx).unwrap().unwrap();
+    }
+
+    #[test]
+    fn test_store_upgraded_key_cb_error() {
+        let (tx, rx) = oneshot::channel();
+        let cb = StoreUpgradedKeyCallback::new_native_binder(tx);
+        assert!(cb.onError("oh no! it failed").is_ok());
+
+        let result = block_on(rx).unwrap();
+        assert_eq!(
+            result.unwrap_err().downcast::<Error>().unwrap(),
+            Error::Km(ErrorCode::UNKNOWN_ERROR)
+        );
+    }
+
+    #[test]
     fn test_get_rkpd_attestation_key() {
+        binder::ProcessState::start_thread_pool();
         let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, 0).unwrap();
         assert!(!key.keyBlob.is_empty());
         assert!(!key.encodedCertChain.is_empty());
     }
 
     #[test]
-    #[ignore]
     fn test_get_rkpd_attestation_key_same_caller() {
+        binder::ProcessState::start_thread_pool();
         let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
         let caller_uid = 0;
 
@@ -364,8 +415,8 @@
     }
 
     #[test]
-    #[ignore]
     fn test_get_rkpd_attestation_key_different_caller() {
+        binder::ProcessState::start_thread_pool();
         let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
 
         // Different callers should be getting different keys.
@@ -377,8 +428,8 @@
     }
 
     #[test]
-    #[ignore]
     fn test_store_rkpd_attestation_key() {
+        binder::ProcessState::start_thread_pool();
         let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
         let key = get_rkpd_attestation_key(&SecurityLevel::TRUSTED_ENVIRONMENT, 0).unwrap();
 
diff --git a/keystore2/test_utils/authorizations.rs b/keystore2/test_utils/authorizations.rs
index 7dcee83..4608bc5 100644
--- a/keystore2/test_utils/authorizations.rs
+++ b/keystore2/test_utils/authorizations.rs
@@ -71,15 +71,6 @@
         self
     }
 
-    /// Add Attestation-ID.
-    pub fn attestation_app_id(mut self, b: Vec<u8>) -> Self {
-        self.0.push(KeyParameter {
-            tag: Tag::ATTESTATION_APPLICATION_ID,
-            value: KeyParameterValue::Blob(b),
-        });
-        self
-    }
-
     /// Add No_auth_required.
     pub fn no_auth_required(mut self) -> Self {
         self.0.push(KeyParameter {
diff --git a/keystore2/test_utils/key_generations.rs b/keystore2/test_utils/key_generations.rs
index 53597af..f9aaabb 100644
--- a/keystore2/test_utils/key_generations.rs
+++ b/keystore2/test_utils/key_generations.rs
@@ -58,8 +58,6 @@
     pub block_mode: Option<BlockMode>,
     /// Attestation challenge.
     pub att_challenge: Option<Vec<u8>>,
-    /// Attestation app id.
-    pub att_app_id: Option<Vec<u8>>,
 }
 
 /// DER-encoded PKCS#8 format RSA key. Generated using:
@@ -338,7 +336,6 @@
     nspace: i64,
     alias: Option<String>,
     att_challenge: Option<&[u8]>,
-    att_app_id: Option<&[u8]>,
 ) -> binder::Result<KeyMetadata> {
     let mut key_attest = false;
     let mut gen_params = AuthSetBuilder::new()
@@ -354,11 +351,6 @@
         gen_params = gen_params.clone().attestation_challenge(challenge.to_vec());
     }
 
-    if let Some(app_id) = att_app_id {
-        key_attest = true;
-        gen_params = gen_params.clone().attestation_app_id(app_id.to_vec());
-    }
-
     match sec_level.generateKey(
         &KeyDescriptor { domain, nspace, alias, blob: None },
         None,
@@ -453,9 +445,6 @@
     if let Some(value) = &key_params.att_challenge {
         gen_params = gen_params.attestation_challenge(value.to_vec())
     }
-    if let Some(value) = &key_params.att_app_id {
-        gen_params = gen_params.attestation_app_id(value.to_vec())
-    }
 
     let key_metadata = sec_level.generateKey(
         &KeyDescriptor { domain, nspace, alias, blob: None },
@@ -468,8 +457,7 @@
     // Must have a public key.
     assert!(key_metadata.certificate.is_some());
 
-    if attest_key.is_none() && key_params.att_challenge.is_some() && key_params.att_app_id.is_some()
-    {
+    if attest_key.is_none() && key_params.att_challenge.is_some() {
         // Should have an attestation record.
         assert!(key_metadata.certificateChain.is_some());
     } else {
@@ -578,7 +566,6 @@
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
     algorithm: Algorithm,
     att_challenge: &[u8],
-    att_app_id: &[u8],
 ) -> binder::Result<KeyMetadata> {
     assert!(algorithm == Algorithm::RSA || algorithm == Algorithm::EC);
 
@@ -597,7 +584,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: Some(att_challenge.to_vec()),
-                att_app_id: Some(att_app_id.to_vec()),
             },
             None,
         )
@@ -607,7 +593,6 @@
         let metadata = generate_ec_attestation_key(
             sec_level,
             att_challenge,
-            att_app_id,
             Digest::SHA_2_256,
             EcCurve::P_256,
         )
@@ -622,7 +607,6 @@
 pub fn generate_ec_attestation_key(
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
     att_challenge: &[u8],
-    att_app_id: &[u8],
     digest: Digest,
     ec_curve: EcCurve,
 ) -> binder::Result<KeyMetadata> {
@@ -633,8 +617,7 @@
         .purpose(KeyPurpose::ATTEST_KEY)
         .ec_curve(ec_curve)
         .digest(digest)
-        .attestation_challenge(att_challenge.to_vec())
-        .attestation_app_id(att_app_id.to_vec());
+        .attestation_challenge(att_challenge.to_vec());
 
     let attestation_key_metadata = sec_level.generateKey(
         &KeyDescriptor {
@@ -662,7 +645,6 @@
     sec_level: &binder::Strong<dyn IKeystoreSecurityLevel>,
     alias: Option<String>,
     att_challenge: &[u8],
-    att_app_id: &[u8],
     attest_key: &KeyDescriptor,
 ) -> binder::Result<KeyMetadata> {
     let ec_gen_params = AuthSetBuilder::new()
@@ -672,8 +654,7 @@
         .purpose(KeyPurpose::VERIFY)
         .digest(Digest::SHA_2_256)
         .ec_curve(EcCurve::P_256)
-        .attestation_challenge(att_challenge.to_vec())
-        .attestation_app_id(att_app_id.to_vec());
+        .attestation_challenge(att_challenge.to_vec());
 
     let ec_key_metadata = sec_level
         .generateKey(
diff --git a/keystore2/tests/keystore2_client_aes_key_tests.rs b/keystore2/tests/keystore2_client_aes_key_tests.rs
index 885cbf5..313f596 100644
--- a/keystore2/tests/keystore2_client_aes_key_tests.rs
+++ b/keystore2/tests/keystore2_client_aes_key_tests.rs
@@ -26,8 +26,7 @@
 };
 
 use crate::keystore2_client_test_utils::{
-    has_trusty_keymint, perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op,
-    SAMPLE_PLAIN_TEXT,
+    perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op, SAMPLE_PLAIN_TEXT,
 };
 
 /// Generate a AES key. Create encrypt and decrypt operations using the generated key.
@@ -393,11 +392,11 @@
     ));
     assert!(result.is_err());
 
-    if has_trusty_keymint() {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::MISSING_MAC_LENGTH));
-    } else {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::UNSUPPORTED_MAC_LENGTH));
-    }
+    let e = result.unwrap_err();
+    assert!(
+        e == Error::Km(ErrorCode::MISSING_MAC_LENGTH)
+            || e == Error::Km(ErrorCode::UNSUPPORTED_MAC_LENGTH)
+    );
 }
 
 /// Generate a AES-GCM key with `MIN_MAC_LENGTH`. Try to create an operation using this
diff --git a/keystore2/tests/keystore2_client_attest_key_tests.rs b/keystore2/tests/keystore2_client_attest_key_tests.rs
index b286b2f..2d44753 100644
--- a/keystore2/tests/keystore2_client_attest_key_tests.rs
+++ b/keystore2/tests/keystore2_client_attest_key_tests.rs
@@ -43,13 +43,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
         let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sec_level, algo, att_challenge, att_app_id)
-                .unwrap();
+            key_generations::generate_attestation_key(&sec_level, algo, att_challenge).unwrap();
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -71,7 +69,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: Some(att_challenge.to_vec()),
-                att_app_id: Some(att_app_id.to_vec()),
             },
             Some(&attestation_key_metadata.key),
         )
@@ -94,13 +91,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
         let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sec_level, algo, att_challenge, att_app_id)
-                .unwrap();
+            key_generations::generate_attestation_key(&sec_level, algo, att_challenge).unwrap();
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -122,7 +117,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: Some(att_challenge.to_vec()),
-                att_app_id: Some(att_app_id.to_vec()),
             },
             Some(&attestation_key_metadata.key),
         )
@@ -146,13 +140,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     for algo in [Algorithm::RSA, Algorithm::EC] {
         // Create attestation key.
         let attestation_key_metadata =
-            key_generations::generate_attestation_key(&sec_level, algo, att_challenge, att_app_id)
-                .unwrap();
+            key_generations::generate_attestation_key(&sec_level, algo, att_challenge).unwrap();
 
         let mut cert_chain: Vec<u8> = Vec::new();
         cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -167,7 +159,6 @@
             &sec_level,
             Some(ec_key_alias),
             att_challenge,
-            att_app_id,
             &attestation_key_metadata.key,
         )
         .unwrap();
@@ -193,13 +184,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     // Create EcCurve::CURVE_25519 attestation key.
     let attestation_key_metadata = key_generations::generate_ec_attestation_key(
         &sec_level,
         att_challenge,
-        att_app_id,
         Digest::NONE,
         EcCurve::CURVE_25519,
     )
@@ -225,7 +214,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: Some(att_challenge.to_vec()),
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&attestation_key_metadata.key),
     )
@@ -327,16 +315,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     // Create RSA attestation key.
-    let attestation_key_metadata = key_generations::generate_attestation_key(
-        &sec_level,
-        Algorithm::RSA,
-        att_challenge,
-        att_app_id,
-    )
-    .unwrap();
+    let attestation_key_metadata =
+        key_generations::generate_attestation_key(&sec_level, Algorithm::RSA, att_challenge)
+            .unwrap();
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -358,7 +341,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&attestation_key_metadata.key),
     ));
@@ -376,7 +358,6 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     let alias = format!("non_attest_key_{}", getuid());
     let non_attest_key_metadata = key_generations::generate_ec_p256_signing_key(
@@ -385,7 +366,6 @@
         -1,
         Some(alias),
         None,
-        None,
     )
     .unwrap();
 
@@ -404,7 +384,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: Some(att_challenge.to_vec()),
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&non_attest_key_metadata.key),
     ));
@@ -421,7 +400,6 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     let alias = "aes_attest_key";
     let sym_key_metadata = key_generations::generate_sym_key(
@@ -450,7 +428,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: Some(att_challenge.to_vec()),
-            att_app_id: Some(att_app_id.to_vec()),
         },
         Some(&sym_key_metadata.key),
     ));
@@ -468,16 +445,11 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
     let att_challenge: &[u8] = b"foo";
-    let att_app_id: &[u8] = b"bar";
 
     // Create attestation key.
-    let attestation_key_metadata = key_generations::generate_attestation_key(
-        &sec_level,
-        Algorithm::RSA,
-        att_challenge,
-        att_app_id,
-    )
-    .unwrap();
+    let attestation_key_metadata =
+        key_generations::generate_attestation_key(&sec_level, Algorithm::RSA, att_challenge)
+            .unwrap();
 
     let mut cert_chain: Vec<u8> = Vec::new();
     cert_chain.extend(attestation_key_metadata.certificate.as_ref().unwrap());
@@ -493,8 +465,7 @@
         .key_size(128)
         .padding_mode(PaddingMode::NONE)
         .block_mode(BlockMode::ECB)
-        .attestation_challenge(att_challenge.to_vec())
-        .attestation_app_id(att_app_id.to_vec());
+        .attestation_challenge(att_challenge.to_vec());
 
     let alias = format!("ks_test_sym_key_attest_{}", getuid());
     let aes_key_metadata = sec_level
diff --git a/keystore2/tests/keystore2_client_ec_key_tests.rs b/keystore2/tests/keystore2_client_ec_key_tests.rs
index 726d61c..c2034de 100644
--- a/keystore2/tests/keystore2_client_ec_key_tests.rs
+++ b/keystore2/tests/keystore2_client_ec_key_tests.rs
@@ -209,7 +209,6 @@
         key_generations::SELINUX_SHELL_NAMESPACE,
         None,
         None,
-        None,
     )
     .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_grant_key_tests.rs b/keystore2/tests/keystore2_client_grant_key_tests.rs
index 827a0de..7c75734 100644
--- a/keystore2/tests/keystore2_client_grant_key_tests.rs
+++ b/keystore2/tests/keystore2_client_grant_key_tests.rs
@@ -44,7 +44,6 @@
         key_generations::SELINUX_SHELL_NAMESPACE,
         Some(alias),
         None,
-        None,
     )
     .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_import_keys_tests.rs b/keystore2/tests/keystore2_client_import_keys_tests.rs
index c8f94b6..ecba402 100644
--- a/keystore2/tests/keystore2_client_import_keys_tests.rs
+++ b/keystore2/tests/keystore2_client_import_keys_tests.rs
@@ -35,7 +35,7 @@
 use crate::ffi_test_utils::{create_wrapped_key, create_wrapped_key_additional_auth_data};
 
 use crate::keystore2_client_test_utils::{
-    encrypt_secure_key, encrypt_transport_key, has_trusty_keymint,
+    encrypt_secure_key, encrypt_transport_key, has_default_keymint,
     perform_sample_asym_sign_verify_op, perform_sample_hmac_sign_verify_op,
     perform_sample_sym_key_decrypt_op, perform_sample_sym_key_encrypt_op, SAMPLE_PLAIN_TEXT,
 };
@@ -286,7 +286,7 @@
         key_generations::RSA_2048_KEY,
     ));
 
-    if has_trusty_keymint() {
+    if has_default_keymint() {
         assert!(result.is_err());
         assert_eq!(Error::Km(ErrorCode::INCOMPATIBLE_PURPOSE), result.unwrap_err());
     } else {
diff --git a/keystore2/tests/keystore2_client_list_entries_tests.rs b/keystore2/tests/keystore2_client_list_entries_tests.rs
index def9d94..62e3dd0 100644
--- a/keystore2/tests/keystore2_client_list_entries_tests.rs
+++ b/keystore2/tests/keystore2_client_list_entries_tests.rs
@@ -89,7 +89,6 @@
                 key_generations::SELINUX_SHELL_NAMESPACE,
                 Some(alias.to_string()),
                 None,
-                None,
             )
             .unwrap();
 
@@ -128,7 +127,6 @@
                     -1,
                     Some(alias.to_string()),
                     None,
-                    None,
                 )
                 .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_operation_tests.rs b/keystore2/tests/keystore2_client_operation_tests.rs
index e1102dd..9714900 100644
--- a/keystore2/tests/keystore2_client_operation_tests.rs
+++ b/keystore2/tests/keystore2_client_operation_tests.rs
@@ -307,7 +307,6 @@
         key_generations::SELINUX_SHELL_NAMESPACE,
         Some(alias),
         None,
-        None,
     )
     .unwrap();
 
diff --git a/keystore2/tests/keystore2_client_rsa_key_tests.rs b/keystore2/tests/keystore2_client_rsa_key_tests.rs
index 3139c2b..ad176a4 100644
--- a/keystore2/tests/keystore2_client_rsa_key_tests.rs
+++ b/keystore2/tests/keystore2_client_rsa_key_tests.rs
@@ -13,8 +13,8 @@
 // limitations under the License.
 
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    BlockMode::BlockMode, Digest::Digest, ErrorCode::ErrorCode, KeyPurpose::KeyPurpose,
-    PaddingMode::PaddingMode, SecurityLevel::SecurityLevel,
+    Digest::Digest, ErrorCode::ErrorCode, KeyPurpose::KeyPurpose, PaddingMode::PaddingMode,
+    SecurityLevel::SecurityLevel,
 };
 use android_system_keystore2::aidl::android::system::keystore2::{
     CreateOperationResponse::CreateOperationResponse, Domain::Domain,
@@ -25,9 +25,7 @@
     authorizations, get_keystore_service, key_generations, key_generations::Error,
 };
 
-use crate::keystore2_client_test_utils::{
-    delete_app_key, has_trusty_keymint, perform_sample_sign_operation, ForcedOp,
-};
+use crate::keystore2_client_test_utils::{delete_app_key, perform_sample_sign_operation, ForcedOp};
 
 /// This macro is used for creating signing key operation tests using digests and paddings
 /// for various key sizes.
@@ -59,7 +57,6 @@
                 stringify!($test_name),
                 $padding,
                 None,
-                None,
             );
         }
     };
@@ -73,7 +70,6 @@
                 stringify!($test_name),
                 $padding,
                 $mgf_digest,
-                Some(BlockMode::ECB),
             );
         }
     };
@@ -133,7 +129,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
@@ -170,18 +165,16 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
     ));
     assert!(result.is_err());
 
-    if has_trusty_keymint() {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::UNKNOWN_ERROR));
-    } else {
-        assert_eq!(result.unwrap_err(), Error::Km(ErrorCode::INCOMPATIBLE_DIGEST));
-    }
+    let e = result.unwrap_err();
+    assert!(
+        e == Error::Km(ErrorCode::UNKNOWN_ERROR) || e == Error::Km(ErrorCode::INCOMPATIBLE_DIGEST)
+    );
 
     delete_app_key(&keystore2, alias).unwrap();
 }
@@ -193,7 +186,6 @@
     alias: &str,
     padding: PaddingMode,
     mgf_digest: Option<Digest>,
-    block_mode: Option<BlockMode>,
 ) {
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
@@ -209,9 +201,8 @@
             padding: Some(padding),
             digest,
             mgf_digest,
-            block_mode,
+            block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1559,7 +1550,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
@@ -1592,7 +1582,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1624,7 +1613,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1654,7 +1642,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::SIGN,
         ForcedOp(false),
@@ -1684,7 +1671,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1714,7 +1700,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::AGREE_KEY,
         ForcedOp(false),
@@ -1747,7 +1732,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: None,
-                att_app_id: None,
             },
             KeyPurpose::DECRYPT,
             ForcedOp(false),
@@ -1781,7 +1765,6 @@
                 mgf_digest: None,
                 block_mode: None,
                 att_challenge: None,
-                att_app_id: None,
             },
             KeyPurpose::SIGN,
             ForcedOp(false),
@@ -1813,7 +1796,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::ENCRYPT,
         ForcedOp(false),
@@ -1845,7 +1827,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1876,7 +1857,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         KeyPurpose::DECRYPT,
         ForcedOp(false),
@@ -1906,7 +1886,6 @@
             mgf_digest: None,
             block_mode: None,
             att_challenge: None,
-            att_app_id: None,
         },
         None,
     ));
diff --git a/keystore2/tests/keystore2_client_test_utils.rs b/keystore2/tests/keystore2_client_test_utils.rs
index 59819df..56995e4 100644
--- a/keystore2/tests/keystore2_client_test_utils.rs
+++ b/keystore2/tests/keystore2_client_test_utils.rs
@@ -84,7 +84,8 @@
     };
 }
 
-pub fn has_trusty_keymint() -> bool {
+/// Indicate whether the default device is KeyMint (rather than Keymaster).
+pub fn has_default_keymint() -> bool {
     binder::is_declared("android.hardware.security.keymint.IKeyMintDevice/default")
         .expect("Could not check for declared keymint interface")
 }
@@ -102,10 +103,9 @@
     let keystore2 = get_keystore_service();
     let sec_level = keystore2.getSecurityLevel(SecurityLevel::TRUSTED_ENVIRONMENT).unwrap();
 
-    let key_metadata = key_generations::generate_ec_p256_signing_key(
-        &sec_level, domain, nspace, alias, None, None,
-    )
-    .unwrap();
+    let key_metadata =
+        key_generations::generate_ec_p256_signing_key(&sec_level, domain, nspace, alias, None)
+            .unwrap();
 
     sec_level.createOperation(
         &key_metadata.key,
diff --git a/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs b/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs
index 32ecd03..63122fe 100644
--- a/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs
+++ b/keystore2/tests/legacy_blobs/keystore2_legacy_blob_tests.rs
@@ -165,14 +165,12 @@
                 .unwrap();
             // Generate Key BLOB and prepare legacy keystore blob files.
             let att_challenge: &[u8] = b"foo";
-            let att_app_id: &[u8] = b"bar";
             let key_metadata = key_generations::generate_ec_p256_signing_key(
                 &sec_level,
                 Domain::BLOB,
                 SELINUX_SHELL_NAMESPACE,
                 None,
                 Some(att_challenge),
-                Some(att_app_id),
             )
             .expect("Failed to generate key blob");
 
@@ -424,14 +422,12 @@
                 .unwrap();
             // Generate Key BLOB and prepare legacy keystore blob files.
             let att_challenge: &[u8] = b"foo";
-            let att_app_id: &[u8] = b"bar";
             let key_metadata = key_generations::generate_ec_p256_signing_key(
                 &sec_level,
                 Domain::BLOB,
                 SELINUX_SHELL_NAMESPACE,
                 None,
                 Some(att_challenge),
-                Some(att_app_id),
             )
             .expect("Failed to generate key blob");