Merge "Ability for KS2 to select remotely provisioned key"
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 40860be..57ca7aa 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -581,11 +581,14 @@
 
 /// This type represents a certificate chain with a private key corresponding to the leaf
 /// certificate. TODO(jbires): This will be used in a follow-on CL, for now it's used in the tests.
-#[allow(dead_code)]
 pub struct CertificateChain {
-    private_key: ZVec,
-    batch_cert: ZVec,
-    cert_chain: ZVec,
+    /// A KM key blob
+    pub private_key: ZVec,
+    /// A batch cert for private_key
+    pub batch_cert: Vec<u8>,
+    /// A full certificate chain from root signing authority to private_key, including batch_cert
+    /// for convenience.
+    pub cert_chain: Vec<u8>,
 }
 
 /// This type represents a Keystore 2.0 key entry.
@@ -1914,8 +1917,8 @@
             }
             Ok(Some(CertificateChain {
                 private_key: ZVec::try_from(km_blob)?,
-                batch_cert: ZVec::try_from(batch_cert_blob)?,
-                cert_chain: ZVec::try_from(cert_chain_blob)?,
+                batch_cert: batch_cert_blob,
+                cert_chain: cert_chain_blob,
             }))
             .no_gc()
         })
@@ -3212,8 +3215,8 @@
         assert_eq!(true, chain.is_some());
         let cert_chain = chain.unwrap();
         assert_eq!(cert_chain.private_key.to_vec(), loaded_values.priv_key);
-        assert_eq!(cert_chain.batch_cert.to_vec(), loaded_values.batch_cert);
-        assert_eq!(cert_chain.cert_chain.to_vec(), loaded_values.cert_chain);
+        assert_eq!(cert_chain.batch_cert, loaded_values.batch_cert);
+        assert_eq!(cert_chain.cert_chain, loaded_values.cert_chain);
         Ok(())
     }
 
@@ -3306,8 +3309,8 @@
             db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace, &KEYSTORE_UUID)?;
         assert!(cert_chain.is_some());
         let value = cert_chain.unwrap();
-        assert_eq!(entry_values.batch_cert, value.batch_cert.to_vec());
-        assert_eq!(entry_values.cert_chain, value.cert_chain.to_vec());
+        assert_eq!(entry_values.batch_cert, value.batch_cert);
+        assert_eq!(entry_values.cert_chain, value.cert_chain);
         assert_eq!(entry_values.priv_key, value.private_key.to_vec());
 
         cert_chain = db.retrieve_attestation_key_and_cert_chain(
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index d606b6a..d6cc680 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -22,20 +22,184 @@
 use std::collections::HashMap;
 
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    IRemotelyProvisionedComponent::IRemotelyProvisionedComponent, MacedPublicKey::MacedPublicKey,
-    ProtectedData::ProtectedData, SecurityLevel::SecurityLevel,
+    Algorithm::Algorithm, AttestationKey::AttestationKey, Certificate::Certificate,
+    IRemotelyProvisionedComponent::IRemotelyProvisionedComponent, KeyParameter::KeyParameter,
+    KeyParameterValue::KeyParameterValue, MacedPublicKey::MacedPublicKey,
+    ProtectedData::ProtectedData, SecurityLevel::SecurityLevel, Tag::Tag,
 };
 use android_security_remoteprovisioning::aidl::android::security::remoteprovisioning::{
     AttestationPoolStatus::AttestationPoolStatus, IRemoteProvisioning::BnRemoteProvisioning,
     IRemoteProvisioning::IRemoteProvisioning,
 };
 use android_security_remoteprovisioning::binder::Strong;
+use android_system_keystore2::aidl::android::system::keystore2::{
+    Domain::Domain, KeyDescriptor::KeyDescriptor,
+};
 use anyhow::{Context, Result};
+use keystore2_crypto::parse_subject_from_certificate;
+use std::sync::atomic::{AtomicBool, Ordering};
 
-use crate::error::{self, map_or_log_err, map_rem_prov_error};
+use crate::database::{CertificateChain, KeystoreDB, Uuid};
+use crate::error::{self, map_or_log_err, map_rem_prov_error, Error};
 use crate::globals::{get_keymint_device, get_remotely_provisioned_component, DB};
 use crate::utils::Asp;
 
