Make 'password unlock' depend on user management infrastructure.

This CL addresses an existing TODO item in password unlock functionality
that builds on top of user management.

Bug: 176123105
Test: TBD
Change-Id: If81c52466f90b65ae8e7d300b48ab5e353c15353
diff --git a/keystore2/src/authorization.rs b/keystore2/src/authorization.rs
index fbaa9eb..3190541 100644
--- a/keystore2/src/authorization.rs
+++ b/keystore2/src/authorization.rs
@@ -16,8 +16,9 @@
 
 use crate::error::Error as KeystoreError;
 use crate::error::map_or_log_err;
-use crate::globals::{DB, ENFORCEMENTS, LEGACY_BLOB_LOADER, SUPER_KEY};
+use crate::globals::{ENFORCEMENTS, SUPER_KEY, DB};
 use crate::permission::KeystorePerm;
+use crate::super_key::UserState;
 use crate::utils::check_keystore_permission;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
     HardwareAuthToken::HardwareAuthToken,
@@ -65,22 +66,21 @@
                     .context("In on_lock_screen_event: Unlock with password.")?;
                 ENFORCEMENTS.set_device_locked(user_id, false);
                 // Unlock super key.
-                DB.with::<_, Result<()>>(|db| {
-                    let mut db = db.borrow_mut();
-                    //TODO - b/176123105 - Once the user management API is implemented, unlock is
-                    //allowed only if the user is added. Then the two tasks handled by the
-                    //unlock_user_key will be split into two methods. For now, unlock_user_key
-                    //method is used as it is, which created a super key for the user if one does
-                    //not exists, in addition to unlocking the existing super key of the user/
-                    SUPER_KEY.unlock_user_key(
-                        &mut db,
-                        user_id as u32,
-                        user_password,
-                        &LEGACY_BLOB_LOADER,
-                    )?;
-                    Ok(())
-                })
-                .context("In on_lock_screen_event.")?;
+                if let UserState::Uninitialized = DB
+                    .with(|db| {
+                        UserState::get_with_password_unlock(
+                            &mut db.borrow_mut(),
+                            &SUPER_KEY,
+                            user_id as u32,
+                            user_password,
+                        )
+                    })
+                    .context("In on_lock_screen_event: Unlock with password.")?
+                {
+                    log::info!(
+                        "In on_lock_screen_event. Trying to unlock when LSKF is uninitialized."
+                    );
+                }
 
                 Ok(())
             }
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 2663c6e..9767d32 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -1167,6 +1167,32 @@
         .context("In store_super_key.")
     }
 
+    /// Loads super key of a given user, if exists
+    pub fn load_super_key(&mut self, user_id: u32) -> Result<Option<(KeyIdGuard, KeyEntry)>> {
+        self.with_transaction(TransactionBehavior::Immediate, |tx| {
+            let key_descriptor = KeyDescriptor {
+                domain: Domain::APP,
+                nspace: user_id as u64 as i64,
+                alias: Some(String::from("USER_SUPER_KEY")),
+                blob: None,
+            };
+            let id = Self::load_key_entry_id(&tx, &key_descriptor, KeyType::Super);
+            match id {
+                Ok(id) => {
+                    let key_entry = Self::load_key_components(&tx, KeyEntryLoadBits::KM, id)
+                        .context("In load_super_key. Failed to load key entry.")?;
+                    Ok(Some((KEY_ID_LOCK.get(id), key_entry)))
+                }
+                Err(error) => match error.root_cause().downcast_ref::<KsError>() {
+                    Some(KsError::Rc(ResponseCode::KEY_NOT_FOUND)) => Ok(None),
+                    _ => Err(error).context("In load_super_key."),
+                },
+            }
+            .no_gc()
+        })
+        .context("In load_super_key.")
+    }
+
     /// Atomically loads a key entry and associated metadata or creates it using the
     /// callback create_new_key callback. The callback is called during a database
     /// transaction. This means that implementers should be mindful about using
@@ -4730,16 +4756,7 @@
         //check if super key exists
         assert!(db.key_exists(Domain::APP, 1, "USER_SUPER_KEY", KeyType::Super)?);
 
-        //load the super key from the database
-        let tx = db.conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
-        let key_descriptor = KeyDescriptor {
-            domain: Domain::APP,
-            nspace: 1,
-            alias: Some(String::from("USER_SUPER_KEY")),
-            blob: None,
-        };
-        let id = KeystoreDB::load_key_entry_id(&tx, &key_descriptor, KeyType::Super)?;
-        let key_entry = KeystoreDB::load_key_components(&tx, KeyEntryLoadBits::KM, id)?;
+        let (_, key_entry) = db.load_super_key(1)?.unwrap();
         let loaded_super_key = SuperKeyManager::extract_super_key_from_key_entry(key_entry, &pw)?;
 
         let decrypted_secret_bytes = keystore2_crypto::aes_gcm_decrypt(
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index 2df5343..dd29d7e 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -221,6 +221,34 @@
     }
 
     /// Checks if user has already setup LSKF (i.e. a super key is persisted in the database or the
+    /// legacy database). If not, return Uninitialized state.
+    /// Otherwise, decrypt the super key from the password and return LskfUnlocked state.
+    pub fn check_and_unlock_super_key(
+        &self,
+        db: &mut KeystoreDB,
+        user_id: u32,
+        pw: &[u8],
+    ) -> Result<UserState> {
+        let result = db
+            .load_super_key(user_id)
+            .context("In check_and_unlock_super_key. Failed to load super key")?;
+
+        match result {
+            Some((_, entry)) => {
+                let super_key = self
+                    .populate_cache_from_super_key_blob(user_id, entry, pw)
+                    .context("In check_and_unlock_super_key.")?;
+                Ok(UserState::LskfUnlocked(super_key))
+            }
+            None =>
+            //TODO: 159371296. Try to load and populate super key from legacy key database.
+            {
+                Ok(UserState::Uninitialized)
+            }
+        }
+    }
+
+    /// Checks if user has already setup LSKF (i.e. a super key is persisted in the database or the
     /// legacy database). If so, return LskfLocked state.
     /// If the password is provided, generate a new super key, encrypt with the password,
     /// store in the database and populate the super key cache for the new user
@@ -234,7 +262,6 @@
     ) -> Result<UserState> {
         let super_key_exists_in_db = Self::super_key_exists_in_db_for_user(db, user_id)
             .context("In check_and_initialize_super_key. Failed to check if super key exists.")?;
-
         if super_key_exists_in_db {
             Ok(UserState::LskfLocked)
         } else {
@@ -504,6 +531,29 @@
         }
     }
 
+    /// Queries user state when serving password unlock requests.
+    pub fn get_with_password_unlock(
+        db: &mut KeystoreDB,
+        skm: &SuperKeyManager,
+        user_id: u32,
+        password: &[u8],
+    ) -> Result<UserState> {
+        match skm.get_per_boot_key_by_user_id(user_id) {
+            Some(super_key) => {
+                log::info!("In get_with_password_unlock. Trying to unlock when already unlocked.");
+                Ok(UserState::LskfUnlocked(super_key))
+            }
+            None => {
+                //Check if a super key exists in the database or legacy database.
+                //If not, return Uninitialized state.
+                //Otherwise, try to unlock the super key and if successful,
+                //return LskfUnlocked state
+                skm.check_and_unlock_super_key(db, user_id, password)
+                    .context("In get_with_password_unlock. Failed to unlock super key.")
+            }
+        }
+    }
+
     /// Delete all the keys created on behalf of the user.
     /// If 'keep_non_super_encrypted_keys' is set to true, delete only the super key and super
     /// encrypted keys.