Keystore 2.0: Make key type an explict argument.

This refactor makes key type an explicit to relevant database function
to make it harder to implicitly use the wrong type.

Ignore-AOSP-First: No automerge path from AOSP.
Bug: 187862706
Bug: 189470584
Test: Regression tested with keystore2_test.
Change-Id: I9e1416743093f0a1ab86fd9351aed97f106ee819
Merged-In: I9e1416743093f0a1ab86fd9351aed97f106ee819
diff --git a/keystore2/src/boot_level_keys.rs b/keystore2/src/boot_level_keys.rs
index ddac1f8..0df3a45 100644
--- a/keystore2/src/boot_level_keys.rs
+++ b/keystore2/src/boot_level_keys.rs
@@ -14,7 +14,11 @@
 
 //! Offer keys based on the "boot level" for superencryption.
 
-use crate::{database::KeystoreDB, key_parameter::KeyParameterValue, raw_device::KeyMintDevice};
+use crate::{
+    database::{KeyType, KeystoreDB},
+    key_parameter::KeyParameterValue,
+    raw_device::KeyMintDevice,
+};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     Algorithm::Algorithm, Digest::Digest, KeyParameter::KeyParameter as KmKeyParameter,
     KeyParameterValue::KeyParameterValue as KmKeyParameterValue, KeyPurpose::KeyPurpose,
@@ -67,7 +71,7 @@
     }
 
     let (key_id_guard, key_entry) = km_dev