+/// Contains helper functions to check if remote provisioning is enabled on the system and, if so,
+/// to assign and retrieve attestation keys and certificate chains.
+#[derive(Default)]
+pub struct RemProvState {
+    security_level: SecurityLevel,
+    km_uuid: Uuid,
+    is_hal_present: AtomicBool,
+}
+
+impl RemProvState {
+    /// Creates a RemProvState struct.
+    pub fn new(security_level: SecurityLevel, km_uuid: Uuid) -> Self {
+        Self { security_level, km_uuid, is_hal_present: AtomicBool::new(true) }
+    }
+
+    /// Checks if remote provisioning is enabled and partially caches the result. On a hybrid system
+    /// remote provisioning can flip from being disabled to enabled depending on responses from the
+    /// server, so unfortunately caching the presence or absence of the HAL is not enough to fully
+    /// make decisions about the state of remote provisioning during runtime.
+    fn check_rem_prov_enabled(&self, db: &mut KeystoreDB) -> Result<bool> {
+        if !self.is_hal_present.load(Ordering::Relaxed)
+            || get_remotely_provisioned_component(&self.security_level).is_err()
+        {
+            self.is_hal_present.store(false, Ordering::Relaxed);
+            return Ok(false);
+        }
+        // To check if remote provisioning is enabled on a system that supports both remote
+        // provisioning and factory provisioned keys, we only need to check if there are any
+        // keys at all generated to indicate if the app has gotten the signal to begin filling
+        // the key pool from the server.
+        let pool_status = db
+            .get_attestation_pool_status(0 /* date */, &self.km_uuid)
+            .context("In check_rem_prov_enabled: failed to get attestation pool status.")?;
+        Ok(pool_status.total != 0)
+    }
+
+    /// Fetches a remote provisioning attestation key and certificate chain inside of the
+    /// returned `CertificateChain` struct if one exists for the given caller_uid. If one has not
+    /// been assigned, this function will assign it. If there are no signed attestation keys
+    /// available to be assigned, it will return the ResponseCode `OUT_OF_KEYS`
+    fn get_rem_prov_attest_key(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        db: &mut KeystoreDB,
+    ) -> Result<Option<CertificateChain>> {
+        match key.domain {
+            Domain::APP => {
+                // Attempt to get an Attestation Key once. If it fails, then the app doesn't
+                // have a valid chain assigned to it. The helper function will return None after
+                // attempting to assign a key. An error will be thrown if the pool is simply out
+                // of usable keys. Then another attempt to fetch the just-assigned key will be
+                // made. If this fails too, something is very wrong.
+                self.get_rem_prov_attest_key_helper(key, caller_uid, db)
+                    .context("In get_rem_prov_attest_key: Failed to get a key")?
+                    .map_or_else(
+                        || self.get_rem_prov_attest_key_helper(key, caller_uid, db),
+                        |v| Ok(Some(v)),
+                    )
+                    .context(concat!(
+                        "In get_rem_prov_attest_key: Failed to get a key after",
+                        "attempting to assign one."
+                    ))?
+                    .map_or_else(
+                        || {
+                            Err(Error::sys()).context(concat!(
+                                "In get_rem_prov_attest_key: Attempted to assign a ",
+                                "key and failed silently. Something is very wrong."
+                            ))
+                        },
+                        |cert_chain| Ok(Some(cert_chain)),
+                    )
+            }
+            _ => Ok(None),
+        }
+    }
+
+    /// Returns None if an AttestationKey fails to be assigned. Errors if no keys are available.
+    fn get_rem_prov_attest_key_helper(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        db: &mut KeystoreDB,
+    ) -> Result<Option<CertificateChain>> {
+        let cert_chain = db
+            .retrieve_attestation_key_and_cert_chain(key.domain, caller_uid as i64, &self.km_uuid)
+            .context("In get_rem_prov_attest_key_helper: Failed to retrieve a key + cert chain")?;
+        match cert_chain {
+            Some(cert_chain) => Ok(Some(cert_chain)),
+            // Either this app needs to be assigned a key, or the pool is empty. An error will
+            // be thrown if there is no key available to assign. This will indicate that the app
+            // should be nudged to provision more keys so keystore can retry.
+            None => {
+                db.assign_attestation_key(key.domain, caller_uid as i64, &self.km_uuid)
+                    .context("In get_rem_prov_attest_key_helper: Failed to assign a key")?;
+                Ok(None)
+            }
+        }
+    }
+
+    fn is_asymmetric_key(&self, params: &[KeyParameter]) -> bool {
+        params.iter().any(|kp| {
+            matches!(
+                kp,
+                KeyParameter {
+                    tag: Tag::ALGORITHM,
+                    value: KeyParameterValue::Algorithm(Algorithm::RSA)
+                } | KeyParameter {
+                    tag: Tag::ALGORITHM,
+                    value: KeyParameterValue::Algorithm(Algorithm::EC)
+                }
+            )
+        })
+    }
+
+    /// Checks to see (1) if the key in question should be attested to based on the algorithm and
+    /// (2) if remote provisioning is present and enabled on the system. If these conditions are
+    /// met, it makes an attempt to fetch the attestation key assigned to the `caller_uid`.
+    ///
+    /// It returns the ResponseCode `OUT_OF_KEYS` if there is not one key currently assigned to the
+    /// `caller_uid` and there are none available to assign.
+    pub fn get_remote_provisioning_key_and_certs(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        params: &[KeyParameter],
+        db: &mut KeystoreDB,
+    ) -> Result<(Option<AttestationKey>, Option<Certificate>)> {
+        if !self.is_asymmetric_key(params) || !self.check_rem_prov_enabled(db)? {
+            // There is no remote provisioning component for this security level on the
+            // device. Return None so the underlying KM instance knows to use its
+            // factory provisioned key instead. Alternatively, it's not an asymmetric key
+            // and therefore will not be attested.
+            Ok((None, None))
+        } else {
+            match self.get_rem_prov_attest_key(&key, caller_uid, db).context(concat!(
+                "In get_remote_provisioning_key_and_certs: Failed to get ",
+                "attestation key"
+            ))? {
+                Some(cert_chain) => Ok((
+                    Some(AttestationKey {
+                        keyBlob: cert_chain.private_key.to_vec(),
+                        attestKeyParams: vec![],
+                        issuerSubjectName: parse_subject_from_certificate(&cert_chain.batch_cert)
+                            .context(concat!(
+                            "In get_remote_provisioning_key_and_certs: Failed to ",
+                            "parse subject."
+                        ))?,
+                    }),
+                    Some(Certificate { encodedCertificate: cert_chain.cert_chain }),
+                )),
+                None => Ok((None, None)),
+            }
+        }
+    }
+}
 /// Implementation of the IRemoteProvisioning service.
 #[derive(Default)]
 pub struct RemoteProvisioningService {
@@ -126,7 +290,21 @@
             protected_data,
         ))
         .context("In generate_csr: Failed to generate csr")?;
