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,
         }
     }