Keystore 2.0: Use preferred KM instance for level zero key.

Prefer KM4.1 and higher over KM4.0 and lower, but prefer TEE over
Strongbox if TEE meets the minimal requirements.

Ignore-AOSP-First: No automerge path from AOSP.
Bug: 187862706
Test: Manually tested by observing logs during boot.
Merged-In: I1d27c80ef7c869b84b6d0c1a5d8eec287c242f6c
Change-Id: I1d27c80ef7c869b84b6d0c1a5d8eec287c242f6c
Merged-In: I1d27c80ef7c869b84b6d0c1a5d8eec287c242f6c
diff --git a/keystore2/src/boot_level_keys.rs b/keystore2/src/boot_level_keys.rs
index 3084195..5b7c3c3 100644
--- a/keystore2/src/boot_level_keys.rs
+++ b/keystore2/src/boot_level_keys.rs
@@ -14,6 +14,7 @@
 
 //! Offer keys based on the "boot level" for superencryption.
 
+use crate::{database::KeystoreDB, key_parameter::KeyParameterValue, raw_device::KeyMintDevice};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     Algorithm::Algorithm, Digest::Digest, KeyPurpose::KeyPurpose, SecurityLevel::SecurityLevel,
 };
@@ -21,26 +22,46 @@
 use keystore2_crypto::{hkdf_expand, ZVec, AES_256_KEY_LENGTH};
 use std::{collections::VecDeque, convert::TryFrom};
 
