Upgrade remotely provisioned keys if necessary.

This change applies a key upgrade path to the RKP keys when they are
used in key generation. Without this, RKP keys will fail after a device
receives an update due to version mismatch in KeyMint.

Test: atest keystore2_test && atest RemoteProvisionerUnitTests
Change-Id: I5dddc8fa1fe7fe9d7dd559b337089d607fcc735a
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index 639fe1e..be23ae5 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -45,7 +45,7 @@
 use std::collections::BTreeMap;
 use std::sync::atomic::{AtomicBool, Ordering};
 
-use crate::database::{CertificateChain, KeystoreDB, Uuid};
+use crate::database::{CertificateChain, KeyIdGuard, 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::metrics_store::log_rkp_error_stats;
@@ -73,6 +73,11 @@
         Self { security_level, km_uuid, is_hal_present: AtomicBool::new(true) }
     }
 
+    /// Returns the uuid for the KM instance attached to this RemProvState struct.
+    pub fn get_uuid(&self) -> Uuid {
+        self.km_uuid
+    }
+
     /// 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
@@ -121,7 +126,7 @@
         caller_uid: u32,
         params: &[KeyParameter],
         db: &mut KeystoreDB,
-    ) -> Result<Option<(AttestationKey, Certificate)>> {
+    ) -> Result<Option<(KeyIdGuard, AttestationKey, 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
@@ -142,7 +147,8 @@
                     Ok(None)
                 }
                 Ok(v) => match v {
-                    Some(cert_chain) => Ok(Some((
+                    Some((guard, cert_chain)) => Ok(Some((
+                        guard,
                         AttestationKey {
                             keyBlob: cert_chain.private_key.to_vec(),
                             attestKeyParams: vec![],
@@ -433,7 +439,7 @@
     caller_uid: u32,
     db: &mut KeystoreDB,
     km_uuid: &Uuid,
-) -> Result<Option<CertificateChain>> {
+) -> Result<Option<(KeyIdGuard, CertificateChain)>> {
     match domain {
         Domain::APP => {
             // Attempt to get an Attestation Key once. If it fails, then the app doesn't
@@ -458,7 +464,7 @@
                             "key and failed silently. Something is very wrong."
                         ))
                     },
-                    |cert_chain| Ok(Some(cert_chain)),
+                    |(guard, cert_chain)| Ok(Some((guard, cert_chain))),
                 )
         }
         _ => Ok(None),
@@ -471,12 +477,12 @@
     caller_uid: u32,
     db: &mut KeystoreDB,
     km_uuid: &Uuid,
-) -> Result<Option<CertificateChain>> {
-    let cert_chain = db
+) -> Result<Option<(KeyIdGuard, CertificateChain)>> {
+    let guard_and_chain = db
         .retrieve_attestation_key_and_cert_chain(domain, caller_uid as i64, 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)),
+    match guard_and_chain {
+        Some((guard, cert_chain)) => Ok(Some((guard, 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.
@@ -598,10 +604,11 @@
             .context(format!("In get_attestation_key: unknown irpc id '{}'", irpc_id))?;
         let (_, _, km_uuid) = get_keymint_device(sec_level)?;
 
-        let cert_chain = get_rem_prov_attest_key(Domain::APP, caller_uid as u32, db, &km_uuid)
-            .context("In get_attestation_key")?;
-        match cert_chain {
-            Some(chain) => Ok(RemotelyProvisionedKey {
+        let guard_and_cert_chain =
+            get_rem_prov_attest_key(Domain::APP, caller_uid as u32, db, &km_uuid)
+                .context("In get_attestation_key")?;
+        match guard_and_cert_chain {
+            Some((_, chain)) => Ok(RemotelyProvisionedKey {
                 keyBlob: chain.private_key.to_vec(),
                 encodedCertChain: chain.cert_chain,
             }),