Keystore 2.0: Enforcement for limited use keys.

Bug: b/174140443
Test: atest keystore2_test
Change-Id: I6433b7c7f305d67d8e77277990f732d634801f10
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index ffec931..344fe2f 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -42,7 +42,7 @@
 //! callbacks.
 
 use crate::db_utils::{self, SqlField};
-use crate::error::{Error as KsError, ResponseCode};
+use crate::error::{Error as KsError, ErrorCode, ResponseCode};
 use crate::impl_metadata; // This is in db_utils.rs
 use crate::key_parameter::{KeyParameter, Tag};
 use crate::permission::KeyPermSet;
@@ -1332,6 +1332,43 @@
         Ok(parameters)
     }
 
+    /// Decrements the usage count of a limited use key. This function first checks whether the
+    /// usage has been exhausted, if not, decreases the usage count. If the usage count reaches
+    /// zero, the key also gets marked unreferenced and scheduled for deletion.
+    /// Returns Ok(true) if the key was marked unreferenced as a hint to the garbage collector.
+    pub fn check_and_update_key_usage_count(&mut self, key_id: i64) -> Result<bool> {
+        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+            let limit: Option<i32> = tx
+                .query_row(
+                    "SELECT data FROM persistent.keyparameter WHERE keyentryid = ? AND tag = ?;",
+                    params![key_id, Tag::USAGE_COUNT_LIMIT.0],
+                    |row| row.get(0),
+                )
+                .optional()
+                .context("Trying to load usage count")?;
+
+            let limit = limit
+                .ok_or(KsError::Km(ErrorCode::INVALID_KEY_BLOB))
+                .context("The Key no longer exists. Key is exhausted.")?;
+
+            tx.execute(
+                "UPDATE persistent.keyparameter
+                 SET data = data - 1
+                 WHERE keyentryid = ? AND tag = ? AND data > 0;",
+                params![key_id, Tag::USAGE_COUNT_LIMIT.0],
+            )
+            .context("Failed to update key usage count.")?;
+
+            match limit {
+                1 => Self::mark_unreferenced(tx, key_id)
+                    .context("Trying to mark limited use key for deletion."),
+                0 => Err(KsError::Km(ErrorCode::INVALID_KEY_BLOB)).context("Key is exhausted."),
+                _ => Ok(false),
+            }
+        })
+        .context("In check_and_update_key_usage_count.")
+    }
+
     /// Load a key entry by the given key descriptor.
     /// It uses the `check_permission` callback to verify if the access is allowed
     /// given the key access tuple read from the database using `load_access_tuple`.
@@ -2223,7 +2260,7 @@
     #[test]
     fn test_insert_and_load_full_keyentry_domain_app() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_domain_app")?
             .0;
         let (_key_guard, key_entry) = db
@@ -2240,7 +2277,7 @@
                 |_k, _av| Ok(()),
             )
             .unwrap();
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(
             KeyDescriptor {
@@ -2280,7 +2317,7 @@
     #[test]
     fn test_insert_and_load_full_keyentry_domain_selinux() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_domain_selinux")?
             .0;
         let (_key_guard, key_entry) = db
@@ -2297,7 +2334,7 @@
                 |_k, _av| Ok(()),
             )
             .unwrap();
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(
             KeyDescriptor {
@@ -2337,7 +2374,7 @@
     #[test]
     fn test_insert_and_load_full_keyentry_domain_key_id() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_domain_key_id")?
             .0;
         let (_, key_entry) = db
@@ -2350,7 +2387,7 @@
             )
             .unwrap();
 
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(
             KeyDescriptor { domain: Domain::KEY_ID, nspace: key_id, alias: None, blob: None },
@@ -2378,9 +2415,57 @@
     }
 
     #[test]
+    fn test_check_and_update_key_usage_count_with_limited_use_key() -> Result<()> {
+        let mut db = new_test_db()?;
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, Some(123))
+            .context("test_check_and_update_key_usage_count_with_limited_use_key")?
+            .0;
+        // Update the usage count of the limited use key.
+        db.check_and_update_key_usage_count(key_id)?;
+
+        let (_key_guard, key_entry) = db.load_key_entry(
+            KeyDescriptor { domain: Domain::KEY_ID, nspace: key_id, alias: None, blob: None },
+            KeyType::Client,
+            KeyEntryLoadBits::BOTH,
+            1,
+            |_k, _av| Ok(()),
+        )?;
+
+        // The usage count is decremented now.
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, Some(122)));
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_check_and_update_key_usage_count_with_exhausted_limited_use_key() -> Result<()> {
+        let mut db = new_test_db()?;
+        let key_id = make_test_key_entry(&mut db, Domain::SELINUX, 1, TEST_ALIAS, Some(1))
+            .context("test_check_and_update_key_usage_count_with_exhausted_limited_use_key")?
+            .0;
+        // Update the usage count of the limited use key.
+        db.check_and_update_key_usage_count(key_id).expect(concat!(
+            "In test_check_and_update_key_usage_count_with_exhausted_limited_use_key: ",
+            "This should succeed."
+        ));
+
+        // Try to update the exhausted limited use key.
+        let e = db.check_and_update_key_usage_count(key_id).expect_err(concat!(
+            "In test_check_and_update_key_usage_count_with_exhausted_limited_use_key: ",
+            "This should fail."
+        ));
+        assert_eq!(
+            &KsError::Km(ErrorCode::INVALID_KEY_BLOB),
+            e.root_cause().downcast_ref::<KsError>().unwrap()
+        );
+
+        Ok(())
+    }
+
+    #[test]
     fn test_insert_and_load_full_keyentry_from_grant() -> Result<()> {
         let mut db = new_test_db()?;
-        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS)
+        let key_id = make_test_key_entry(&mut db, Domain::APP, 1, TEST_ALIAS, None)
             .context("test_insert_and_load_full_keyentry_from_grant")?
             .0;
 