-        Ok(mac)
+        let mut cose_mac_0 = Vec::<u8>::new();
+        // TODO(b/180392379): Replace this manual CBOR generation with the cbor-serde crate as well.
+        //                    This generates an array consisting of the mac and the public key Maps.
+        //                    Just generate the actual MacedPublicKeys structure when the crate is
+        //                    available.
+        cose_mac_0.push((0b100_00000 | (keys_to_sign.len() + 1)) as u8);
+        cose_mac_0.push(0b010_11000); //push mac
+        cose_mac_0.push(mac.len() as u8);
+        cose_mac_0.append(&mut mac);
+        for maced_public_key in keys_to_sign {
+            if maced_public_key.macedKey.len() > 83 + 8 {
+                cose_mac_0.extend_from_slice(&maced_public_key.macedKey[8..83 + 8]);
+            }
+        }
+        Ok(cose_mac_0)
     }
 
     /// Provisions a certificate chain for a key whose CSR was included in generate_csr. The
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index b187d3b..aee3b82 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -18,7 +18,7 @@
 
 use crate::globals::get_keymint_device;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    Algorithm::Algorithm, AttestationKey::AttestationKey,
+    Algorithm::Algorithm, AttestationKey::AttestationKey, Certificate::Certificate,
     HardwareAuthenticatorType::HardwareAuthenticatorType, IKeyMintDevice::IKeyMintDevice,
     KeyCreationResult::KeyCreationResult, KeyFormat::KeyFormat,
     KeyMintHardwareInfo::KeyMintHardwareInfo, KeyParameter::KeyParameter,