-        .lookup_or_generate_key(db, &key_desc, &params, |key_characteristics| {
+        .lookup_or_generate_key(db, &key_desc, KeyType::Client, &params, |key_characteristics| {
             key_characteristics.iter().any(|kc| {
                 if kc.securityLevel == km_dev.security_level() {
                     kc.authorizations.iter().any(|a| {
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 3b7ed38..a6c16e9 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -1481,12 +1481,13 @@
         &mut self,
         domain: &Domain,
         namespace: &i64,
+        key_type: KeyType,
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
         let _wp = wd::watch_millis("KeystoreDB::create_key_entry", 500);
 
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            Self::create_key_entry_internal(tx, domain, namespace, km_uuid).no_gc()
+            Self::create_key_entry_internal(tx, domain, namespace, key_type, km_uuid).no_gc()
         })
         .context("In create_key_entry.")
     }
@@ -1495,6 +1496,7 @@
         tx: &Transaction,
         domain: &Domain,
         namespace: &i64,
+        key_type: KeyType,
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
         match *domain {
@@ -1512,7 +1514,7 @@
                      VALUES(?, ?, ?, ?, NULL, ?, ?);",
                     params![
                         id,
-                        KeyType::Client,
+                        key_type,
                         domain.0 as u32,
                         *namespace,
                         KeyLifeCycle::Existing,
@@ -2103,6 +2105,7 @@
         alias: &str,
         domain: &Domain,
         namespace: &i64,
+        key_type: KeyType,
     ) -> Result<bool> {
         match *domain {
             Domain::APP | Domain::SELINUX => {}
@@ -2117,15 +2120,15 @@
             .execute(
                 "UPDATE persistent.keyentry
                  SET alias = NULL, domain = NULL, namespace = NULL, state = ?
-                 WHERE alias = ? AND domain = ? AND namespace = ?;",
-                params![KeyLifeCycle::Unreferenced, alias, domain.0 as u32, namespace],
+                 WHERE alias = ? AND domain = ? AND namespace = ? AND key_type = ?;",
+                params![KeyLifeCycle::Unreferenced, alias, domain.0 as u32, namespace, key_type],
             )
             .context("In rebind_alias: Failed to rebind existing entry.")?;
         let result = tx
             .execute(
                 "UPDATE persistent.keyentry
                     SET alias = ?, state = ?
-                    WHERE id = ? AND domain = ? AND namespace = ? AND state = ?;",
+                    WHERE id = ? AND domain = ? AND namespace = ? AND state = ? AND key_type = ?;",
                 params![
                     alias,
                     KeyLifeCycle::Live,
@@ -2133,6 +2136,7 @@
                     domain.0 as u32,
                     *namespace,
                     KeyLifeCycle::Existing,
+                    key_type,
                 ],
             )
             .context("In rebind_alias: Failed to set alias.")?;
@@ -2215,9 +2219,11 @@
     /// fields, and rebinds the given alias to the new key.
     /// The boolean returned is a hint for the garbage collector. If true, a key was replaced,
     /// is now unreferenced and needs to be collected.
+    #[allow(clippy::clippy::too_many_arguments)]
     pub fn store_new_key(
         &mut self,
         key: &KeyDescriptor,
+        key_type: KeyType,
         params: &[KeyParameter],
         blob_info: &(&[u8], &BlobMetaData),
         cert_info: &CertificateInfo,
@@ -2237,7 +2243,7 @@
             }
         };
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            let key_id = Self::create_key_entry_internal(tx, &domain, namespace, km_uuid)
+            let key_id = Self::create_key_entry_internal(tx, &domain, namespace, key_type, km_uuid)
                 .context("Trying to create new key entry.")?;
             let (blob, blob_metadata) = *blob_info;
             Self::set_blob_internal(
@@ -2265,7 +2271,7 @@
             Self::insert_keyparameter_internal(tx, &key_id, params)
                 .context("Trying to insert key parameters.")?;
             metadata.store_in_db(key_id.id(), tx).context("Trying to insert key metadata.")?;
-            let need_gc = Self::rebind_alias(tx, &key_id, &alias, &domain, namespace)
+            let need_gc = Self::rebind_alias(tx, &key_id, &alias, &domain, namespace, key_type)
                 .context("Trying to rebind alias.")?;
             Ok(key_id).do_gc(need_gc)
         })
@@ -2278,6 +2284,7 @@
     pub fn store_new_certificate(
         &mut self,
         key: &KeyDescriptor,
+        key_type: KeyType,
         cert: &[u8],
         km_uuid: &Uuid,
     ) -> Result<KeyIdGuard> {
@@ -2295,7 +2302,7 @@
             }
         };
         self.with_transaction(TransactionBehavior::Immediate, |tx| {
-            let key_id = Self::create_key_entry_internal(tx, &domain, namespace, km_uuid)
+            let key_id = Self::create_key_entry_internal(tx, &domain, namespace, key_type, km_uuid)
                 .context("Trying to create new key entry.")?;
 
             Self::set_blob_internal(
@@ -2314,7 +2321,7 @@
 
             metadata.store_in_db(key_id.id(), tx).context("Trying to insert key metadata.")?;
 
-            let need_gc = Self::rebind_alias(tx, &key_id, &alias, &domain, namespace)
+            let need_gc = Self::rebind_alias(tx, &key_id, &alias, &domain, namespace, key_type)
                 .context("Trying to rebind alias.")?;
             Ok(key_id).do_gc(need_gc)
         })
@@ -3256,7 +3263,7 @@
         namespace: i64,
     ) -> Result<bool> {
         db.with_transaction(TransactionBehavior::Immediate, |tx| {
-            KeystoreDB::rebind_alias(tx, newid, alias, &domain, &namespace).no_gc()
+            KeystoreDB::rebind_alias(tx, newid, alias, &domain, &namespace, KeyType::Client).no_gc()
         })
         .context("In rebind_alias.")
     }
@@ -3380,7 +3387,7 @@
         let temp_dir = TempDir::new("persistent_db_test")?;
         let mut db = KeystoreDB::new(temp_dir.path(), None)?;
 
-        db.create_key_entry(&Domain::APP, &100, &KEYSTORE_UUID)?;
+        db.create_key_entry(&Domain::APP, &100, KeyType::Client, &KEYSTORE_UUID)?;
         let entries = get_keyentry(&db)?;
         assert_eq!(entries.len(), 1);
 
@@ -3399,8 +3406,8 @@
 
         let mut db = new_test_db()?;
 
-        db.create_key_entry(&Domain::APP, &100, &KEYSTORE_UUID)?;
-        db.create_key_entry(&Domain::SELINUX, &101, &KEYSTORE_UUID)?;
+        db.create_key_entry(&Domain::APP, &100, KeyType::Client, &KEYSTORE_UUID)?;
+        db.create_key_entry(&Domain::SELINUX, &101, KeyType::Client, &KEYSTORE_UUID)?;
 
         let entries = get_keyentry(&db)?;
         assert_eq!(entries.len(), 2);
@@ -3409,15 +3416,15 @@
 
         // Test that we must pass in a valid Domain.
         check_result_is_error_containing_string(
-            db.create_key_entry(&Domain::GRANT, &102, &KEYSTORE_UUID),
+            db.create_key_entry(&Domain::GRANT, &102, KeyType::Client, &KEYSTORE_UUID),
             "Domain Domain(1) must be either App or SELinux.",
         );
         check_result_is_error_containing_string(
-            db.create_key_entry(&Domain::BLOB, &103, &KEYSTORE_UUID),
+            db.create_key_entry(&Domain::BLOB, &103, KeyType::Client, &KEYSTORE_UUID),
             "Domain Domain(3) must be either App or SELinux.",
         );
         check_result_is_error_containing_string(
-            db.create_key_entry(&Domain::KEY_ID, &104, &KEYSTORE_UUID),
+            db.create_key_entry(&Domain::KEY_ID, &104, KeyType::Client, &KEYSTORE_UUID),
             "Domain Domain(4) must be either App or SELinux.",
         );
 
@@ -3585,7 +3592,7 @@
         let mut db = new_test_db()?;
         load_attestation_key_pool(&mut db, 45 /* expiration */, 1 /* namespace */, 0x02)?;
         load_attestation_key_pool(&mut db, 80 /* expiration */, 2 /* namespace */, 0x03)?;
-        db.create_key_entry(&Domain::APP, &42, &KEYSTORE_UUID)?;
+        db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
         let result = db.delete_all_attestation_keys()?;
 
         // Give the garbage collector half a second to catch up.
@@ -3606,8 +3613,8 @@
         }
 
         let mut db = new_test_db()?;
-        db.create_key_entry(&Domain::APP, &42, &KEYSTORE_UUID)?;
-        db.create_key_entry(&Domain::APP, &42, &KEYSTORE_UUID)?;
+        db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
+        db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
         let entries = get_keyentry(&db)?;
         assert_eq!(entries.len(), 2);
         assert_eq!(
@@ -3943,6 +3950,7 @@
                 alias: Some(TEST_ALIAS.to_string()),
                 blob: None,
             },
+            KeyType::Client,
             TEST_CERT_BLOB,
             &KEYSTORE_UUID,
         )
@@ -5104,7 +5112,7 @@
         alias: &str,
         max_usage_count: Option<i32>,
     ) -> Result<KeyIdGuard> {
-        let key_id = db.create_key_entry(&domain, &namespace, &KEYSTORE_UUID)?;
+        let key_id = db.create_key_entry(&domain, &namespace, KeyType::Client, &KEYSTORE_UUID)?;
         let mut blob_metadata = BlobMetaData::new();
         blob_metadata.add(BlobMetaEntry::EncryptedBy(EncryptedBy::Password));
         blob_metadata.add(BlobMetaEntry::Salt(vec![1, 2, 3]));
@@ -5394,7 +5402,7 @@
         let mut db = new_test_db()?;
         let mut working_stats = get_storage_stats_map(&mut db);
 
-        let key_id = db.create_key_entry(&Domain::APP, &42, &KEYSTORE_UUID)?;
+        let key_id = db.create_key_entry(&Domain::APP, &42, KeyType::Client, &KEYSTORE_UUID)?;
         assert_storage_increased(
             &mut db,
             vec![
diff --git a/keystore2/src/legacy_migrator.rs b/keystore2/src/legacy_migrator.rs
index d5647cd..f92fd45 100644
--- a/keystore2/src/legacy_migrator.rs
+++ b/keystore2/src/legacy_migrator.rs
@@ -14,11 +14,11 @@
 
 //! This module acts as a bridge between the legacy key database and the keystore2 database.
 
-use crate::error::Error;
 use crate::key_parameter::KeyParameterValue;
 use crate::legacy_blob::BlobValue;
 use crate::utils::{uid_to_android_user, watchdog as wd};
 use crate::{async_task::AsyncTask, legacy_blob::LegacyBlobLoader};
+use crate::{database::KeyType, error::Error};
 use crate::{
     database::{
         BlobMetaData, BlobMetaEntry, CertificateInfo, DateTime, EncryptedBy, KeyMetaData,
@@ -523,6 +523,7 @@
                 self.db
                     .store_new_key(
                         &key,
+                        KeyType::Client,
                         &params,
                         &(&blob, &blob_metadata),
                         &CertificateInfo::new(user_cert, ca_cert),
@@ -535,7 +536,7 @@
             None => {
                 if let Some(ca_cert) = ca_cert {
                     self.db
-                        .store_new_certificate(&key, &ca_cert, &KEYSTORE_UUID)
+                        .store_new_certificate(&key, KeyType::Client, &ca_cert, &KEYSTORE_UUID)
                         .context("In check_and_migrate: Failed to insert new certificate.")?;
                     Ok(())
                 } else {
diff --git a/keystore2/src/raw_device.rs b/keystore2/src/raw_device.rs
index 4e86d39..cd54915 100644
--- a/keystore2/src/raw_device.rs
+++ b/keystore2/src/raw_device.rs
@@ -101,6 +101,7 @@
         &self,
         db: &mut KeystoreDB,
         key_desc: &KeyDescriptor,
+        key_type: KeyType,
         creator: F,
     ) -> Result<()>
     where
@@ -120,6 +121,7 @@
 
         db.store_new_key(
             &key_desc,
+            key_type,
             &key_parameters,
             &(&creation_result.keyBlob, &blob_metadata),
             &CertificateInfo::new(None, None),
@@ -144,11 +146,10 @@
     fn lookup_from_desc(
         db: &mut KeystoreDB,
         key_desc: &KeyDescriptor,
+        key_type: KeyType,
     ) -> Result<(KeyIdGuard, KeyEntry)> {
-        db.load_key_entry(&key_desc, KeyType::Client, KeyEntryLoadBits::KM, AID_KEYSTORE, |_, _| {
-            Ok(())
-        })
-        .context("In lookup_from_desc: load_key_entry failed")
+        db.load_key_entry(&key_desc, key_type, KeyEntryLoadBits::KM, AID_KEYSTORE, |_, _| Ok(()))
+            .context("In lookup_from_desc: load_key_entry failed.")
     }
 
     /// Look up the key in the database, and return None if it is absent.
@@ -170,6 +171,7 @@
         &self,
         db: &mut KeystoreDB,
         key_desc: &KeyDescriptor,
+        key_type: KeyType,
         params: &[KeyParameter],
         validate_characteristics: F,
     ) -> Result<(KeyIdGuard, KeyBlob)>
@@ -181,7 +183,7 @@
         // - because the caller needs to hold a lock in any case
         // - because it avoids holding database locks during slow
         //   KeyMint operations
-        let lookup = Self::not_found_is_none(Self::lookup_from_desc(db, key_desc))
+        let lookup = Self::not_found_is_none(Self::lookup_from_desc(db, key_desc, key_type))
             .context("In lookup_or_generate_key: first lookup failed")?;
 
         if let Some((key_id_guard, mut key_entry)) = lookup {
@@ -226,9 +228,11 @@
             };
         }
 
-        self.create_and_store_key(db, &key_desc, |km_dev| km_dev.generateKey(&params, None))
-            .context("In lookup_or_generate_key: generate_and_store_key failed")?;
-        Self::lookup_from_desc(db, key_desc)
+        self.create_and_store_key(db, &key_desc, key_type, |km_dev| {
+            km_dev.generateKey(&params, None)
+        })
+        .context("In lookup_or_generate_key: generate_and_store_key failed")?;
+        Self::lookup_from_desc(db, key_desc, key_type)
             .and_then(|(key_id_guard, mut key_entry)| {
                 Ok((
                     key_id_guard,
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 00b26e4..f78d98b 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -176,6 +176,7 @@
                     let key_id = db
                         .store_new_key(
                             &key,
+                            KeyType::Client,
                             &key_parameters,
                             &(&key_blob, &blob_metadata),
                             &cert_info,
diff --git a/keystore2/src/service.rs b/keystore2/src/service.rs
index 1f61729..d65743d 100644
--- a/keystore2/src/service.rs
+++ b/keystore2/src/service.rs
@@ -241,8 +241,13 @@
             check_key_permission(KeyPerm::rebind(), &key, &None)
                 .context("Caller does not have permission to insert this certificate.")?;
 
-            db.store_new_certificate(&key, certificate_chain.unwrap(), &KEYSTORE_UUID)
-                .context("Failed to insert new certificate.")?;
+            db.store_new_certificate(
+                &key,
+                KeyType::Client,
+                certificate_chain.unwrap(),
+                &KEYSTORE_UUID,
+            )
+            .context("Failed to insert new certificate.")?;
             Ok(())
         })
         .context("In update_subcomponent.")
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index 7449f20..9fb267a 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -956,13 +956,23 @@
                     }
                     let key_params: Vec<KmKeyParameter> =
                         key_params.into_iter().map(|x| x.into()).collect();
-                    km_dev.create_and_store_key(db, &key_desc, |dev| {
-                        let _wp = wd::watch_millis(
-                            "In lock_screen_lock_bound_key: calling importKey.",
-                            500,
-                        );
-                        dev.importKey(key_params.as_slice(), KeyFormat::RAW, &encrypting_key, None)
-                    })?;
+                    km_dev.create_and_store_key(
+                        db,
+                        &key_desc,
+                        KeyType::Client, /* TODO Should be Super b/189470584 */
+                        |dev| {
+                            let _wp = wd::watch_millis(
+                                "In lock_screen_lock_bound_key: calling importKey.",
+                                500,
+                            );
+                            dev.importKey(
+                                key_params.as_slice(),
+                                KeyFormat::RAW,
+                                &encrypting_key,
+                                None,
+                            )
+                        },
+                    )?;
                     entry.biometric_unlock = Some(BiometricUnlock {
                         sids: unlocking_sids.into(),
                         key_desc,