@@ -2415,7 +2500,7 @@
             )
             .unwrap();
 
-        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+        assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
 
         db.unbind_key(granted_key.clone(), KeyType::Client, 2, |_, _| Ok(())).unwrap();
 
@@ -2444,7 +2529,7 @@
             let temp_dir = Arc::new(TempDir::new("id_lock_test")?);
             let temp_dir_clone = temp_dir.clone();
             let mut db = KeystoreDB::new(temp_dir.path())?;
-            let key_id = make_test_key_entry(&mut db, Domain::APP, 33, KEY_LOCK_TEST_ALIAS)
+            let key_id = make_test_key_entry(&mut db, Domain::APP, 33, KEY_LOCK_TEST_ALIAS, None)
                 .context("test_insert_and_load_full_keyentry_domain_app")?
                 .0;
             let (_key_guard, key_entry) = db
@@ -2461,7 +2546,7 @@
                     |_k, _av| Ok(()),
                 )
                 .unwrap();
-            assert_eq!(key_entry, make_test_key_entry_test_vector(key_id));
+            assert_eq!(key_entry, make_test_key_entry_test_vector(key_id, None));
             let state = Arc::new(AtomicU8::new(1));
             let state2 = state.clone();
 
@@ -2543,8 +2628,8 @@
         let list_o_keys: Vec<(i64, i64)> = LIST_O_ENTRIES
             .iter()
             .map(|(domain, ns, alias)| {
-                let entry =
-                    make_test_key_entry(&mut db, *domain, *ns, *alias).unwrap_or_else(|e| {
+                let entry = make_test_key_entry(&mut db, *domain, *ns, *alias, None)
+                    .unwrap_or_else(|e| {
                         panic!("Failed to insert {:?} {} {}. Error {:?}", domain, ns, alias, e)
                     });
                 (entry.id(), *ns)
@@ -2652,8 +2737,8 @@
     // Note: The parameters and SecurityLevel associations are nonsensical. This
     // collection is only used to check if the parameters are preserved as expected by the
     // database.
-    fn make_test_params() -> Vec<KeyParameter> {
-        vec![
+    fn make_test_params(max_usage_count: Option<i32>) -> Vec<KeyParameter> {
+        let mut params = vec![
             KeyParameter::new(KeyParameterValue::Invalid, SecurityLevel::TRUSTED_ENVIRONMENT),
             KeyParameter::new(
                 KeyParameterValue::KeyPurpose(KeyPurpose::SIGN),
@@ -2868,7 +2953,14 @@
                 KeyParameterValue::ConfirmationToken(vec![5u8, 5u8, 5u8, 5u8]),
                 SecurityLevel::TRUSTED_ENVIRONMENT,
             ),
-        ]
+        ];
+        if let Some(value) = max_usage_count {
+            params.push(KeyParameter::new(
+                KeyParameterValue::UsageCountLimit(value),
+                SecurityLevel::SOFTWARE,
+            ));
+        }
+        params
     }
 
     fn make_test_key_entry(
@@ -2876,12 +2968,16 @@
         domain: Domain,
         namespace: i64,
         alias: &str,
+        max_usage_count: Option<i32>,
     ) -> Result<KeyIdGuard> {
         let key_id = db.create_key_entry(domain, namespace)?;
         db.insert_blob(&key_id, SubComponentType::KEY_BLOB, TEST_KEY_BLOB)?;
         db.insert_blob(&key_id, SubComponentType::CERT, TEST_CERT_BLOB)?;
         db.insert_blob(&key_id, SubComponentType::CERT_CHAIN, TEST_CERT_CHAIN_BLOB)?;
-        db.insert_keyparameter(&key_id, &make_test_params())?;
+
+        let params = make_test_params(max_usage_count);
+        db.insert_keyparameter(&key_id, &params)?;
+
         let mut metadata = KeyMetaData::new();
         metadata.add(KeyMetaEntry::EncryptedBy(EncryptedBy::Password));
         metadata.add(KeyMetaEntry::Salt(vec![1, 2, 3]));
@@ -2892,7 +2988,9 @@
         Ok(key_id)
     }
 
-    fn make_test_key_entry_test_vector(key_id: i64) -> KeyEntry {
+    fn make_test_key_entry_test_vector(key_id: i64, max_usage_count: Option<i32>) -> KeyEntry {
+        let params = make_test_params(max_usage_count);
+
         let mut metadata = KeyMetaData::new();
         metadata.add(KeyMetaEntry::EncryptedBy(EncryptedBy::Password));
         metadata.add(KeyMetaEntry::Salt(vec![1, 2, 3]));
@@ -2905,7 +3003,7 @@
             cert: Some(TEST_CERT_BLOB.to_vec()),
             cert_chain: Some(TEST_CERT_CHAIN_BLOB.to_vec()),
             sec_level: SecurityLevel::TRUSTED_ENVIRONMENT,
-            parameters: make_test_params(),
+            parameters: params,
             metadata,
         }
     }
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 22b4bed..274ddf8 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -14,10 +14,13 @@
 
 //! This is the Keystore 2.0 Enforcements module.
 // TODO: more description to follow.
-use crate::database::{AuthTokenEntry, MonotonicRawTime};
 use crate::error::{map_binder_status, Error, ErrorCode};
 use crate::globals::{get_timestamp_service, ASYNC_TASK, DB, ENFORCEMENTS};
 use crate::key_parameter::{KeyParameter, KeyParameterValue};
+use crate::{
+    database::{AuthTokenEntry, MonotonicRawTime},
+    gc::Gc,
+};
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     Algorithm::Algorithm, ErrorCode::ErrorCode as Ec, HardwareAuthToken::HardwareAuthToken,
     HardwareAuthenticatorType::HardwareAuthenticatorType,
@@ -128,6 +131,8 @@
 #[derive(Debug)]
 pub struct AuthInfo {
     state: DeferredAuthState,
+    /// An optional key id required to update the usage count if the key usage is limited.
+    key_usage_limited: Option<i64>,
 }
 
 struct TokenReceiverMap {
@@ -251,14 +256,46 @@
         }
     }
 
+    /// This function is the authorization hook called before operation update.
+    /// It returns the auth tokens required by the operation to commence update.
+    pub fn before_update(&mut self) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
+        self.get_auth_tokens()
+    }
+
+    /// This function is the authorization hook called before operation finish.
+    /// It returns the auth tokens required by the operation to commence finish.
+    pub fn before_finish(&mut self) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
+        self.get_auth_tokens()
+    }
+
+    /// This function is the authorization hook called after finish succeeded.
+    /// As of this writing it checks if the key was a limited use key. If so it updates the
+    /// use counter of the key in the database. When the use counter is depleted, the key gets
+    /// marked for deletion and the garbage collector is notified.
+    pub fn after_finish(&self) -> Result<()> {
+        if let Some(key_id) = self.key_usage_limited {
+            // On the last successful use, the key gets deleted. In this case we
+            // have to notify the garbage collector.
+            let need_gc = DB
+                .with(|db| {
+                    db.borrow_mut()
+                        .check_and_update_key_usage_count(key_id)
+                        .context("Trying to update key usage count.")
+                })
+                .context("In after_finish.")?;
+            if need_gc {
+                Gc::notify_gc();
+            }
+        }
+        Ok(())
+    }
+
     /// This function returns the auth tokens as needed by the ongoing operation or fails
     /// with ErrorCode::KEY_USER_NOT_AUTHENTICATED. If this was called for the first time
     /// after a deferred authorization was requested by finalize_create_authorization, this
     /// function may block on the generation of a time stamp token. It then moves the
     /// tokens into the DeferredAuthState::Token state for future use.
-    pub fn get_auth_tokens(
-        &mut self,
-    ) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
+    fn get_auth_tokens(&mut self) -> Result<(Option<HardwareAuthToken>, Option<TimeStampToken>)> {
         let deferred_tokens = if let DeferredAuthState::Waiting(ref auth_request) = self.state {
             let mut state = auth_request.lock().unwrap();
             Some(state.get_auth_tokens().context("In AuthInfo::get_auth_tokens.")?)
@@ -325,14 +362,18 @@
     pub fn authorize_create(
         &self,
         purpose: KeyPurpose,
-        key_params: Option<&[KeyParameter]>,
+        key_properties: Option<&(i64, Vec<KeyParameter>)>,
         op_params: &[KmKeyParameter],
         requires_timestamp: bool,
     ) -> Result<(Option<HardwareAuthToken>, AuthInfo)> {
-        let key_params = if let Some(k) = key_params {
-            k
-        } else {
-            return Ok((None, AuthInfo { state: DeferredAuthState::NoAuthRequired }));
+        let (key_id, key_params) = match key_properties {
+            Some((key_id, key_params)) => (*key_id, key_params),
+            None => {
+                return Ok((
+                    None,
+                    AuthInfo { state: DeferredAuthState::NoAuthRequired, key_usage_limited: None },
+                ))
+            }
         };
 
         match purpose {
@@ -377,6 +418,7 @@
         let mut key_time_out: Option<i64> = None;
         let mut allow_while_on_body = false;
         let mut unlocked_device_required = false;
+        let mut key_usage_limited: Option<i64> = None;
 
         // iterate through key parameters, recording information we need for authorization
         // enforcements later, or enforcing authorizations in place, where applicable
@@ -437,6 +479,12 @@
                 KeyParameterValue::AllowWhileOnBody => {
                     allow_while_on_body = true;
                 }
+                KeyParameterValue::UsageCountLimit(_) => {
+                    // We don't examine the limit here because this is enforced on finish.
+                    // Instead, we store the key_id so that finish can look up the key
+                    // in the database again and check and update the counter.
+                    key_usage_limited = Some(key_id);
+                }
                 // NOTE: as per offline discussion, sanitizing key parameters and rejecting
                 // create operation if any non-allowed tags are present, is not done in
                 // authorize_create (unlike in legacy keystore where AuthorizeBegin is rejected if
@@ -491,7 +539,10 @@
         }
 
         if !unlocked_device_required && no_auth_required {
-            return Ok((None, AuthInfo { state: DeferredAuthState::NoAuthRequired }));
+            return Ok((
+                None,
+                AuthInfo { state: DeferredAuthState::NoAuthRequired, key_usage_limited },
+            ));
         }
 
         let has_sids = !user_secure_ids.is_empty();
@@ -565,7 +616,7 @@
             (None, _, true) => (None, DeferredAuthState::OpAuthRequired),
             (None, _, false) => (None, DeferredAuthState::NoAuthRequired),
         })
-        .map(|(hat, state)| (hat, AuthInfo { state }))
+        .map(|(hat, state)| (hat, AuthInfo { state, key_usage_limited }))
     }
 
     fn find_auth_token<F>(p: F) -> Result<Option<(AuthTokenEntry, MonotonicRawTime)>>
diff --git a/keystore2/src/key_parameter.rs b/keystore2/src/key_parameter.rs
index 71a17fe..93de6f2 100644
--- a/keystore2/src/key_parameter.rs
+++ b/keystore2/src/key_parameter.rs
@@ -840,6 +840,9 @@
     /// Maximum number of times that a key may be used between system reboots
     #[key_param(tag = MAX_USES_PER_BOOT, field = Integer)]
     MaxUsesPerBoot(i32),
+    /// The number of times that a limited use key can be used
+    #[key_param(tag = USAGE_COUNT_LIMIT, field = Integer)]
+    UsageCountLimit(i32),
     /// ID of the Android user that is permitted to use the key
     #[key_param(tag = USER_ID, field = Integer)]
     UserID(i32),
diff --git a/keystore2/src/km_compat/km_compat_type_conversion.h b/keystore2/src/km_compat/km_compat_type_conversion.h
index 97c61ce..895f749 100644
--- a/keystore2/src/km_compat/km_compat_type_conversion.h
+++ b/keystore2/src/km_compat/km_compat_type_conversion.h
@@ -389,6 +389,9 @@
             return V4_0::makeKeyParameter(V4_0::TAG_MAX_USES_PER_BOOT, v->get());
         }
         break;
+    case KMV1::Tag::USAGE_COUNT_LIMIT:
+        // Does not exist in KM < KeyMint 1.0.
+        break;
     case KMV1::Tag::USER_ID:
         if (auto v = KMV1::authorizationValue(KMV1::TAG_USER_ID, kp)) {
             return V4_0::makeKeyParameter(V4_0::TAG_USER_ID, v->get());
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 30e6d55..829987d 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -341,7 +341,7 @@
             .auth_info
             .lock()
             .unwrap()
-            .get_auth_tokens()
+            .before_update()
             .context("In update_aad: Trying to get auth tokens.")?;
 
         self.update_outcome(
@@ -377,7 +377,7 @@
             .auth_info
             .lock()
             .unwrap()
-            .get_auth_tokens()
+            .before_update()
             .context("In update: Trying to get auth tokens.")?;
 
         self.update_outcome(
@@ -423,7 +423,7 @@
             .auth_info
             .lock()
             .unwrap()
-            .get_auth_tokens()
+            .before_finish()
             .context("In finish: Trying to get auth tokens.")?;
 
         let output = self
@@ -440,6 +440,8 @@
             )
             .context("In finish: KeyMint::finish failed.")?;
 
+        self.auth_info.lock().unwrap().after_finish().context("In finish.")?;
+
         // At this point the operation concluded successfully.
         *outcome = Outcome::Success;
 
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index d8787bd..dad241a 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -34,7 +34,7 @@
 use crate::globals::ENFORCEMENTS;
 use crate::key_parameter::KeyParameter as KsKeyParam;
 use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
-use crate::utils::{check_key_permission, Asp};
+use crate::utils::{check_key_permission, uid_to_android_user, Asp};
 use crate::{database::KeyIdGuard, globals::DB};
 use crate::{
     database::{DateTime, KeyMetaData, KeyMetaEntry, KeyType},
@@ -48,7 +48,6 @@
 use crate::{
     error::{self, map_km_error, map_or_log_err, Error, ErrorCode},
     utils::key_characteristics_to_internal,
-    utils::uid_to_android_user,
 };
 use anyhow::{Context, Result};
 use binder::{IBinder, Interface, ThreadState};
@@ -173,7 +172,7 @@
         // so that we can use it by reference like the blob provided by the key descriptor.
         // Otherwise, we would have to clone the blob from the key descriptor.
         let scoping_blob: Vec<u8>;
-        let (km_blob, key_id_guard, key_parameters) = match key.domain {
+        let (km_blob, key_properties, key_id_guard) = match key.domain {
             Domain::BLOB => {
                 check_key_permission(KeyPerm::use_(), key, &None)
                     .context("In create_operation: checking use permission for Domain::BLOB.")?;
@@ -212,7 +211,11 @@
                         ))
                     }
                 };
-                (&scoping_blob, Some(key_id_guard), Some(key_entry.into_key_parameters()))
+                (
+                    &scoping_blob,
+                    Some((key_id_guard.id(), key_entry.into_key_parameters())),
+                    Some(key_id_guard),
+                )
             }
         };
 
@@ -229,8 +232,8 @@
         let (immediate_hat, mut auth_info) = ENFORCEMENTS
             .authorize_create(
                 purpose,
-                key_parameters.as_deref(),
-                operation_parameters,
+                key_properties.as_ref(),
+                operation_parameters.as_ref(),
                 // TODO b/178222844 Replace this with the configuration returned by
                 //      KeyMintDevice::getHardwareInfo.
                 //      For now we assume that strongbox implementations need secure timestamps.