-use crate::{database::KeystoreDB, key_parameter::KeyParameterValue, raw_device::KeyMintDevice};
+fn get_preferred_km_instance_for_level_zero_key() -> Result<KeyMintDevice> {
+    let tee = KeyMintDevice::get(SecurityLevel::TRUSTED_ENVIRONMENT)
+        .context("In get_preferred_km_instance_for_level_zero_key: Get TEE instance failed.")?;
+    if tee.version() >= KeyMintDevice::KEY_MASTER_V4_1 {
+        Ok(tee)
+    } else {
+        match KeyMintDevice::get_or_none(SecurityLevel::STRONGBOX).context(
+            "In get_preferred_km_instance_for_level_zero_key: Get Strongbox instance failed.",
+        )? {
+            Some(strongbox) if strongbox.version() >= KeyMintDevice::KEY_MASTER_V4_1 => {
+                Ok(strongbox)
+            }
+            _ => Ok(tee),
+        }
+    }
+}
 
 /// This is not thread safe; caller must hold a lock before calling.
 /// In practice the caller is SuperKeyManager and the lock is the
 /// Mutex on its internal state.
 pub fn get_level_zero_key(db: &mut KeystoreDB) -> Result<ZVec> {
+    let km_dev = get_preferred_km_instance_for_level_zero_key()
+        .context("In get_level_zero_key: get preferred KM instance failed")?;
+
     let key_desc = KeyMintDevice::internal_descriptor("boot_level_key".to_string());
-    let params = [
+    let mut params = vec![
         KeyParameterValue::Algorithm(Algorithm::HMAC).into(),
         KeyParameterValue::Digest(Digest::SHA_2_256).into(),
         KeyParameterValue::KeySize(256).into(),
         KeyParameterValue::MinMacLength(256).into(),
         KeyParameterValue::KeyPurpose(KeyPurpose::SIGN).into(),
         KeyParameterValue::NoAuthRequired.into(),
-        KeyParameterValue::MaxUsesPerBoot(1).into(),
     ];
-    // We use TRUSTED_ENVIRONMENT here because it is the authority on when
-    // the device has rebooted.
-    let km_dev: KeyMintDevice = KeyMintDevice::get(SecurityLevel::TRUSTED_ENVIRONMENT)
-        .context("In get_level_zero_key: KeyMintDevice::get failed")?;
+
+    if km_dev.version() >= KeyMintDevice::KEY_MASTER_V4_1 {
+        params.push(KeyParameterValue::EarlyBootOnly.into());
+    } else {
+        params.push(KeyParameterValue::MaxUsesPerBoot(1).into())
+    }
+
     let (key_id_guard, key_entry) = km_dev
         .lookup_or_generate_key(db, &key_desc, &params)
         .context("In get_level_zero_key: lookup_or_generate_key failed")?;
diff --git a/keystore2/src/key_parameter.rs b/keystore2/src/key_parameter.rs
index 549f574..de142a7 100644
--- a/keystore2/src/key_parameter.rs
+++ b/keystore2/src/key_parameter.rs
@@ -825,6 +825,9 @@
     /// When deleted, the key is guaranteed to be permanently deleted and unusable
     #[key_param(tag = ROLLBACK_RESISTANCE, field = BoolValue)]
     RollbackResistance,
+    /// The Key shall only be used during the early boot stage
+    #[key_param(tag = EARLY_BOOT_ONLY, field = BoolValue)]
+    EarlyBootOnly,
     /// The date and time at which the key becomes active
     #[key_param(tag = ACTIVE_DATETIME, field = DateTime)]
     ActiveDateTime(i64),
diff --git a/keystore2/src/raw_device.rs b/keystore2/src/raw_device.rs
index 9e6ef41..9cc69ef 100644
--- a/keystore2/src/raw_device.rs
+++ b/keystore2/src/raw_device.rs
@@ -19,16 +19,15 @@
         BlobMetaData, BlobMetaEntry, CertificateInfo, DateTime, KeyEntry, KeyEntryLoadBits,
         KeyIdGuard, KeyMetaData, KeyMetaEntry, KeyType, KeystoreDB, SubComponentType, Uuid,
     },
-    error::{map_km_error, Error},
+    error::{map_km_error, Error, ErrorCode},
     globals::get_keymint_device,
     super_key::KeyBlob,
     utils::{key_characteristics_to_internal, watchdog as wd, Asp, AID_KEYSTORE},
 };
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    BeginResult::BeginResult, ErrorCode::ErrorCode, HardwareAuthToken::HardwareAuthToken,
-    IKeyMintDevice::IKeyMintDevice, IKeyMintOperation::IKeyMintOperation,
-    KeyCreationResult::KeyCreationResult, KeyParameter::KeyParameter, KeyPurpose::KeyPurpose,
-    SecurityLevel::SecurityLevel,
+    BeginResult::BeginResult, HardwareAuthToken::HardwareAuthToken, IKeyMintDevice::IKeyMintDevice,
+    IKeyMintOperation::IKeyMintOperation, KeyCreationResult::KeyCreationResult,
+    KeyParameter::KeyParameter, KeyPurpose::KeyPurpose, SecurityLevel::SecurityLevel,
 };
 use android_system_keystore2::aidl::android::system::keystore2::{
     Domain::Domain, KeyDescriptor::KeyDescriptor, ResponseCode::ResponseCode,
@@ -48,14 +47,39 @@
 pub struct KeyMintDevice {
     asp: Asp,
     km_uuid: Uuid,
+    version: i32,
 }
 
 impl KeyMintDevice {
+    /// Version number of KeyMasterDevice@V4_0
+    pub const KEY_MASTER_V4_0: i32 = 40;
+    /// Version number of KeyMasterDevice@V4_1
+    pub const KEY_MASTER_V4_1: i32 = 41;
+    /// Version number of KeyMintDevice@V1
+    pub const KEY_MINT_V1: i32 = 100;
+
     /// Get a [`KeyMintDevice`] for the given [`SecurityLevel`]
     pub fn get(security_level: SecurityLevel) -> Result<KeyMintDevice> {
-        let (asp, _hw_info, km_uuid) = get_keymint_device(&security_level)
+        let (asp, hw_info, km_uuid) = get_keymint_device(&security_level)
             .context("In KeyMintDevice::get: get_keymint_device failed")?;
-        Ok(KeyMintDevice { asp, km_uuid })
+
+        Ok(KeyMintDevice { asp, km_uuid, version: hw_info.versionNumber })
+    }
+
+    /// Get a [`KeyMintDevice`] for the given [`SecurityLevel`], return
+    /// [`None`] if the error `HARDWARE_TYPE_UNAVAILABLE` is returned
+    pub fn get_or_none(security_level: SecurityLevel) -> Result<Option<KeyMintDevice>> {
+        KeyMintDevice::get(security_level).map(Some).or_else(|e| {
+            match e.root_cause().downcast_ref::<Error>() {
+                Some(Error::Km(ErrorCode::HARDWARE_TYPE_UNAVAILABLE)) => Ok(None),
+                _ => Err(e),
+            }
+        })
+    }
+
+    /// Returns the version of the underlying KeyMint/KeyMaster device.
+    pub fn version(&self) -> i32 {
+        self.version
     }
 
     /// Create a KM key and store in the database.
@@ -145,14 +169,24 @@
         //   KeyMint operations
         let lookup = Self::not_found_is_none(Self::lookup_from_desc(db, key_desc))
             .context("In lookup_or_generate_key: first lookup failed")?;
-        if let Some(result) = lookup {
-            Ok(result)
-        } else {
-            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)
-                .context("In lookup_or_generate_key: second lookup failed")
+
+        if let Some((key_id_guard, key_entry)) = lookup {
+            // If the key is associated with a different km instance
+            // or if there is no blob metadata for some reason the key entry
+            // is considered corrupted and needs to be replaced with a new one.
+            if key_entry
+                .key_blob_info()
+                .as_ref()
+                .map(|(_, blob_metadata)| Some(&self.km_uuid) == blob_metadata.km_uuid())
+                .unwrap_or(false)
+            {
+                return Ok((key_id_guard, key_entry));
+            }
         }
+        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)
+            .context("In lookup_or_generate_key: second lookup failed")
     }
 
     /// Call the passed closure; if it returns `KEY_REQUIRES_UPGRADE`, call upgradeKey, and