Snap for 8363780 from 6a50860489ff3085edfb034731c8b94249fbe7f6 to sc-qpr3-releaseandroid-12.1.0_r9android-12.1.0_r8android-12.1.0_r7android-12.1.0_r22android-12.1.0_r21android-12.1.0_r20android-12.1.0_r19android-12.1.0_r11android-12.1.0_r10android12-qpr3-s7-releaseandroid12-qpr3-s6-releaseandroid12-qpr3-s5-releaseandroid12-qpr3-s4-releaseandroid12-qpr3-s3-releaseandroid12-qpr3-s2-releaseandroid12-qpr3-s1-releaseandroid12-qpr3-release

Change-Id: I4d726fb2a82dc1c5db9ab999790d408d88ce094c

4 files changed, 243 insertions, 107 deletions

diff --git a/keystore2/src/attestation_key_utils.rs b/keystore2/src/attestation_key_utils.rs
index ca00539b..b2bc86c2 100644
--- a/keystore2/src/attestation_key_utils.rs
+++ b/keystore2/src/attestation_key_utils.rs
@@ -35,6 +35,7 @@ use keystore2_crypto::parse_subject_from_certificate;
 /// handled quite differently, thus the different representations.
 pub enum AttestationKeyInfo {
     RemoteProvisioned {
+        key_id_guard: KeyIdGuard,
         attestation_key: AttestationKey,
         attestation_certs: Certificate,
     },
@@ -66,8 +67,12 @@ pub fn get_attest_key_info(
                 "Trying to get remotely provisioned attestation key."
             ))
             .map(|result| {
-                result.map(|(attestation_key, attestation_certs)| {
-                    AttestationKeyInfo::RemoteProvisioned { attestation_key, attestation_certs }
+                result.map(|(key_id_guard, attestation_key, attestation_certs)| {
+                    AttestationKeyInfo::RemoteProvisioned {
+                        key_id_guard,
+                        attestation_key,
+                        attestation_certs,
+                    }
                 })
             }),
         None => Ok(None),
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 8acd775e..e1a704c5 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -323,6 +323,8 @@ pub static KEYSTORE_UUID: Uuid = Uuid([
     0x41, 0xe3, 0xb9, 0xce, 0x27, 0x58, 0x4e, 0x91, 0xbc, 0xfd, 0xa5, 0x5d, 0x91, 0x85, 0xab, 0x11,
 ]);
 
+static EXPIRATION_BUFFER_MS: i64 = 20000;
+
 /// Indicates how the sensitive part of this key blob is encrypted.
 #[derive(Debug, Eq, PartialEq, Ord, PartialOrd)]
 pub enum EncryptedBy {
@@ -1939,8 +1941,11 @@ impl KeystoreDB {
                 )?
                 .collect::<rusqlite::Result<Vec<(i64, DateTime)>>>()
                 .context("Failed to get date metadata")?;
+            // Calculate curr_time with a discount factor to avoid a key that's milliseconds away
+            // from expiration dodging this delete call.
             let curr_time = DateTime::from_millis_epoch(
-                SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_millis() as i64,
+                SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_millis() as i64
+                    + EXPIRATION_BUFFER_MS,
             );
             let mut num_deleted = 0;
             for id in key_ids_to_check.iter().filter(|kt| kt.1 < curr_time).map(|kt| kt.0) {
@@ -2049,6 +2054,41 @@ impl KeystoreDB {
         .context("In get_attestation_pool_status: ")
     }
 
+    fn query_kid_for_attestation_key_and_cert_chain(
+        &self,
+        tx: &Transaction,
+        domain: Domain,
+        namespace: i64,
+        km_uuid: &Uuid,
+    ) -> Result<Option<i64>> {
+        let mut stmt = tx.prepare(
+            "SELECT id
+             FROM persistent.keyentry
+             WHERE key_type = ?
+                   AND domain = ?
+                   AND namespace = ?
+                   AND state = ?
+                   AND km_uuid = ?;",
+        )?;
+        let rows = stmt
+            .query_map(
+                params![
+                    KeyType::Attestation,
+                    domain.0 as u32,
+                    namespace,
+                    KeyLifeCycle::Live,
+                    km_uuid
+                ],
+                |row| row.get(0),
+            )?
+            .collect::<rusqlite::Result<Vec<i64>>>()
+            .context("query failed.")?;
+        if rows.is_empty() {
+            return Ok(None);
+        }
+        Ok(Some(rows[0]))
+    }
+
     /// Fetches the private key and corresponding certificate chain assigned to a
     /// domain/namespace pair. Will either return nothing if the domain/namespace is
     /// not assigned, or one CertificateChain.
@@ -2057,7 +2097,7 @@ impl KeystoreDB {
         domain: Domain,
         namespace: i64,
         km_uuid: &Uuid,
-    ) -> Result<Option<CertificateChain>> {
+    ) -> Result<Option<(KeyIdGuard, CertificateChain)>> {
         let _wp = wd::watch_millis("KeystoreDB::retrieve_attestation_key_and_cert_chain", 500);
 
         match domain {
@@ -2067,69 +2107,70 @@ impl KeystoreDB {
                     .context(format!("Domain {:?} must be either App or SELinux.", domain));
             }
         }
-        self.with_transaction(TransactionBehavior::Deferred, |tx| {
-            let mut stmt = tx.prepare(
-                "SELECT subcomponent_type, blob
-             FROM persistent.blobentry
-             WHERE keyentryid IN
-                (SELECT id
-                 FROM persistent.keyentry
-                 WHERE key_type = ?
-                       AND domain = ?
-                       AND namespace = ?
-                       AND state = ?
-                       AND km_uuid = ?);",
-            )?;
-            let rows = stmt
-                .query_map(
-                    params![
-                        KeyType::Attestation,
-                        domain.0 as u32,
-                        namespace,
-                        KeyLifeCycle::Live,
-                        km_uuid
-                    ],
-                    |row| Ok((row.get(0)?, row.get(1)?)),
-                )?
-                .collect::<rusqlite::Result<Vec<(SubComponentType, Vec<u8>)>>>()
-                .context("query failed.")?;
-            if rows.is_empty() {
-                return Ok(None).no_gc();
-            } else if rows.len() != 3 {
-                return Err(KsError::sys()).context(format!(
-                    concat!(
-                        "Expected to get a single attestation",
-                        "key, cert, and cert chain for a total of 3 entries, but instead got {}."
-                    ),
-                    rows.len()
-                ));
-            }
-            let mut km_blob: Vec<u8> = Vec::new();
-            let mut cert_chain_blob: Vec<u8> = Vec::new();
-            let mut batch_cert_blob: Vec<u8> = Vec::new();
-            for row in rows {
-                let sub_type: SubComponentType = row.0;
-                match sub_type {
-                    SubComponentType::KEY_BLOB => {
-                        km_blob = row.1;
-                    }
-                    SubComponentType::CERT_CHAIN => {
-                        cert_chain_blob = row.1;
-                    }
-                    SubComponentType::CERT => {
-                        batch_cert_blob = row.1;
-                    }
-                    _ => Err(KsError::sys()).context("Unknown or incorrect subcomponent type.")?,
+
+        self.delete_expired_attestation_keys().context(
+            "In retrieve_attestation_key_and_cert_chain: failed to prune expired attestation keys",
+        )?;
+        let tx = self.conn.unchecked_transaction().context(
+            "In retrieve_attestation_key_and_cert_chain: Failed to initialize transaction.",
+        )?;
+        let key_id: i64;
+        match self.query_kid_for_attestation_key_and_cert_chain(&tx, domain, namespace, km_uuid)? {
+            None => return Ok(None),
+            Some(kid) => key_id = kid,
+        }
+        tx.commit()
+            .context("In retrieve_attestation_key_and_cert_chain: Failed to commit keyid query")?;
+        let key_id_guard = KEY_ID_LOCK.get(key_id);
+        let tx = self.conn.unchecked_transaction().context(
+            "In retrieve_attestation_key_and_cert_chain: Failed to initialize transaction.",
+        )?;
+        let mut stmt = tx.prepare(
+            "SELECT subcomponent_type, blob
+            FROM persistent.blobentry
+            WHERE keyentryid = ?;",
+        )?;
+        let rows = stmt
+            .query_map(params![key_id_guard.id()], |row| Ok((row.get(0)?, row.get(1)?)))?
+            .collect::<rusqlite::Result<Vec<(SubComponentType, Vec<u8>)>>>()
+            .context("query failed.")?;
+        if rows.is_empty() {
+            return Ok(None);
+        } else if rows.len() != 3 {
+            return Err(KsError::sys()).context(format!(
+                concat!(
+                    "Expected to get a single attestation",
+                    "key, cert, and cert chain for a total of 3 entries, but instead got {}."
+                ),
+                rows.len()
+            ));
+        }
+        let mut km_blob: Vec<u8> = Vec::new();
+        let mut cert_chain_blob: Vec<u8> = Vec::new();
+        let mut batch_cert_blob: Vec<u8> = Vec::new();
+        for row in rows {
+            let sub_type: SubComponentType = row.0;
+            match sub_type {
+                SubComponentType::KEY_BLOB => {
+                    km_blob = row.1;
+                }
+                SubComponentType::CERT_CHAIN => {
+                    cert_chain_blob = row.1;
                 }
+                SubComponentType::CERT => {
+                    batch_cert_blob = row.1;
+                }
+                _ => Err(KsError::sys()).context("Unknown or incorrect subcomponent type.")?,
             }
-            Ok(Some(CertificateChain {
+        }
+        Ok(Some((
+            key_id_guard,
+            CertificateChain {
                 private_key: ZVec::try_from(km_blob)?,
                 batch_cert: batch_cert_blob,
                 cert_chain: cert_chain_blob,
-            }))
-            .no_gc()
-        })
-        .context("In retrieve_attestation_key_and_cert_chain:")
+            },
+        )))
     }
 
     /// Updates the alias column of the given key id `newid` with the given alias,
@@ -3509,7 +3550,10 @@ mod tests {
     #[test]
     fn test_store_signed_attestation_certificate_chain() -> Result<()> {
         let mut db = new_test_db()?;
-        let expiration_date: i64 = 20;
+        let expiration_date: i64 =
+            SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_millis() as i64
+                + EXPIRATION_BUFFER_MS
+                + 10000;
         let namespace: i64 = 30;
         let base_byte: u8 = 1;
         let loaded_values =
@@ -3517,7 +3561,7 @@ mod tests {
         let chain =
             db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace, &KEYSTORE_UUID)?;
         assert_eq!(true, chain.is_some());
-        let cert_chain = chain.unwrap();
+        let (_, cert_chain) = chain.unwrap();
         assert_eq!(cert_chain.private_key.to_vec(), loaded_values.priv_key);
         assert_eq!(cert_chain.batch_cert, loaded_values.batch_cert);
         assert_eq!(cert_chain.cert_chain, loaded_values.cert_chain);
@@ -3586,7 +3630,9 @@ mod tests {
             TempDir::new("test_remove_expired_certs_").expect("Failed to create temp dir.");
         let mut db = new_test_db_with_gc(temp_dir.path(), |_, _| Ok(()))?;
         let expiration_date: i64 =
-            SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_millis() as i64 + 10000;
+            SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_millis() as i64
+                + EXPIRATION_BUFFER_MS
+                + 10000;
         let namespace: i64 = 30;
         let namespace_del1: i64 = 45;
         let namespace_del2: i64 = 60;
@@ -3597,7 +3643,7 @@ mod tests {
             0x01, /* base_byte */
         )?;
         load_attestation_key_pool(&mut db, 45, namespace_del1, 0x02)?;
-        load_attestation_key_pool(&mut db, 60, namespace_del2, 0x03)?;
+        load_attestation_key_pool(&mut db, expiration_date - 10001, namespace_del2, 0x03)?;
 
         let blob_entry_row_count: u32 = db
             .conn
@@ -3612,7 +3658,7 @@ mod tests {
         let mut cert_chain =
             db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace, &KEYSTORE_UUID)?;
         assert!(cert_chain.is_some());
-        let value = cert_chain.unwrap();
+        let (_, value) = cert_chain.unwrap();
         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());
@@ -3644,6 +3690,73 @@ mod tests {
         Ok(())
     }
 
+    fn compare_rem_prov_values(
+        expected: &RemoteProvValues,
+        actual: Option<(KeyIdGuard, CertificateChain)>,
+    ) {
+        assert!(actual.is_some());
+        let (_, value) = actual.unwrap();
+        assert_eq!(expected.batch_cert, value.batch_cert);
+        assert_eq!(expected.cert_chain, value.cert_chain);
+        assert_eq!(expected.priv_key, value.private_key.to_vec());
+    }
+
+    #[test]
+    fn test_dont_remove_valid_certs() -> Result<()> {
+        let temp_dir =
+            TempDir::new("test_remove_expired_certs_").expect("Failed to create temp dir.");
+        let mut db = new_test_db_with_gc(temp_dir.path(), |_, _| Ok(()))?;
+        let expiration_date: i64 =
+            SystemTime::now().duration_since(SystemTime::UNIX_EPOCH)?.as_millis() as i64
+                + EXPIRATION_BUFFER_MS
+                + 10000;
+        let namespace1: i64 = 30;
+        let namespace2: i64 = 45;
+        let namespace3: i64 = 60;
+        let entry_values1 = load_attestation_key_pool(
+            &mut db,
+            expiration_date,
+            namespace1,
+            0x01, /* base_byte */
+        )?;
+        let entry_values2 =
+            load_attestation_key_pool(&mut db, expiration_date + 40000, namespace2, 0x02)?;
+        let entry_values3 =
+            load_attestation_key_pool(&mut db, expiration_date - 9000, namespace3, 0x03)?;
+
+        let blob_entry_row_count: u32 = db
+            .conn
+            .query_row("SELECT COUNT(id) FROM persistent.blobentry;", NO_PARAMS, |row| row.get(0))
+            .expect("Failed to get blob entry row count.");
+        // We expect 9 rows here because there are three blobs per attestation key, i.e.,
+        // one key, one certificate chain, and one certificate.
+        assert_eq!(blob_entry_row_count, 9);
+
+        let mut cert_chain =
+            db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace1, &KEYSTORE_UUID)?;
+        compare_rem_prov_values(&entry_values1, cert_chain);
+
+        cert_chain =
+            db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace2, &KEYSTORE_UUID)?;
+        compare_rem_prov_values(&entry_values2, cert_chain);
+
+        cert_chain =
+            db.retrieve_attestation_key_and_cert_chain(Domain::APP, namespace3, &KEYSTORE_UUID)?;
+        compare_rem_prov_values(&entry_values3, cert_chain);
+
+        // Give the garbage collector half a second to catch up.
+        std::thread::sleep(Duration::from_millis(500));
+
+        let blob_entry_row_count: u32 = db
+            .conn
+            .query_row("SELECT COUNT(id) FROM persistent.blobentry;", NO_PARAMS, |row| row.get(0))
+            .expect("Failed to get blob entry row count.");
+        // There shound be 9 blob entries left, because all three keys are valid with
+        // three blobs each.
+        assert_eq!(blob_entry_row_count, 9);
+
+        Ok(())
+    }
     #[test]
     fn test_delete_all_attestation_keys() -> Result<()> {
         let mut db = new_test_db()?;
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index 212bf399..369be7c0 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -40,7 +40,7 @@ use anyhow::{Context, Result};
 use keystore2_crypto::parse_subject_from_certificate;
 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;
@@ -62,6 +62,11 @@ impl RemProvState {
         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
@@ -92,7 +97,7 @@ impl RemProvState {
         key: &KeyDescriptor,
         caller_uid: u32,
         db: &mut KeystoreDB,
-    ) -> Result<Option<CertificateChain>> {
+    ) -> Result<Option<(KeyIdGuard, CertificateChain)>> {
         match key.domain {
             Domain::APP => {
                 // Attempt to get an Attestation Key once. If it fails, then the app doesn't
@@ -117,7 +122,7 @@ impl RemProvState {
                                 "key and failed silently. Something is very wrong."
                             ))
                         },
-                        |cert_chain| Ok(Some(cert_chain)),
+                        |(guard, cert_chain)| Ok(Some((guard, cert_chain))),
                     )
             }
             _ => Ok(None),
@@ -130,12 +135,12 @@ impl RemProvState {
         key: &KeyDescriptor,
         caller_uid: u32,
         db: &mut KeystoreDB,
-    ) -> Result<Option<CertificateChain>> {
+    ) -> Result<Option<(KeyIdGuard, 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)),
+            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.
@@ -174,7 +179,7 @@ impl RemProvState {
         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
@@ -195,7 +200,8 @@ impl RemProvState {
                     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![],
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))