Add tests for super_key.rs

The new tests are focused on unlocking, resetting and removing a user.
The tests verify that keys are deleted when necessary and that the user
state transitions properly.

Bug: 280502317
Test: atest keystore2_test on cuttlefish
Change-Id: Idae5d99fb289045bb277ba6c93ab62cfd9aed6fb
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index d7c939c..10cadfe 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -4660,7 +4660,7 @@
         params
     }
 
-    fn make_test_key_entry(
+    pub fn make_test_key_entry(
         db: &mut KeystoreDB,
         domain: Domain,
         namespace: i64,
@@ -4719,7 +4719,7 @@
         }
     }
 
-    fn make_bootlevel_key_entry(
+    pub fn make_bootlevel_key_entry(
         db: &mut KeystoreDB,
         domain: Domain,
         namespace: i64,
diff --git a/keystore2/src/legacy_importer.rs b/keystore2/src/legacy_importer.rs
index 9eb702d..325c213 100644
--- a/keystore2/src/legacy_importer.rs
+++ b/keystore2/src/legacy_importer.rs
@@ -102,6 +102,11 @@
         }
     }
 
+    #[cfg(test)]
+    pub fn set_empty(&mut self) {
+        self.state = AtomicU8::new(Self::STATE_EMPTY);
+    }
+
     /// The legacy importer must be initialized deferred, because keystore starts very early.
     /// At this time the data partition may not be mounted. So we cannot open database connections
     /// until we get actual key load requests. This sets the function that the legacy loader
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index fc82538..2e8b60f 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -490,7 +490,8 @@
     ) -> Result<Arc<SuperKey>> {
         let super_key = Self::extract_super_key_from_key_entry(algorithm, entry, pw, None)
             .context(ks_err!("Failed to extract super key from key entry"))?;
-        self.install_per_boot_key_for_user(user_id, super_key.clone())?;
+        self.install_per_boot_key_for_user(user_id, super_key.clone())
+            .context(ks_err!("Failed to install per boot key for user!"))?;
         Ok(super_key)
     }
 
