Keystore 2.0: Add support for resetting legacy user keys.

Test: N/A
Bug: 159371296
Change-Id: I2e8adbf17ae953f17950591d72432ec3da7b4fee
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index f9710f9..db06bff 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -795,6 +795,7 @@
 pub struct PerBootDbKeepAlive(Connection);
 
 impl KeystoreDB {
+    const UNASSIGNED_KEY_ID: i64 = -1i64;
     const PERBOOT_DB_FILE_NAME: &'static str = &"file:perboot.sqlite?mode=memory&cache=shared";
 
     /// The alias of the user super key.
@@ -1460,6 +1461,24 @@
         .context("In set_blob.")
     }
 
+    /// Why would we insert a deleted blob? This weird function is for the purpose of legacy
+    /// key migration in the case where we bulk delete all the keys of an app or even a user.
+    /// We use this to insert key blobs into the database which can then be garbage collected
+    /// lazily by the key garbage collector.
+    pub fn set_deleted_blob(&mut self, blob: &[u8], blob_metadata: &BlobMetaData) -> Result<()> {
+        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+            Self::set_blob_internal(
+                &tx,
+                Self::UNASSIGNED_KEY_ID,
+                SubComponentType::KEY_BLOB,
+                Some(blob),
+                Some(blob_metadata),
+            )
+            .need_gc()
+        })
+        .context("In set_deleted_blob.")
+    }
+
     fn set_blob_internal(
         tx: &Transaction,
         key_id: i64,
@@ -2746,7 +2765,10 @@
     // otherwise return the id.
     fn insert_with_retry(inserter: impl Fn(i64) -> rusqlite::Result<usize>) -> Result<i64> {
         loop {
-            let newid: i64 = random();
+            let newid: i64 = match random() {
+                Self::UNASSIGNED_KEY_ID => continue, // UNASSIGNED_KEY_ID cannot be assigned.
+                i => i,
+            };
             match inserter(newid) {
                 // If the id already existed, try again.
                 Err(rusqlite::Error::SqliteFailure(
diff --git a/keystore2/src/legacy_blob.rs b/keystore2/src/legacy_blob.rs
index 1981022..b51f644 100644
--- a/keystore2/src/legacy_blob.rs
+++ b/keystore2/src/legacy_blob.rs
@@ -27,6 +27,7 @@
 };
 use anyhow::{Context, Result};
 use keystore2_crypto::{aes_gcm_decrypt, derive_key_from_password, ZVec};
+use std::collections::{HashMap, HashSet};
 use std::{convert::TryInto, fs::File, path::Path, path::PathBuf};
 use std::{
     fs,
@@ -724,6 +725,31 @@
         Ok(result)
     }
 
+    /// List all keystore entries belonging to the given user. Returns a map of UIDs
+    /// to sets of decoded aliases.
+    pub fn list_keystore_entries_for_user(
+        &self,
+        user_id: u32,
+    ) -> Result<HashMap<u32, HashSet<String>>> {
+        let user_entries = self
+            .list_user(user_id)
+            .context("In list_keystore_entries_for_user: Trying to list user.")?;
+
+        let result =
+            user_entries.into_iter().fold(HashMap::<u32, HashSet<String>>::new(), |mut acc, v| {
+                if let Some(sep_pos) = v.find('_') {
+                    if let Ok(uid) = v[0..sep_pos].parse::<u32>() {
+                        if let Some(alias) = Self::extract_alias(&v[sep_pos + 1..]) {
+                            let entry = acc.entry(uid).or_default();
+                            entry.insert(alias);
+                        }
+                    }
+                }
+                acc
+            });
+        Ok(result)
+    }
+
     /// List all keystore entries belonging to the given uid.
     pub fn list_keystore_entries_for_uid(&self, uid: u32) -> Result<Vec<String>> {
         let user_id = uid_to_android_user(uid);
diff --git a/keystore2/src/legacy_migrator.rs b/keystore2/src/legacy_migrator.rs
index 60c6bca..9ffe86c 100644
--- a/keystore2/src/legacy_migrator.rs
+++ b/keystore2/src/legacy_migrator.rs
@@ -19,6 +19,7 @@
     KeystoreDB, Uuid, KEYSTORE_UUID,
 };
 use crate::error::Error;
+use crate::key_parameter::KeyParameterValue;
 use crate::legacy_blob::BlobValue;
 use crate::utils::uid_to_android_user;
 use crate::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader};
@@ -66,6 +67,11 @@
     }
 }
 
+enum BulkDeleteRequest {
+    Uid(u32),
+    User(u32),
+}
+
 struct LegacyMigratorState {
     recently_migrated: HashSet<RecentMigration>,
     recently_migrated_super_key: HashSet<u32>,
@@ -356,6 +362,31 @@
         }
     }
 
+    /// Deletes all keys belonging to the given uid, migrating them into the database
+    /// for subsequent garbage collection if necessary.
+    pub fn bulk_delete_uid(&self, uid: u32, keep_non_super_encrypted_keys: bool) -> Result<()> {
+        let result = self.do_serialized(move |migrator_state| {
+            migrator_state.bulk_delete(BulkDeleteRequest::Uid(uid), keep_non_super_encrypted_keys)
+        });
+
+        result.unwrap_or(Ok(()))
+    }
+
+    /// Deletes all keys belonging to the given android user, migrating them into the database
+    /// for subsequent garbage collection if necessary.
+    pub fn bulk_delete_user(
+        &self,
+        user_id: u32,
+        keep_non_super_encrypted_keys: bool,
+    ) -> Result<()> {
+        let result = self.do_serialized(move |migrator_state| {
+            migrator_state
+                .bulk_delete(BulkDeleteRequest::User(user_id), keep_non_super_encrypted_keys)
+        });
+
+        result.unwrap_or(Ok(()))
+    }
+
     /// Queries the legacy database for the presence of a super key for the given user.
     pub fn has_super_key(&self, user_id: u32) -> Result<bool> {
         let result =
@@ -539,6 +570,111 @@
         }
     }
 
