1 files changed, 45 insertions, 33 deletions

diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 70264d3c..76110b3f 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -318,7 +318,7 @@ impl KeystoreSecurityLevel {
                 &*km_dev,
                 key_id_guard,
                 &km_blob,
-                &blob_metadata,
+                blob_metadata.km_uuid().copied(),
                 &operation_parameters,
                 |blob| loop {
                     match map_km_error({
@@ -550,7 +550,7 @@ impl KeystoreSecurityLevel {
                     &*km_dev,
                     Some(key_id_guard),
                     &KeyBlob::Ref(&blob),
-                    &blob_metadata,
+                    blob_metadata.km_uuid().copied(),
                     &params,
                     |blob| {
                         let attest_key = Some(AttestationKey {
@@ -572,23 +572,40 @@ impl KeystoreSecurityLevel {
                 )
                 .context("In generate_key: Using user generated attestation key.")
                 .map(|(result, _)| result),
-            Some(AttestationKeyInfo::RemoteProvisioned { attestation_key, attestation_certs }) => {
-                map_km_error({
-                    let _wp = self.watch_millis(
-                        concat!(
-                            "In KeystoreSecurityLevel::generate_key (RemoteProvisioned): ",
-                            "calling generate_key.",
-                        ),
-                        5000, // Generate can take a little longer.
-                    );
-                    km_dev.generateKey(&params, Some(&attestation_key))
-                })
+            Some(AttestationKeyInfo::RemoteProvisioned {
+                key_id_guard,
+                attestation_key,
+                attestation_certs,
+            }) => self
+                .upgrade_keyblob_if_required_with(
+                    &*km_dev,
+                    Some(key_id_guard),
+                    &KeyBlob::Ref(&attestation_key.keyBlob),
+                    Some(self.rem_prov_state.get_uuid()),
+                    &[],
+                    |blob| {
+                        map_km_error({
+                            let _wp = self.watch_millis(
+                                concat!(
+                                    "In KeystoreSecurityLevel::generate_key (RemoteProvisioned): ",
+                                    "calling generate_key.",
+                                ),
+                                5000, // Generate can take a little longer.
+                            );
+                            let dynamic_attest_key = Some(AttestationKey {
+                                keyBlob: blob.to_vec(),
+                                attestKeyParams: vec![],
+                                issuerSubjectName: attestation_key.issuerSubjectName.clone(),
+                            });
+                            km_dev.generateKey(&params, dynamic_attest_key.as_ref())
+                        })
+                    },
+                )
                 .context("While generating Key with remote provisioned attestation key.")
-                .map(|mut creation_result| {
-                    creation_result.certificateChain.push(attestation_certs);
-                    creation_result
-                })
-            }
+                .map(|(mut result, _)| {
+                    result.certificateChain.push(attestation_certs);
+                    result
+                }),
             None => map_km_error({
                 let _wp = self.watch_millis(
                     concat!(
@@ -773,7 +790,7 @@ impl KeystoreSecurityLevel {
                 &*km_dev,
                 Some(wrapping_key_id_guard),
                 &wrapping_key_blob,
-                &wrapping_blob_metadata,
+                wrapping_blob_metadata.km_uuid().copied(),
                 &[],
                 |wrapping_blob| {
                     let _wp = self.watch_millis(
@@ -799,7 +816,7 @@ impl KeystoreSecurityLevel {
 
     fn store_upgraded_keyblob(
         key_id_guard: KeyIdGuard,
-        km_uuid: Option<&Uuid>,
+        km_uuid: Option<Uuid>,
         key_blob: &KeyBlob,
         upgraded_blob: &[u8],
     ) -> Result<()> {
@@ -809,7 +826,7 @@ impl KeystoreSecurityLevel {
 
         let mut new_blob_metadata = new_blob_metadata.unwrap_or_default();
         if let Some(uuid) = km_uuid {
-            new_blob_metadata.add(BlobMetaEntry::KmUuid(*uuid));
+            new_blob_metadata.add(BlobMetaEntry::KmUuid(uuid));
         }
 
         DB.with(|db| {
@@ -829,7 +846,7 @@ impl KeystoreSecurityLevel {
         km_dev: &dyn IKeyMintDevice,
         mut key_id_guard: Option<KeyIdGuard>,
         key_blob: &KeyBlob,
-        blob_metadata: &BlobMetaData,
+        km_uuid: Option<Uuid>,
         params: &[KeyParameter],
         f: F,
     ) -> Result<(T, Option<Vec<u8>>)>
@@ -845,13 +862,9 @@ impl KeystoreSecurityLevel {
                 if key_id_guard.is_some() {
                     // Unwrap cannot panic, because the is_some was true.
                     let kid = key_id_guard.take().unwrap();
-                    Self::store_upgraded_keyblob(
-                        kid,
-                        blob_metadata.km_uuid(),
-                        key_blob,
-                        upgraded_blob,
+                    Self::store_upgraded_keyblob(kid, km_uuid, key_blob, upgraded_blob).context(
+                        "In upgrade_keyblob_if_required_with: store_upgraded_keyblob failed",
                     )
-                    .context("In upgrade_keyblob_if_required_with: store_upgraded_keyblob failed")
                 } else {
                     Ok(())
                 }
@@ -864,11 +877,10 @@ impl KeystoreSecurityLevel {
         // upgrade was performed above and if one was given in the first place.
         if key_blob.force_reencrypt() {
             if let Some(kid) = key_id_guard {
-                Self::store_upgraded_keyblob(kid, blob_metadata.km_uuid(), key_blob, key_blob)
-                    .context(concat!(
-                        "In upgrade_keyblob_if_required_with: ",
-                        "store_upgraded_keyblob failed in forced reencrypt"
-                    ))?;
+                Self::store_upgraded_keyblob(kid, km_uuid, key_blob, key_blob).context(concat!(
+                    "In upgrade_keyblob_if_required_with: ",
+                    "store_upgraded_keyblob failed in forced reencrypt"
+                ))?;
             }
         }
         Ok((v, upgraded_blob))