@@ -1181,3 +1182,390 @@
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::database::tests::make_bootlevel_key_entry;
+    use crate::database::tests::make_test_key_entry;
+    use crate::database::tests::new_test_db;
+    use rand::prelude::*;
+    const USER_ID: u32 = 0;
+    const TEST_KEY_ALIAS: &str = "TEST_KEY";
+    const TEST_BOOT_KEY_ALIAS: &str = "TEST_BOOT_KEY";
+
+    pub fn generate_password_blob() -> Password<'static> {
+        let mut rng = rand::thread_rng();
+        let mut password = vec![0u8; 64];
+        rng.fill_bytes(&mut password);
+
+        let mut zvec = ZVec::new(64).expect("Failed to create ZVec");
+        zvec[..].copy_from_slice(&password[..]);
+
+        Password::Owned(zvec)
+    }
+
+    fn setup_test(pw: &Password) -> (Arc<RwLock<SuperKeyManager>>, KeystoreDB, LegacyImporter) {
+        let mut keystore_db = new_test_db().unwrap();
+        let mut legacy_importer = LegacyImporter::new(Arc::new(Default::default()));
+        legacy_importer.set_empty();
+        let skm: Arc<RwLock<SuperKeyManager>> = Default::default();
+        assert!(skm
+            .write()
+            .unwrap()
+            .init_user(&mut keystore_db, &legacy_importer, USER_ID, pw)
+            .is_ok());
+        (skm, keystore_db, legacy_importer)
+    }
+
+    fn assert_unlocked(
+        skm: &Arc<RwLock<SuperKeyManager>>,
+        keystore_db: &mut KeystoreDB,
+        legacy_importer: &LegacyImporter,
+        user_id: u32,
+        err_msg: &str,
+    ) {
+        let user_state =
+            skm.write().unwrap().get_user_state(keystore_db, legacy_importer, user_id).unwrap();
+        match user_state {
+            UserState::LskfUnlocked(_) => {}
+            _ => panic!("{}", err_msg),
+        }
+    }
+
+    fn assert_locked(
+        skm: &Arc<RwLock<SuperKeyManager>>,
+        keystore_db: &mut KeystoreDB,
+        legacy_importer: &LegacyImporter,
+        user_id: u32,
+        err_msg: &str,
+    ) {
+        let user_state =
+            skm.write().unwrap().get_user_state(keystore_db, legacy_importer, user_id).unwrap();
+        match user_state {
+            UserState::LskfLocked => {}
+            _ => panic!("{}", err_msg),
+        }
+    }
+
+    fn assert_uninitialized(
+        skm: &Arc<RwLock<SuperKeyManager>>,
+        keystore_db: &mut KeystoreDB,
+        legacy_importer: &LegacyImporter,
+        user_id: u32,
+        err_msg: &str,
+    ) {
+        let user_state =
+            skm.write().unwrap().get_user_state(keystore_db, legacy_importer, user_id).unwrap();
+        match user_state {
+            UserState::Uninitialized => {}
+            _ => panic!("{}", err_msg),
+        }
+    }
+
+    #[test]
+    fn test_init_user() {
+        let pw: Password = generate_password_blob();
+        let (skm, mut keystore_db, legacy_importer) = setup_test(&pw);
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not unlocked after initialization!",
+        );
+    }
+
+    #[test]
+    fn test_unlock_user() {
+        let pw: Password = generate_password_blob();
+        let (skm, mut keystore_db, legacy_importer) = setup_test(&pw);
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not unlocked after initialization!",
+        );
+
+        skm.write().unwrap().data.user_keys.clear();
+        assert_locked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "Clearing the cache did not lock the user!",
+        );
+
+        assert!(skm
+            .write()
+            .unwrap()
+            .unlock_user(&mut keystore_db, &legacy_importer, USER_ID, &pw)
+            .is_ok());
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user did not unlock!",
+        );
+    }
+
+    #[test]
+    fn test_unlock_wrong_password() {
+        let pw: Password = generate_password_blob();
+        let wrong_pw: Password = generate_password_blob();
+        let (skm, mut keystore_db, legacy_importer) = setup_test(&pw);
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not unlocked after initialization!",
+        );
+
+        skm.write().unwrap().data.user_keys.clear();
+        assert_locked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "Clearing the cache did not lock the user!",
+        );
+
+        assert!(skm
+            .write()
+            .unwrap()
+            .unlock_user(&mut keystore_db, &legacy_importer, USER_ID, &wrong_pw)
+            .is_err());
+        assert_locked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was unlocked with an incorrect password!",
+        );
+    }
+
+    #[test]
+    fn test_unlock_user_idempotent() {
+        let pw: Password = generate_password_blob();
+        let (skm, mut keystore_db, legacy_importer) = setup_test(&pw);
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not unlocked after initialization!",
+        );
+
+        skm.write().unwrap().data.user_keys.clear();
+        assert_locked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "Clearing the cache did not lock the user!",
+        );
+
+        for _ in 0..5 {
+            assert!(skm
+                .write()
+                .unwrap()
+                .unlock_user(&mut keystore_db, &legacy_importer, USER_ID, &pw)
+                .is_ok());
+            assert_unlocked(
+                &skm,
+                &mut keystore_db,
+                &legacy_importer,
+                USER_ID,
+                "The user did not unlock!",
+            );
+        }
+    }
+
+    fn test_user_removal(locked: bool) {
+        let pw: Password = generate_password_blob();
+        let (skm, mut keystore_db, legacy_importer) = setup_test(&pw);
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not unlocked after initialization!",
+        );
+
+        assert!(make_test_key_entry(
+            &mut keystore_db,
+            Domain::APP,
+            USER_ID.into(),
+            TEST_KEY_ALIAS,
+            None
+        )
+        .is_ok());
+        assert!(make_bootlevel_key_entry(
+            &mut keystore_db,
+            Domain::APP,
+            USER_ID.into(),
+            TEST_BOOT_KEY_ALIAS,
+            false
+        )
+        .is_ok());
+
+        assert!(keystore_db
+            .key_exists(Domain::APP, USER_ID.into(), TEST_KEY_ALIAS, KeyType::Client)
+            .unwrap());
+        assert!(keystore_db
+            .key_exists(Domain::APP, USER_ID.into(), TEST_BOOT_KEY_ALIAS, KeyType::Client)
+            .unwrap());
+
+        if locked {
+            skm.write().unwrap().data.user_keys.clear();
+            assert_locked(
+                &skm,
+                &mut keystore_db,
+                &legacy_importer,
+                USER_ID,
+                "Clearing the cache did not lock the user!",
+            );
+        }
+
+        assert!(skm
+            .write()
+            .unwrap()
+            .remove_user(&mut keystore_db, &legacy_importer, USER_ID)
+            .is_ok());
+        assert_uninitialized(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not removed!",
+        );
+
+        assert!(!skm
+            .write()
+            .unwrap()
+            .super_key_exists_in_db_for_user(&mut keystore_db, &legacy_importer, USER_ID)
+            .unwrap());
+
+        assert!(!keystore_db
+            .key_exists(Domain::APP, USER_ID.into(), TEST_KEY_ALIAS, KeyType::Client)
+            .unwrap());
+        assert!(!keystore_db
+            .key_exists(Domain::APP, USER_ID.into(), TEST_BOOT_KEY_ALIAS, KeyType::Client)
+            .unwrap());
+    }
+
+    fn test_user_reset(locked: bool) {
+        let pw: Password = generate_password_blob();
+        let (skm, mut keystore_db, legacy_importer) = setup_test(&pw);
+        assert_unlocked(
+            &skm,
+            &mut keystore_db,
+            &legacy_importer,
+            USER_ID,
+            "The user was not unlocked after initialization!",
+        );
+
+        assert!(make_test_key_entry(
+            &mut keystore_db,
+            Domain::APP,
+            USER_ID.into(),
+            TEST_KEY_ALIAS,
+            None
+        )
+        .is_ok());
+        assert!(make_bootlevel_key_entry(
+            &mut keystore_db,
+            Domain::APP,
+            USER_ID.into(),
+            TEST_BOOT_KEY_ALIAS,
+            false
+        )
+        .is_ok());
+        assert!(keystore_db
+            .key_exists(Domain::APP, USER_ID.into(), TEST_KEY_ALIAS, KeyType::Client)
+            .unwrap());
+        assert!(keystore_db
+            .key_exists(Domain::APP, USER_ID.into(), TEST_BOOT_KEY_ALIAS, KeyType::Client)
+            .unwrap());
+
+        if locked {
+            skm.write().unwrap().data.user_keys.clear();
+            assert_locked(
+                &skm,
+                &mut keystore_db,
+                &legacy_importer,
+                USER_ID,
+                "Clearing the cache did not lock the user!",
+            );
+            assert!(skm
+                .write()
+                .unwrap()
+                .reset_user(&mut keystore_db, &legacy_importer, USER_ID)
+                .is_err());
+            assert_locked(
+                &skm,
+                &mut keystore_db,
+                &legacy_importer,
+                USER_ID,
+                "User state should not have changed!",
+            );
+
+            // Keys should still exist.
+            assert!(keystore_db
+                .key_exists(Domain::APP, USER_ID.into(), TEST_KEY_ALIAS, KeyType::Client)
+                .unwrap());
+            assert!(keystore_db
+                .key_exists(Domain::APP, USER_ID.into(), TEST_BOOT_KEY_ALIAS, KeyType::Client)
+                .unwrap());
+        } else {
+            assert!(skm
+                .write()
+                .unwrap()
+                .reset_user(&mut keystore_db, &legacy_importer, USER_ID)
+                .is_ok());
+            assert_uninitialized(
+                &skm,
+                &mut keystore_db,
+                &legacy_importer,
+                USER_ID,
+                "The user was not reset!",
+            );
+            assert!(!skm
+                .write()
+                .unwrap()
+                .super_key_exists_in_db_for_user(&mut keystore_db, &legacy_importer, USER_ID)
+                .unwrap());
+
+            // Auth bound key should no longer exist.
+            assert!(!keystore_db
+                .key_exists(Domain::APP, USER_ID.into(), TEST_KEY_ALIAS, KeyType::Client)
+                .unwrap());
+            assert!(keystore_db
+                .key_exists(Domain::APP, USER_ID.into(), TEST_BOOT_KEY_ALIAS, KeyType::Client)
+                .unwrap());
+        }
+    }
+
+    #[test]
+    fn test_remove_unlocked_user() {
+        test_user_removal(false);
+    }
+
+    #[test]
+    fn test_remove_locked_user() {
+        test_user_removal(true);
+    }
+
+    #[test]
+    fn test_reset_unlocked_user() {
+        test_user_reset(false);
+    }
+
+    #[test]
+    fn test_reset_locked_user() {
+        test_user_reset(true);
+    }
+}