@@ -32,10 +32,11 @@
     KeyMetadata::KeyMetadata, KeyParameters::KeyParameters,
 };
 
-use crate::database::{CertificateInfo, KeyIdGuard};
+use crate::database::{CertificateInfo, KeyIdGuard, KeystoreDB};
 use crate::globals::{DB, ENFORCEMENTS, LEGACY_MIGRATOR, SUPER_KEY};
 use crate::key_parameter::KeyParameter as KsKeyParam;
 use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
+use crate::remote_provisioning::RemProvState;
 use crate::super_key::{KeyBlob, SuperKeyManager};
 use crate::utils::{check_key_permission, uid_to_android_user, Asp};
 use crate::{
@@ -63,6 +64,7 @@
     hw_info: KeyMintHardwareInfo,
     km_uuid: Uuid,
     operation_db: OperationDb,
+    rem_prov_state: RemProvState,
 }
 
 // Blob of 32 zeroes used as empty masking key.
@@ -88,6 +90,7 @@
             hw_info,
             km_uuid,
             operation_db: OperationDb::new(),
+            rem_prov_state: RemProvState::new(security_level, km_uuid),
         });
         result.as_binder().set_requesting_sid(true);
         Ok((result, km_uuid))
@@ -405,28 +408,56 @@
 
         // generate_key requires the rebind permission.
         check_key_permission(KeyPerm::rebind(), &key, &None).context("In generate_key.")?;
-
-        let attest_key = match attest_key_descriptor {
-            None => None,
-            Some(key) => Some(
-                self.get_attest_key(key, caller_uid)
-                    .context("In generate_key: Trying to load attest key")?,
-            ),
-        };
-
+        let (attest_key, cert_chain) = DB
+            .with::<_, Result<(Option<AttestationKey>, Option<Certificate>)>>(|db| {
+                self.get_attest_key_and_cert_chain(
+                    &key,
+                    caller_uid,
+                    attest_key_descriptor,
+                    params,
+                    &mut db.borrow_mut(),
+                )
+            })
+            .context("In generate_key: Trying to get an attestation key")?;
         let params = Self::add_certificate_parameters(caller_uid, params, &key)
             .context("In generate_key: Trying to get aaid.")?;
 
         let km_dev: Strong<dyn IKeyMintDevice> = self.keymint.get_interface()?;
         map_km_error(km_dev.addRngEntropy(entropy))
             .context("In generate_key: Trying to add entropy.")?;
-        let creation_result = map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
+        let mut creation_result = map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
             .context("In generate_key: While generating Key")?;
-
+        // The certificate chain ultimately gets flattened into a big DER encoded byte array,
+        // so providing that blob upfront in a single certificate entry should be fine.
+        if let Some(cert) = cert_chain {
+            creation_result.certificateChain.push(cert);
+        }
         let user_id = uid_to_android_user(caller_uid);
         self.store_new_key(key, creation_result, user_id, Some(flags)).context("In generate_key.")
     }
 
+    fn get_attest_key_and_cert_chain(
+        &self,
+        key: &KeyDescriptor,
+        caller_uid: u32,
+        attest_key_descriptor: Option<&KeyDescriptor>,
+        params: &[KeyParameter],
+        db: &mut KeystoreDB,
+    ) -> Result<(Option<AttestationKey>, Option<Certificate>)> {
+        match attest_key_descriptor {
+            None => self
+                .rem_prov_state
+                .get_remote_provisioning_key_and_certs(&key, caller_uid, params, db),
+            Some(attest_key) => Ok((
+                Some(
+                    self.get_attest_key(&attest_key, caller_uid)
+                        .context("In generate_key: Trying to load attest key")?,
+                ),
+                None,
+            )),
+        }
+    }
+
     fn get_attest_key(&self, key: &KeyDescriptor, caller_uid: u32) -> Result<AttestationKey> {
         let (km_blob, cert) = self
             .load_attest_key_blob_and_cert(&key, caller_uid)