+    /// Key migrator request to be run by do_serialized.
+    /// See LegacyMigrator::bulk_delete_uid and LegacyMigrator::bulk_delete_user.
+    fn bulk_delete(
+        &mut self,
+        bulk_delete_request: BulkDeleteRequest,
+        keep_non_super_encrypted_keys: bool,
+    ) -> Result<()> {
+        let (aliases, user_id) = match bulk_delete_request {
+            BulkDeleteRequest::Uid(uid) => (
+                self.legacy_loader
+                    .list_keystore_entries_for_uid(uid)
+                    .context("In bulk_delete: Trying to get aliases for uid.")
+                    .map(|aliases| {
+                        let mut h = HashMap::<u32, HashSet<String>>::new();
+                        h.insert(uid, aliases.into_iter().collect());
+                        h
+                    })?,
+                uid_to_android_user(uid),
+            ),
+            BulkDeleteRequest::User(user_id) => (
+                self.legacy_loader
+                    .list_keystore_entries_for_user(user_id)
+                    .context("In bulk_delete: Trying to get aliases for user_id.")?,
+                user_id,
+            ),
+        };
+
+        let super_key_id = self
+            .db
+            .load_super_key(user_id)
+            .context("In bulk_delete: Failed to load super key")?
+            .map(|(_, entry)| entry.id());
+
+        for (uid, alias) in aliases
+            .into_iter()
+            .map(|(uid, aliases)| aliases.into_iter().map(move |alias| (uid, alias)))
+            .flatten()
+        {
+            let (km_blob_params, _, _) = self
+                .legacy_loader
+                .load_by_uid_alias(uid, &alias, None)
+                .context("In bulk_delete: Trying to load legacy blob.")?;
+
+            // Determine if the key needs special handling to be deleted.
+            let (need_gc, is_super_encrypted) = km_blob_params
+                .as_ref()
+                .map(|(blob, params)| {
+                    (
+                        params.iter().any(|kp| {
+                            KeyParameterValue::RollbackResistance == *kp.key_parameter_value()
+                        }),
+                        blob.is_encrypted(),
+                    )
+                })
+                .unwrap_or((false, false));
+
+            if keep_non_super_encrypted_keys && !is_super_encrypted {
+                continue;
+            }
+
+            if need_gc {
+                let mark_deleted = match km_blob_params
+                    .map(|(blob, _)| (blob.is_strongbox(), blob.take_value()))
+                {
+                    Some((is_strongbox, BlobValue::Encrypted { iv, tag, data })) => {
+                        let mut blob_metadata = BlobMetaData::new();
+                        if let (Ok(km_uuid), Some(super_key_id)) =
+                            (self.get_km_uuid(is_strongbox), super_key_id)
+                        {
+                            blob_metadata.add(BlobMetaEntry::KmUuid(km_uuid));
+                            blob_metadata.add(BlobMetaEntry::Iv(iv.to_vec()));
+                            blob_metadata.add(BlobMetaEntry::AeadTag(tag.to_vec()));
+                            blob_metadata
+                                .add(BlobMetaEntry::EncryptedBy(EncryptedBy::KeyId(super_key_id)));
+                            Some((LegacyBlob::Vec(data), blob_metadata))
+                        } else {
+                            // Oh well - we tried our best, but if we cannot determine which
+                            // KeyMint instance we have to send this blob to, we cannot
+                            // do more than delete the key from the file system.
+                            // And if we don't know which key wraps this key we cannot
+                            // unwrap it for KeyMint either.
+                            None
+                        }
+                    }
+                    Some((_, BlobValue::Decrypted(data))) => {
+                        Some((LegacyBlob::ZVec(data), BlobMetaData::new()))
+                    }
+                    _ => None,
+                };
+
+                if let Some((blob, blob_metadata)) = mark_deleted {
+                    self.db.set_deleted_blob(&blob, &blob_metadata).context(concat!(
+                        "In bulk_delete: Trying to insert deleted ",
+                        "blob into the database for garbage collection."
+                    ))?;
+                }
+            }
+
+            self.legacy_loader
+                .remove_keystore_entry(uid, &alias)
+                .context("In bulk_delete: Trying to remove migrated key.")?;
+        }
+        Ok(())
+    }
+
     fn has_super_key(&mut self, user_id: u32) -> Result<bool> {
         Ok(self.recently_migrated_super_key.contains(&user_id)
             || self.legacy_loader.has_super_key(user_id))
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index 156d20d..5ee685a 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -451,9 +451,10 @@
         match key_blob_before_upgrade {
             KeyBlob::Sensitive(_, super_key) => {
                 let (key, metadata) = Self::encrypt_with_super_key(key_after_upgrade, super_key)
-                .context(
-                "In reencrypt_on_upgrade_if_required. Failed to re-super-encrypt key on key upgrade.",
-                )?;
+                    .context(concat!(
+                        "In reencrypt_on_upgrade_if_required. ",
+                        "Failed to re-super-encrypt key on key upgrade."
+                    ))?;
                 Ok((KeyBlob::NonSensitive(key), Some(metadata)))
             }
             _ => Ok((KeyBlob::Ref(key_after_upgrade), None)),
@@ -520,8 +521,9 @@
                 if password.is_none() {
                     //transitioning to swiping, delete only the super key in database and cache, and
                     //super-encrypted keys in database (and in KM)
-                    Self::reset_user(db, skm, user_id, true)
-                        .context("In get_with_password_changed.")?;
+                    Self::reset_user(db, skm, legacy_migrator, user_id, true).context(
+                        "In get_with_password_changed: Trying to delete keys from the db.",
+                    )?;
                     //Lskf is now removed in Keystore
                     Ok(UserState::Uninitialized)
                 } else {
@@ -570,10 +572,14 @@
     pub fn reset_user(
         db: &mut KeystoreDB,
         skm: &SuperKeyManager,
+        legacy_migrator: &LegacyMigrator,
         user_id: u32,
         keep_non_super_encrypted_keys: bool,
     ) -> Result<()> {
         // mark keys created on behalf of the user as unreferenced.
+        legacy_migrator
+            .bulk_delete_user(user_id, keep_non_super_encrypted_keys)
+            .context("In reset_user: Trying to delete legacy keys.")?;
         db.unbind_keys_for_user(user_id as u32, keep_non_super_encrypted_keys)
             .context("In reset user. Error in unbinding keys.")?;
 
@@ -583,18 +589,18 @@
     }
 }
 
-/// This enum represents two states a Keymint Blob can be in, w.r.t super encryption.
-/// Sensitive variant represents a Keymint blob that is supposed to be super encrypted,
-/// but unwrapped during usage. Therefore, it has the super key along with the unwrapped key.
-/// Ref variant represents a Keymint blob that is not required to super encrypt or that is
-/// already super encrypted.
+/// This enum represents three states a KeyMint Blob can be in, w.r.t super encryption.
+/// `Sensitive` holds the non encrypted key and a reference to its super key.
+/// `NonSensitive` holds a non encrypted key that is never supposed to be encrypted.
+/// `Ref` holds a reference to a key blob when it does not need to be modified if its
+/// life time allows it.
 pub enum KeyBlob<'a> {
     Sensitive(ZVec, SuperKey),
     NonSensitive(Vec<u8>),
     Ref(&'a [u8]),
 }
 
-/// Deref returns a reference to the key material in both variants.
+/// Deref returns a reference to the key material in any variant.
 impl<'a> Deref for KeyBlob<'a> {
     type Target = [u8];
 
diff --git a/keystore2/src/user_manager.rs b/keystore2/src/user_manager.rs
index 8b7aad9..8e09144 100644
--- a/keystore2/src/user_manager.rs
+++ b/keystore2/src/user_manager.rs
@@ -59,12 +59,12 @@
             .context("In on_user_password_changed.")?
         {
             UserState::LskfLocked => {
-                //error - password can not be changed when the device is locked
+                // Error - password can not be changed when the device is locked
                 Err(KeystoreError::Rc(ResponseCode::LOCKED))
                     .context("In on_user_password_changed. Device is locked.")
             }
             _ => {
-                //LskfLocked is the only error case for password change
+                // LskfLocked is the only error case for password change
                 Ok(())
             }
         }
@@ -74,7 +74,16 @@
         // Check permission. Function should return if this failed. Therefore having '?' at the end
         // is very important.
         check_keystore_permission(KeystorePerm::change_user()).context("In add_or_remove_user.")?;
-        DB.with(|db| UserState::reset_user(&mut db.borrow_mut(), &SUPER_KEY, user_id as u32, false))
+        DB.with(|db| {
+            UserState::reset_user(
+                &mut db.borrow_mut(),
+                &SUPER_KEY,
+                &LEGACY_MIGRATOR,
+                user_id as u32,
+                false,
+            )
+        })
+        .context("In add_or_remove_user: Trying to delete keys from db.")
     }
 }