Merge "-= src_available"
diff --git a/diced/Android.bp b/diced/Android.bp
index 525828e..e13d863 100644
--- a/diced/Android.bp
+++ b/diced/Android.bp
@@ -138,6 +138,24 @@
     init_rc: ["diced.rc"],
 }
 
+rust_binary {
+    name: "diced.microdroid",
+    srcs: ["src/diced_main.rs"],
+    prefer_rlib: true,
+    rustlibs: [
+        "android.hardware.security.dice-V1-rust",
+        "libandroid_logger",
+        "libbinder_rs",
+        "libdiced",
+        "libdiced_open_dice_cbor",
+        "libdiced_sample_inputs",
+        "libdiced_utils",
+        "liblog_rust",
+    ],
+    init_rc: ["diced.microdroid.rc"],
+    bootstrap: true,
+}
+
 rust_test {
     name: "diced_test",
     crate_name: "diced_test",
diff --git a/diced/diced.microdroid.rc b/diced/diced.microdroid.rc
new file mode 100644
index 0000000..2226f47
--- /dev/null
+++ b/diced/diced.microdroid.rc
@@ -0,0 +1,13 @@
+# Start the Diced service.
+#
+# See system/core/init/README.md for information on the init.rc language.
+
+service diced /system/bin/diced.microdroid
+    class main
+    user diced
+    group diced
+    # The diced service must not be allowed to restart.
+    # If it crashes for any reason security critical state is lost.
+    # The only remedy is to restart the device.
+    oneshot
+    writepid /dev/cpuset/foreground/tasks
diff --git a/identity/Android.bp b/identity/Android.bp
index 7b0503a..790a731 100644
--- a/identity/Android.bp
+++ b/identity/Android.bp
@@ -27,6 +27,7 @@
     defaults: [
         "identity_defaults",
         "keymint_use_latest_hal_aidl_ndk_shared",
+        "keymint_use_latest_hal_aidl_cpp_static",
     ],
 
     srcs: [
diff --git a/keystore2/aidl/android/security/remoteprovisioning/IRemotelyProvisionedKeyPool.aidl b/keystore2/aidl/android/security/remoteprovisioning/IRemotelyProvisionedKeyPool.aidl
new file mode 100644
index 0000000..7d45e52
--- /dev/null
+++ b/keystore2/aidl/android/security/remoteprovisioning/IRemotelyProvisionedKeyPool.aidl
@@ -0,0 +1,49 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.security.remoteprovisioning;
+
+import android.security.remoteprovisioning.RemotelyProvisionedKey;
+
+/**
+ * This is the interface providing access to remotely-provisioned attestation keys
+ * for an `IRemotelyProvisionedComponent`.
+ *
+ * @hide
+ */
+interface IRemotelyProvisionedKeyPool {
+
+    /**
+     * Fetches an attestation key for the given uid and `IRemotelyProvisionedComponent`, as
+     * identified by the given id.
+
+     * Callers require the keystore2::get_attestation_key permission.
+     *
+     * ## Error conditions
+     * `android.system.keystore2.ResponseCode::PERMISSION_DENIED` if the caller does not have the
+     *      `keystore2::get_attestation_key` permission
+     *
+     * @param clientUid The client application for which an attestation key is needed.
+     *
+     * @param irpcId The unique identifier for the `IRemotelyProvisionedComponent` for which a key
+     *      is requested. This id may be retrieved from a given component via the
+     *      `IRemotelyProvisionedComponent::getHardwareInfo` function.
+     *
+     * @return A `RemotelyProvisionedKey` parcelable containing a key and certification chain for
+     *      the given `IRemotelyProvisionedComponent`.
+     */
+    RemotelyProvisionedKey getAttestationKey(in int clientUid, in @utf8InCpp String irpcId);
+}
diff --git a/keystore2/aidl/android/security/remoteprovisioning/RemotelyProvisionedKey.aidl b/keystore2/aidl/android/security/remoteprovisioning/RemotelyProvisionedKey.aidl
new file mode 100644
index 0000000..ae21855
--- /dev/null
+++ b/keystore2/aidl/android/security/remoteprovisioning/RemotelyProvisionedKey.aidl
@@ -0,0 +1,42 @@
+/*
+ * Copyright 2021, The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package android.security.remoteprovisioning;
+
+/**
+ * A `RemotelyProvisionedKey` holds an attestation key and the corresponding remotely provisioned
+ * certificate chain.
+ *
+ * @hide
+ */
+@RustDerive(Eq=true, PartialEq=true)
+parcelable RemotelyProvisionedKey {
+    /**
+     * The remotely-provisioned key that may be used to sign attestations. The format of this key
+     * is opaque, and need only be understood by the IRemotelyProvisionedComponent that generated
+     * it.
+     *
+     * Any private key material contained within this blob must be encrypted.
+     */
+    byte[] keyBlob;
+
+    /**
+     * Sequence of DER-encoded X.509 certificates that make up the attestation key's certificate
+     * chain. This is the binary encoding for a chain that is supported by Java's
+     * CertificateFactory.generateCertificates API.
+     */
+    byte[] encodedCertChain;
+}
diff --git a/keystore2/src/authorization.rs b/keystore2/src/authorization.rs
index 64b498f..81790af 100644
--- a/keystore2/src/authorization.rs
+++ b/keystore2/src/authorization.rs
@@ -154,8 +154,10 @@
                     .context("In on_lock_screen_event: Unlock with password.")?;
                 ENFORCEMENTS.set_device_locked(user_id, false);
 
+                let mut skm = SUPER_KEY.write().unwrap();
+
                 DB.with(|db| {
-                    SUPER_KEY.unlock_screen_lock_bound_key(
+                    skm.unlock_screen_lock_bound_key(
                         &mut db.borrow_mut(),
                         user_id as u32,
                         &password,
@@ -166,10 +168,9 @@
                 // Unlock super key.
                 if let UserState::Uninitialized = DB
                     .with(|db| {
-                        UserState::get_with_password_unlock(
+                        skm.unlock_and_get_user_state(
                             &mut db.borrow_mut(),
                             &LEGACY_MIGRATOR,
-                            &SUPER_KEY,
                             user_id as u32,
                             &password,
                         )
@@ -187,8 +188,9 @@
                 check_keystore_permission(KeystorePerm::Unlock)
                     .context("In on_lock_screen_event: Unlock.")?;
                 ENFORCEMENTS.set_device_locked(user_id, false);
+                let mut skm = SUPER_KEY.write().unwrap();
                 DB.with(|db| {
-                    SUPER_KEY.try_unlock_user_with_biometric(&mut db.borrow_mut(), user_id as u32)
+                    skm.try_unlock_user_with_biometric(&mut db.borrow_mut(), user_id as u32)
                 })
                 .context("In on_lock_screen_event: try_unlock_user_with_biometric failed")?;
                 Ok(())
@@ -197,8 +199,9 @@
                 check_keystore_permission(KeystorePerm::Lock)
                     .context("In on_lock_screen_event: Lock")?;
                 ENFORCEMENTS.set_device_locked(user_id, true);
+                let mut skm = SUPER_KEY.write().unwrap();
                 DB.with(|db| {
-                    SUPER_KEY.lock_screen_lock_bound_key(
+                    skm.lock_screen_lock_bound_key(
                         &mut db.borrow_mut(),
                         user_id as u32,
                         unlocking_sids.unwrap_or(&[]),
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 7099f5a..e68b0fd 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -3207,7 +3207,7 @@
 }
 
 #[cfg(test)]
-mod tests {
+pub mod tests {
 
     use super::*;
     use crate::key_parameter::{
@@ -3231,13 +3231,13 @@
     use std::collections::BTreeMap;
     use std::fmt::Write;
     use std::sync::atomic::{AtomicU8, Ordering};
-    use std::sync::Arc;
+    use std::sync::{Arc, RwLock};
     use std::thread;
     use std::time::{Duration, SystemTime};
     #[cfg(disabled)]
     use std::time::Instant;
 
-    fn new_test_db() -> Result<KeystoreDB> {
+    pub fn new_test_db() -> Result<KeystoreDB> {
         let conn = KeystoreDB::make_connection("file::memory:")?;
 
         let mut db = KeystoreDB { conn, gc: None, perboot: Arc::new(perboot::PerbootDB::new()) };
@@ -3251,7 +3251,7 @@
     where
         F: Fn(&Uuid, &[u8]) -> Result<()> + Send + 'static,
     {
-        let super_key: Arc<SuperKeyManager> = Default::default();
+        let super_key: Arc<RwLock<SuperKeyManager>> = Default::default();
 
         let gc_db = KeystoreDB::new(path, None).expect("Failed to open test gc db_connection.");
         let gc = Gc::new_init_with(Default::default(), move || (Box::new(cb), gc_db, super_key));
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 2407525..ade4751 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -602,7 +602,7 @@
         }
 
         if let Some(level) = max_boot_level {
-            if !SUPER_KEY.level_accessible(level) {
+            if !SUPER_KEY.read().unwrap().level_accessible(level) {
                 return Err(Error::Km(Ec::BOOT_LEVEL_EXCEEDED))
                     .context("In authorize_create: boot level is too late.");
             }
diff --git a/keystore2/src/error.rs b/keystore2/src/error.rs
index 42dd3d2..f34c5da 100644
--- a/keystore2/src/error.rs
+++ b/keystore2/src/error.rs
@@ -67,10 +67,15 @@
         Error::Rc(ResponseCode::SYSTEM_ERROR)
     }
 
-    /// Short hand for `Error::Rc(ResponseCode::PERMISSION_DENIED`
+    /// Short hand for `Error::Rc(ResponseCode::PERMISSION_DENIED)`
     pub fn perm() -> Self {
         Error::Rc(ResponseCode::PERMISSION_DENIED)
     }
+
+    /// Short hand for `Error::Rc(ResponseCode::OUT_OF_KEYS)`
+    pub fn out_of_keys() -> Self {
+        Error::Rc(ResponseCode::OUT_OF_KEYS)
+    }
 }
 
 /// Helper function to map the binder status we get from calls into KeyMint
diff --git a/keystore2/src/gc.rs b/keystore2/src/gc.rs
index 25f08c8..341aa0a 100644
--- a/keystore2/src/gc.rs
+++ b/keystore2/src/gc.rs
@@ -27,7 +27,7 @@
 use async_task::AsyncTask;
 use std::sync::{
     atomic::{AtomicU8, Ordering},
-    Arc,
+    Arc, RwLock,
 };
 
 pub struct Gc {
@@ -47,7 +47,7 @@
         F: FnOnce() -> (
                 Box<dyn Fn(&Uuid, &[u8]) -> Result<()> + Send + 'static>,
                 KeystoreDB,
-                Arc<SuperKeyManager>,
+                Arc<RwLock<SuperKeyManager>>,
             ) + Send
             + 'static,
     {
@@ -87,7 +87,7 @@
     invalidate_key: Box<dyn Fn(&Uuid, &[u8]) -> Result<()> + Send + 'static>,
     db: KeystoreDB,
     async_task: std::sync::Weak<AsyncTask>,
-    super_key: Arc<SuperKeyManager>,
+    super_key: Arc<RwLock<SuperKeyManager>>,
     notified: Arc<AtomicU8>,
 }
 
@@ -121,6 +121,8 @@
             if let Some(uuid) = blob_metadata.km_uuid() {
                 let blob = self
                     .super_key
+                    .read()
+                    .unwrap()
                     .unwrap_key_if_required(&blob_metadata, &blob)
                     .context("In process_one_key: Trying to unwrap to-be-deleted blob.")?;
                 (self.invalidate_key)(uuid, &*blob)
diff --git a/keystore2/src/globals.rs b/keystore2/src/globals.rs
index 7028aae..2819314 100644
--- a/keystore2/src/globals.rs
+++ b/keystore2/src/globals.rs
@@ -156,7 +156,7 @@
     pub static ref DB_PATH: RwLock<PathBuf> = RwLock::new(
         Path::new("/data/misc/keystore").to_path_buf());
     /// Runtime database of unwrapped super keys.
-    pub static ref SUPER_KEY: Arc<SuperKeyManager> = Default::default();
+    pub static ref SUPER_KEY: Arc<RwLock<SuperKeyManager>> = Default::default();
     /// Map of KeyMint devices.
     static ref KEY_MINT_DEVICES: Mutex<DevicesMap<dyn IKeyMintDevice>> = Default::default();
     /// Timestamp service.
diff --git a/keystore2/src/legacy_blob.rs b/keystore2/src/legacy_blob.rs
index 7454cca..b801ed3 100644
--- a/keystore2/src/legacy_blob.rs
+++ b/keystore2/src/legacy_blob.rs
@@ -1340,7 +1340,7 @@
             CACERT_NON_AUTHBOUND,
         )?;
 
-        let key_manager: SuperKeyManager = Default::default();
+        let mut key_manager: SuperKeyManager = Default::default();
         let mut db = crate::database::KeystoreDB::new(temp_dir.path(), None)?;
         let legacy_blob_loader = LegacyBlobLoader::new(temp_dir.path());
 
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 8b629b1..134f707 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -40,7 +40,6 @@
 pub mod security_level;
 pub mod service;
 pub mod shared_secret_negotiation;
-pub mod try_insert;
 pub mod utils;
 
 mod attestation_key_utils;
diff --git a/keystore2/src/maintenance.rs b/keystore2/src/maintenance.rs
index d5feee1..9925e42 100644
--- a/keystore2/src/maintenance.rs
+++ b/keystore2/src/maintenance.rs
@@ -21,7 +21,7 @@
 use crate::globals::get_keymint_device;
 use crate::globals::{DB, LEGACY_MIGRATOR, SUPER_KEY};
 use crate::permission::{KeyPerm, KeystorePerm};
-use crate::super_key::UserState;
+use crate::super_key::{SuperKeyManager, UserState};
 use crate::utils::{
     check_key_permission, check_keystore_permission, list_key_entries, watchdog as wd,
 };
@@ -70,24 +70,25 @@
     }
 
     fn on_user_password_changed(user_id: i32, password: Option<Password>) -> Result<()> {
-        //Check permission. Function should return if this failed. Therefore having '?' at the end
-        //is very important.
+        // Check permission. Function should return if this failed. Therefore having '?' at the end
+        // is very important.
         check_keystore_permission(KeystorePerm::ChangePassword)
             .context("In on_user_password_changed.")?;
 
+        let mut skm = SUPER_KEY.write().unwrap();
+
         if let Some(pw) = password.as_ref() {
             DB.with(|db| {
-                SUPER_KEY.unlock_screen_lock_bound_key(&mut db.borrow_mut(), user_id as u32, pw)
+                skm.unlock_screen_lock_bound_key(&mut db.borrow_mut(), user_id as u32, pw)
             })
             .context("In on_user_password_changed: unlock_screen_lock_bound_key failed")?;
         }
 
         match DB
             .with(|db| {
-                UserState::get_with_password_changed(
+                skm.reset_or_init_user_and_get_user_state(
                     &mut db.borrow_mut(),
                     &LEGACY_MIGRATOR,
-                    &SUPER_KEY,
                     user_id as u32,
                     password.as_ref(),
                 )
@@ -110,10 +111,10 @@
         // Check permission. Function should return if this failed. Therefore having '?' at the end
         // is very important.
         check_keystore_permission(KeystorePerm::ChangeUser).context("In add_or_remove_user.")?;
+
         DB.with(|db| {
-            UserState::reset_user(
+            SUPER_KEY.write().unwrap().reset_user(
                 &mut db.borrow_mut(),
-                &SUPER_KEY,
                 &LEGACY_MIGRATOR,
                 user_id as u32,
                 false,
@@ -145,7 +146,11 @@
         check_keystore_permission(KeystorePerm::GetState).context("In get_state.")?;
         let state = DB
             .with(|db| {
-                UserState::get(&mut db.borrow_mut(), &LEGACY_MIGRATOR, &SUPER_KEY, user_id as u32)
+                SUPER_KEY.read().unwrap().get_user_state(
+                    &mut db.borrow_mut(),
+                    &LEGACY_MIGRATOR,
+                    user_id as u32,
+                )
             })
             .context("In get_state. Trying to get UserState.")?;
 
@@ -202,7 +207,9 @@
             .context("In early_boot_ended. Checking permission")?;
         log::info!("In early_boot_ended.");
 
-        if let Err(e) = DB.with(|db| SUPER_KEY.set_up_boot_level_cache(&mut db.borrow_mut())) {
+        if let Err(e) =
+            DB.with(|db| SuperKeyManager::set_up_boot_level_cache(&SUPER_KEY, &mut db.borrow_mut()))
+        {
             log::error!("SUPER_KEY.set_up_boot_level_cache failed:\n{:?}\n:(", e);
         }
         Maintenance::call_on_all_security_levels("earlyBootEnded", |dev| dev.earlyBootEnded())
diff --git a/keystore2/src/permission.rs b/keystore2/src/permission.rs
index e6d61b0..1e6f10a 100644
--- a/keystore2/src/permission.rs
+++ b/keystore2/src/permission.rs
@@ -149,6 +149,9 @@
         /// introduced for migrating keys when an app leaves a sharedUserId.
         #[selinux(name = migrate_any_key)]
         MigrateAnyKey,
+        /// Checked on calls to IRemotelyProvisionedKeyPool::getAttestationKey
+        #[selinux(name = get_attestation_key)]
+        GetAttestationKey,
     }
 );
 
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index 132ffbe..fadd252 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -30,11 +30,13 @@
 };
 use android_security_remoteprovisioning::aidl::android::security::remoteprovisioning::{
     AttestationPoolStatus::AttestationPoolStatus, IRemoteProvisioning::BnRemoteProvisioning,
-    IRemoteProvisioning::IRemoteProvisioning, ImplInfo::ImplInfo,
+    IRemoteProvisioning::IRemoteProvisioning,
+    IRemotelyProvisionedKeyPool::IRemotelyProvisionedKeyPool, ImplInfo::ImplInfo,
+    RemotelyProvisionedKey::RemotelyProvisionedKey,
 };
 use android_security_remoteprovisioning::binder::{BinderFeatures, Strong};
 use android_system_keystore2::aidl::android::system::keystore2::{
-    Domain::Domain, KeyDescriptor::KeyDescriptor,
+    Domain::Domain, KeyDescriptor::KeyDescriptor, ResponseCode::ResponseCode,
 };
 use anyhow::{Context, Result};
 use keystore2_crypto::parse_subject_from_certificate;
@@ -46,7 +48,8 @@
 use crate::error::{self, map_or_log_err, map_rem_prov_error, Error};
 use crate::globals::{get_keymint_device, get_remotely_provisioned_component, DB};
 use crate::metrics_store::log_rkp_error_stats;
-use crate::utils::watchdog as wd;
+use crate::permission::KeystorePerm;
+use crate::utils::{check_keystore_permission, watchdog as wd};
 use android_security_metrics::aidl::android::security::metrics::RkpError::RkpError as MetricsRkpError;
 
 /// Contains helper functions to check if remote provisioning is enabled on the system and, if so,
@@ -90,70 +93,6 @@
         Ok(pool_status.total != 0)
     }
 
-    /// Fetches a remote provisioning attestation key and certificate chain inside of the
-    /// returned `CertificateChain` struct if one exists for the given caller_uid. If one has not
-    /// been assigned, this function will assign it. If there are no signed attestation keys
-    /// available to be assigned, it will return the ResponseCode `OUT_OF_KEYS`
-    fn get_rem_prov_attest_key(
-        &self,
-        key: &KeyDescriptor,
-        caller_uid: u32,
-        db: &mut KeystoreDB,
-    ) -> Result<Option<CertificateChain>> {
-        match key.domain {
-            Domain::APP => {
-                // Attempt to get an Attestation Key once. If it fails, then the app doesn't
-                // have a valid chain assigned to it. The helper function will return None after
-                // attempting to assign a key. An error will be thrown if the pool is simply out
-                // of usable keys. Then another attempt to fetch the just-assigned key will be
-                // made. If this fails too, something is very wrong.
-                self.get_rem_prov_attest_key_helper(key, caller_uid, db)
-                    .context("In get_rem_prov_attest_key: Failed to get a key")?
-                    .map_or_else(
-                        || self.get_rem_prov_attest_key_helper(key, caller_uid, db),
-                        |v| Ok(Some(v)),
-                    )
-                    .context(concat!(
-                        "In get_rem_prov_attest_key: Failed to get a key after",
-                        "attempting to assign one."
-                    ))?
-                    .map_or_else(
-                        || {
-                            Err(Error::sys()).context(concat!(
-                                "In get_rem_prov_attest_key: Attempted to assign a ",
-                                "key and failed silently. Something is very wrong."
-                            ))
-                        },
-                        |cert_chain| Ok(Some(cert_chain)),
-                    )
-            }
-            _ => Ok(None),
-        }
-    }
-
-    /// Returns None if an AttestationKey fails to be assigned. Errors if no keys are available.
-    fn get_rem_prov_attest_key_helper(
-        &self,
-        key: &KeyDescriptor,
-        caller_uid: u32,
-        db: &mut KeystoreDB,
-    ) -> Result<Option<CertificateChain>> {
-        let cert_chain = db
-            .retrieve_attestation_key_and_cert_chain(key.domain, caller_uid as i64, &self.km_uuid)
-            .context("In get_rem_prov_attest_key_helper: Failed to retrieve a key + cert chain")?;
-        match cert_chain {
-            Some(cert_chain) => Ok(Some(cert_chain)),
-            // Either this app needs to be assigned a key, or the pool is empty. An error will
-            // be thrown if there is no key available to assign. This will indicate that the app
-            // should be nudged to provision more keys so keystore can retry.
-            None => {
-                db.assign_attestation_key(key.domain, caller_uid as i64, &self.km_uuid)
-                    .context("In get_rem_prov_attest_key_helper: Failed to assign a key")?;
-                Ok(None)
-            }
-        }
-    }
-
     fn is_asymmetric_key(&self, params: &[KeyParameter]) -> bool {
         params.iter().any(|kp| {
             matches!(
@@ -189,7 +128,7 @@
             // and therefore will not be attested.
             Ok(None)
         } else {
-            match self.get_rem_prov_attest_key(key, caller_uid, db) {
+            match get_rem_prov_attest_key(key.domain, caller_uid, db, &self.km_uuid) {
                 Err(e) => {
                     log::error!(
                         concat!(
@@ -233,9 +172,9 @@
     fn get_dev_by_sec_level(
         &self,
         sec_level: &SecurityLevel,
-    ) -> Result<Strong<dyn IRemotelyProvisionedComponent>> {
+    ) -> Result<&dyn IRemotelyProvisionedComponent> {
         if let Some(dev) = self.device_by_sec_level.get(sec_level) {
-            Ok(dev.clone())
+            Ok(dev.as_ref())
         } else {
             Err(error::Error::sys()).context(concat!(
                 "In get_dev_by_sec_level: Remote instance for requested security level",
@@ -244,6 +183,22 @@
         }
     }
 
+    fn get_dev_by_unique_id(
+        &self,
+        unique_id: &str,
+    ) -> Result<(SecurityLevel, &dyn IRemotelyProvisionedComponent)> {
+        for (sec_level, dev) in &self.device_by_sec_level {
+            if dev.getHardwareInfo()?.uniqueId == Some(unique_id.to_string()) {
+                return Ok((*sec_level, dev.as_ref()));
+            }
+        }
+
+        Err(error::Error::sys()).context(format!(
+            "In get_dev_by_unique_id: Instance for requested unique id '{}' not found",
+            unique_id
+        ))
+    }
+
     /// Creates a new instance of the remote provisioning service
     pub fn new_native_binder() -> Result<Strong<dyn IRemoteProvisioning>> {
         let mut result: Self = Default::default();
@@ -346,23 +301,21 @@
     /// here.
     pub fn provision_cert_chain(
         &self,
+        db: &mut KeystoreDB,
         public_key: &[u8],
         batch_cert: &[u8],
         certs: &[u8],
         expiration_date: i64,
         sec_level: SecurityLevel,
     ) -> Result<()> {
-        DB.with::<_, Result<()>>(|db| {
-            let mut db = db.borrow_mut();
-            let (_, _, uuid) = get_keymint_device(&sec_level)?;
-            db.store_signed_attestation_certificate_chain(
-                public_key,
-                batch_cert,
-                certs, /* DER encoded certificate chain */
-                expiration_date,
-                &uuid,
-            )
-        })
+        let (_, _, uuid) = get_keymint_device(&sec_level)?;
+        db.store_signed_attestation_certificate_chain(
+            public_key,
+            batch_cert,
+            certs, /* DER encoded certificate chain */
+            expiration_date,
+            &uuid,
+        )
     }
 
     fn parse_cose_mac0_for_coords(data: &[u8]) -> Result<Vec<u8>> {
@@ -429,19 +382,25 @@
     /// `is_test_mode` indicates whether or not the returned public key should be marked as being
     /// for testing in order to differentiate them from private keys. If the call is successful,
     /// the key pair is then added to the database.
-    pub fn generate_key_pair(&self, is_test_mode: bool, sec_level: SecurityLevel) -> Result<()> {
+    pub fn generate_key_pair(
+        &self,
+        db: &mut KeystoreDB,
+        is_test_mode: bool,
+        sec_level: SecurityLevel,
+    ) -> Result<()> {
         let (_, _, uuid) = get_keymint_device(&sec_level)?;
-        let dev = self.get_dev_by_sec_level(&sec_level)?;
+        let dev = self.get_dev_by_sec_level(&sec_level).context(format!(
+            "In generate_key_pair: Failed to get device for security level {:?}",
+            sec_level
+        ))?;
         let mut maced_key = MacedPublicKey { macedKey: Vec::new() };
         let priv_key =
             map_rem_prov_error(dev.generateEcdsaP256KeyPair(is_test_mode, &mut maced_key))
                 .context("In generate_key_pair: Failed to generated ECDSA keypair.")?;
         let raw_key = Self::parse_cose_mac0_for_coords(&maced_key.macedKey)
             .context("In generate_key_pair: Failed to parse raw key")?;
-        DB.with::<_, Result<()>>(|db| {
-            let mut db = db.borrow_mut();
-            db.create_attestation_key_entry(&maced_key.macedKey, &raw_key, &priv_key, &uuid)
-        })
+        db.create_attestation_key_entry(&maced_key.macedKey, &raw_key, &priv_key, &uuid)
+            .context("In generate_key_pair: Failed to insert attestation key entry")
     }
 
     /// Checks the security level of each available IRemotelyProvisionedComponent hal and returns
@@ -462,6 +421,35 @@
             db.delete_all_attestation_keys()
         })
     }
+
+    /// Fetches a remotely provisioned certificate chain and key for the given client uid that
+    /// was provisioned using the IRemotelyProvisionedComponent with the given id. The same key
+    /// will be returned for a given caller_uid on every request. If there are no attestation keys
+    /// available, `OUT_OF_KEYS` is returned.
+    fn get_attestation_key(
+        &self,
+        db: &mut KeystoreDB,
+        caller_uid: i32,
+        irpc_id: &str,
+    ) -> Result<RemotelyProvisionedKey> {
+        log::info!("get_attestation_key(self, {}, {}", caller_uid, irpc_id);
+
+        let (sec_level, _) = self.get_dev_by_unique_id(irpc_id)?;
+        let (_, _, km_uuid) = get_keymint_device(&sec_level)?;
+
+        let cert_chain = get_rem_prov_attest_key(Domain::APP, caller_uid as u32, db, &km_uuid)
+            .context("In get_attestation_key")?;
+        match cert_chain {
+            Some(chain) => Ok(RemotelyProvisionedKey {
+                keyBlob: chain.private_key.to_vec(),
+                encodedCertChain: chain.cert_chain,
+            }),
+            // It should be impossible to get `None`, but handle it just in case as a
+            // precaution against future behavioral changes in `get_rem_prov_attest_key`.
+            None => Err(error::Error::Rc(ResponseCode::OUT_OF_KEYS))
+                .context("In get_attestation_key: No available attestation keys"),
+        }
+    }
 }
 
 /// Populates the AttestationPoolStatus parcelable with information about how many
@@ -480,6 +468,70 @@
     })
 }
 
+/// Fetches a remote provisioning attestation key and certificate chain inside of the
+/// returned `CertificateChain` struct if one exists for the given caller_uid. If one has not
+/// been assigned, this function will assign it. If there are no signed attestation keys
+/// available to be assigned, it will return the ResponseCode `OUT_OF_KEYS`
+fn get_rem_prov_attest_key(
+    domain: Domain,
+    caller_uid: u32,
+    db: &mut KeystoreDB,
+    km_uuid: &Uuid,
+) -> Result<Option<CertificateChain>> {
+    match domain {
+        Domain::APP => {
+            // Attempt to get an Attestation Key once. If it fails, then the app doesn't
+            // have a valid chain assigned to it. The helper function will return None after
+            // attempting to assign a key. An error will be thrown if the pool is simply out
+            // of usable keys. Then another attempt to fetch the just-assigned key will be
+            // made. If this fails too, something is very wrong.
+            get_rem_prov_attest_key_helper(domain, caller_uid, db, km_uuid)
+                .context("In get_rem_prov_attest_key: Failed to get a key")?
+                .map_or_else(
+                    || get_rem_prov_attest_key_helper(domain, caller_uid, db, km_uuid),
+                    |v| Ok(Some(v)),
+                )
+                .context(concat!(
+                    "In get_rem_prov_attest_key: Failed to get a key after",
+                    "attempting to assign one."
+                ))?
+                .map_or_else(
+                    || {
+                        Err(Error::sys()).context(concat!(
+                            "In get_rem_prov_attest_key: Attempted to assign a ",
+                            "key and failed silently. Something is very wrong."
+                        ))
+                    },
+                    |cert_chain| Ok(Some(cert_chain)),
+                )
+        }
+        _ => Ok(None),
+    }
+}
+
+/// Returns None if an AttestationKey fails to be assigned. Errors if no keys are available.
+fn get_rem_prov_attest_key_helper(
+    domain: Domain,
+    caller_uid: u32,
+    db: &mut KeystoreDB,
+    km_uuid: &Uuid,
+) -> Result<Option<CertificateChain>> {
+    let cert_chain = db
+        .retrieve_attestation_key_and_cert_chain(domain, caller_uid as i64, km_uuid)
+        .context("In get_rem_prov_attest_key_helper: Failed to retrieve a key + cert chain")?;
+    match cert_chain {
+        Some(cert_chain) => Ok(Some(cert_chain)),
+        // Either this app needs to be assigned a key, or the pool is empty. An error will
+        // be thrown if there is no key available to assign. This will indicate that the app
+        // should be nudged to provision more keys so keystore can retry.
+        None => {
+            db.assign_attestation_key(domain, caller_uid as i64, km_uuid)
+                .context("In get_rem_prov_attest_key_helper: Failed to assign a key")?;
+            Ok(None)
+        }
+    }
+}
+
 impl binder::Interface for RemoteProvisioningService {}
 
 // Implementation of IRemoteProvisioning. See AIDL spec at
@@ -528,15 +580,29 @@
         sec_level: SecurityLevel,
     ) -> binder::Result<()> {
         let _wp = wd::watch_millis("IRemoteProvisioning::provisionCertChain", 500);
-        map_or_log_err(
-            self.provision_cert_chain(public_key, batch_cert, certs, expiration_date, sec_level),
-            Ok,
-        )
+        DB.with::<_, binder::Result<()>>(|db| {
+            map_or_log_err(
+                self.provision_cert_chain(
+                    &mut db.borrow_mut(),
+                    public_key,
+                    batch_cert,
+                    certs,
+                    expiration_date,
+                    sec_level,
+                ),
+                Ok,
+            )
+        })
     }
 
     fn generateKeyPair(&self, is_test_mode: bool, sec_level: SecurityLevel) -> binder::Result<()> {
         let _wp = wd::watch_millis("IRemoteProvisioning::generateKeyPair", 500);
-        map_or_log_err(self.generate_key_pair(is_test_mode, sec_level), Ok)
+        DB.with::<_, binder::Result<()>>(|db| {
+            map_or_log_err(
+                self.generate_key_pair(&mut db.borrow_mut(), is_test_mode, sec_level),
+                Ok,
+            )
+        })
     }
 
     fn getImplementationInfo(&self) -> binder::Result<Vec<ImplInfo>> {
@@ -550,11 +616,126 @@
     }
 }
 
+// Implementation of IRemotelyProvisionedKeyPool. See AIDL spec at
+// :aidl/android/security/remoteprovisioning/IRemotelyProvisionedKeyPool.aidl
+impl IRemotelyProvisionedKeyPool for RemoteProvisioningService {
+    fn getAttestationKey(
+        &self,
+        caller_uid: i32,
+        irpc_id: &str,
+    ) -> binder::Result<RemotelyProvisionedKey> {
+        let _wp = wd::watch_millis("IRemotelyProvisionedKeyPool::getAttestationKey", 500);
+        map_or_log_err(check_keystore_permission(KeystorePerm::GetAttestationKey), Ok)?;
+        DB.with::<_, binder::Result<RemotelyProvisionedKey>>(|db| {
+            map_or_log_err(self.get_attestation_key(&mut db.borrow_mut(), caller_uid, irpc_id), Ok)
+        })
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
     use serde_cbor::Value;
     use std::collections::BTreeMap;
+    use std::sync::{Arc, Mutex};
+    use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+        RpcHardwareInfo::RpcHardwareInfo,
+    };
+
+    #[derive(Default)]
+    struct MockRemotelyProvisionedComponentValues {
+        hw_info: RpcHardwareInfo,
+        private_key: Vec<u8>,
+        maced_public_key: Vec<u8>,
+    }
+
+    // binder::Interface requires the Send trait, so we have to use a Mutex even though the test
+    // is single threaded.
+    #[derive(Default)]
+    struct MockRemotelyProvisionedComponent(Arc<Mutex<MockRemotelyProvisionedComponentValues>>);
+
+    impl binder::Interface for MockRemotelyProvisionedComponent {}
+
+    impl IRemotelyProvisionedComponent for MockRemotelyProvisionedComponent {
+        fn getHardwareInfo(&self) -> binder::Result<RpcHardwareInfo> {
+            Ok(self.0.lock().unwrap().hw_info.clone())
+        }
+
+        fn generateEcdsaP256KeyPair(
+            &self,
+            test_mode: bool,
+            maced_public_key: &mut MacedPublicKey,
+        ) -> binder::Result<Vec<u8>> {
+            assert!(test_mode);
+            maced_public_key.macedKey = self.0.lock().unwrap().maced_public_key.clone();
+            Ok(self.0.lock().unwrap().private_key.clone())
+        }
+
+        fn generateCertificateRequest(
+            &self,
+            _test_mode: bool,
+            _keys_to_sign: &[MacedPublicKey],
+            _eek: &[u8],
+            _challenge: &[u8],
+            _device_info: &mut DeviceInfo,
+            _protected_data: &mut ProtectedData,
+        ) -> binder::Result<Vec<u8>> {
+            Err(binder::StatusCode::INVALID_OPERATION.into())
+        }
+    }
+
+    // Hard coded cert that can be parsed -- the content doesn't matter for testing, only that it's valid.
+    fn get_fake_cert() -> Vec<u8> {
+        vec![
+            0x30, 0x82, 0x01, 0xbb, 0x30, 0x82, 0x01, 0x61, 0xa0, 0x03, 0x02, 0x01, 0x02, 0x02,
+            0x14, 0x3a, 0xd5, 0x67, 0xce, 0xfe, 0x93, 0xe1, 0xea, 0xb7, 0xe4, 0xbf, 0x64, 0x19,
+            0xa4, 0x11, 0xe1, 0x87, 0x40, 0x20, 0x37, 0x30, 0x0a, 0x06, 0x08, 0x2a, 0x86, 0x48,
+            0xce, 0x3d, 0x04, 0x03, 0x02, 0x30, 0x33, 0x31, 0x0b, 0x30, 0x09, 0x06, 0x03, 0x55,
+            0x04, 0x06, 0x13, 0x02, 0x55, 0x54, 0x31, 0x13, 0x30, 0x11, 0x06, 0x03, 0x55, 0x04,
+            0x08, 0x0c, 0x0a, 0x53, 0x6f, 0x6d, 0x65, 0x2d, 0x53, 0x74, 0x61, 0x74, 0x65, 0x31,
+            0x0f, 0x30, 0x0d, 0x06, 0x03, 0x55, 0x04, 0x0a, 0x0c, 0x06, 0x47, 0x6f, 0x6f, 0x67,
+            0x6c, 0x65, 0x30, 0x1e, 0x17, 0x0d, 0x32, 0x31, 0x31, 0x32, 0x31, 0x30, 0x32, 0x32,
+            0x30, 0x38, 0x35, 0x32, 0x5a, 0x17, 0x0d, 0x34, 0x39, 0x30, 0x34, 0x32, 0x36, 0x32,
+            0x32, 0x30, 0x38, 0x35, 0x32, 0x5a, 0x30, 0x33, 0x31, 0x0b, 0x30, 0x09, 0x06, 0x03,
+            0x55, 0x04, 0x06, 0x13, 0x02, 0x55, 0x54, 0x31, 0x13, 0x30, 0x11, 0x06, 0x03, 0x55,
+            0x04, 0x08, 0x0c, 0x0a, 0x53, 0x6f, 0x6d, 0x65, 0x2d, 0x53, 0x74, 0x61, 0x74, 0x65,
+            0x31, 0x0f, 0x30, 0x0d, 0x06, 0x03, 0x55, 0x04, 0x0a, 0x0c, 0x06, 0x47, 0x6f, 0x6f,
+            0x67, 0x6c, 0x65, 0x30, 0x59, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48, 0xce, 0x3d,
+            0x02, 0x01, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03, 0x01, 0x07, 0x03, 0x42,
+            0x00, 0x04, 0x1e, 0xac, 0x0c, 0xe0, 0x0d, 0xc5, 0x25, 0x84, 0x1b, 0xd2, 0x77, 0x2d,
+            0xe7, 0xba, 0xf1, 0xde, 0xa7, 0xf6, 0x39, 0x7f, 0x38, 0x91, 0xbf, 0xa4, 0x58, 0xf5,
+            0x62, 0x6b, 0xce, 0x06, 0xcf, 0xb9, 0x73, 0x91, 0x0d, 0x8a, 0x60, 0xa0, 0xc6, 0xa2,
+            0x22, 0xe6, 0x51, 0x2e, 0x58, 0xd6, 0x43, 0x02, 0x80, 0x43, 0x44, 0x29, 0x38, 0x9a,
+            0x99, 0xf3, 0xa4, 0xdd, 0xd0, 0xb4, 0x6f, 0x8b, 0x44, 0x2d, 0xa3, 0x53, 0x30, 0x51,
+            0x30, 0x1d, 0x06, 0x03, 0x55, 0x1d, 0x0e, 0x04, 0x16, 0x04, 0x14, 0xdb, 0x13, 0x68,
+            0xe0, 0x0e, 0x47, 0x10, 0xf8, 0xcb, 0x88, 0x83, 0xfe, 0x42, 0x3c, 0xd9, 0x3f, 0x1a,
+            0x33, 0xe9, 0xaa, 0x30, 0x1f, 0x06, 0x03, 0x55, 0x1d, 0x23, 0x04, 0x18, 0x30, 0x16,
+            0x80, 0x14, 0xdb, 0x13, 0x68, 0xe0, 0x0e, 0x47, 0x10, 0xf8, 0xcb, 0x88, 0x83, 0xfe,
+            0x42, 0x3c, 0xd9, 0x3f, 0x1a, 0x33, 0xe9, 0xaa, 0x30, 0x0f, 0x06, 0x03, 0x55, 0x1d,
+            0x13, 0x01, 0x01, 0xff, 0x04, 0x05, 0x30, 0x03, 0x01, 0x01, 0xff, 0x30, 0x0a, 0x06,
+            0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x04, 0x03, 0x02, 0x03, 0x48, 0x00, 0x30, 0x45,
+            0x02, 0x20, 0x10, 0xdf, 0x40, 0xc3, 0x20, 0x54, 0x36, 0xb5, 0xc9, 0x3c, 0x70, 0xe3,
+            0x55, 0x37, 0xd2, 0x04, 0x51, 0xeb, 0x0f, 0x18, 0x83, 0xd0, 0x58, 0xa1, 0x08, 0x77,
+            0x8d, 0x4d, 0xa4, 0x20, 0xee, 0x33, 0x02, 0x21, 0x00, 0x8d, 0xe3, 0xa6, 0x6c, 0x0d,
+            0x86, 0x25, 0xdc, 0x59, 0x0d, 0x21, 0x43, 0x22, 0x3a, 0xb9, 0xa1, 0x73, 0x28, 0xc9,
+            0x16, 0x9e, 0x91, 0x15, 0xc4, 0xc3, 0xd7, 0xeb, 0xe5, 0xce, 0xdc, 0x1c, 0x1b,
+        ]
+    }
+
+    // Generate a fake COSE_Mac0 with a key that's just `byte` repeated
+    fn generate_maced_pubkey(byte: u8) -> Vec<u8> {
+        vec![
+            0x84, 0x43, 0xA1, 0x01, 0x05, 0xA0, 0x58, 0x4D, 0xA5, 0x01, 0x02, 0x03, 0x26, 0x20,
+            0x01, 0x21, 0x58, 0x20, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte,
+            byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte,
+            byte, byte, byte, byte, byte, byte, byte, byte, 0x22, 0x58, 0x20, byte, byte, byte,
+            byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte,
+            byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte,
+            byte, 0x58, 0x20, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte,
+            byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte, byte,
+            byte, byte, byte, byte, byte, byte, byte,
+        ]
+    }
 
     #[test]
     fn test_parse_cose_mac0_for_coords_raw_bytes() -> Result<()> {
@@ -653,4 +834,178 @@
         assert!(extracted_payload.is_err());
         Ok(())
     }
+
+    #[test]
+    #[ignore] // b/215746308
+    fn test_get_attestation_key_no_keys_provisioned() {
+        let mut db = crate::database::tests::new_test_db().unwrap();
+        let mock_rpc = Box::<MockRemotelyProvisionedComponent>::default();
+        mock_rpc.0.lock().unwrap().hw_info.uniqueId = Some(String::from("mallory"));
+
+        let mut service: RemoteProvisioningService = Default::default();
+        service
+            .device_by_sec_level
+            .insert(SecurityLevel::TRUSTED_ENVIRONMENT, Strong::new(mock_rpc));
+
+        assert_eq!(
+            service
+                .get_attestation_key(&mut db, 0, "mallory")
+                .unwrap_err()
+                .downcast::<error::Error>()
+                .unwrap(),
+            error::Error::Rc(ResponseCode::OUT_OF_KEYS)
+        );
+    }
+
+    #[test]
+    #[ignore] // b/215746308
+    fn test_get_attestation_key() {
+        let mut db = crate::database::tests::new_test_db().unwrap();
+        let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
+        let irpc_id = "paul";
+        let caller_uid = 0;
+
+        let mock_rpc = Box::<MockRemotelyProvisionedComponent>::default();
+        let mock_values = mock_rpc.0.clone();
+        let mut service: RemoteProvisioningService = Default::default();
+        service.device_by_sec_level.insert(sec_level, Strong::new(mock_rpc));
+
+        mock_values.lock().unwrap().hw_info.uniqueId = Some(String::from(irpc_id));
+        mock_values.lock().unwrap().private_key = vec![8, 6, 7, 5, 3, 0, 9];
+        mock_values.lock().unwrap().maced_public_key = generate_maced_pubkey(0x11);
+        service.generate_key_pair(&mut db, true, sec_level).unwrap();
+
+        let public_key = RemoteProvisioningService::parse_cose_mac0_for_coords(
+            mock_values.lock().unwrap().maced_public_key.as_slice(),
+        )
+        .unwrap();
+        let batch_cert = get_fake_cert();
+        let certs = &[5, 6, 7, 8];
+        assert!(service
+            .provision_cert_chain(
+                &mut db,
+                public_key.as_slice(),
+                batch_cert.as_slice(),
+                certs,
+                0,
+                sec_level
+            )
+            .is_ok());
+
+        // ensure we got the key we expected
+        let first_key = service
+            .get_attestation_key(&mut db, caller_uid, irpc_id)
+            .context("get first key")
+            .unwrap();
+        assert_eq!(first_key.keyBlob, mock_values.lock().unwrap().private_key);
+        assert_eq!(first_key.encodedCertChain, certs);
+
+        // ensure that multiple calls get the same key
+        assert_eq!(
+            first_key,
+            service
+                .get_attestation_key(&mut db, caller_uid, irpc_id)
+                .context("get second key")
+                .unwrap()
+        );
+
+        // no more keys for new clients
+        assert_eq!(
+            service
+                .get_attestation_key(&mut db, caller_uid + 1, irpc_id)
+                .unwrap_err()
+                .downcast::<error::Error>()
+                .unwrap(),
+            error::Error::Rc(ResponseCode::OUT_OF_KEYS)
+        );
+    }
+
+    #[test]
+    #[ignore] // b/215746308
+    fn test_get_attestation_key_gets_different_key_for_different_client() {
+        let mut db = crate::database::tests::new_test_db().unwrap();
+        let sec_level = SecurityLevel::TRUSTED_ENVIRONMENT;
+        let irpc_id = "ringo";
+        let first_caller = 0;
+        let second_caller = first_caller + 1;
+
+        let mock_rpc = Box::<MockRemotelyProvisionedComponent>::default();
+        let mock_values = mock_rpc.0.clone();
+        let mut service: RemoteProvisioningService = Default::default();
+        service.device_by_sec_level.insert(sec_level, Strong::new(mock_rpc));
+
+        // generate two distinct keys and provision them with certs
+        mock_values.lock().unwrap().hw_info.uniqueId = Some(String::from(irpc_id));
+        mock_values.lock().unwrap().private_key = vec![3, 1, 4, 1, 5];
+        mock_values.lock().unwrap().maced_public_key = generate_maced_pubkey(0x11);
+        assert!(service.generate_key_pair(&mut db, true, sec_level).is_ok());
+        let public_key = RemoteProvisioningService::parse_cose_mac0_for_coords(
+            mock_values.lock().unwrap().maced_public_key.as_slice(),
+        )
+        .unwrap();
+        assert!(service
+            .provision_cert_chain(
+                &mut db,
+                public_key.as_slice(),
+                get_fake_cert().as_slice(),
+                &[1],
+                0,
+                sec_level
+            )
+            .is_ok());
+
+        mock_values.lock().unwrap().hw_info.uniqueId = Some(String::from(irpc_id));
+        mock_values.lock().unwrap().private_key = vec![9, 0, 2, 1, 0];
+        mock_values.lock().unwrap().maced_public_key = generate_maced_pubkey(0x22);
+        assert!(service.generate_key_pair(&mut db, true, sec_level).is_ok());
+        let public_key = RemoteProvisioningService::parse_cose_mac0_for_coords(
+            mock_values.lock().unwrap().maced_public_key.as_slice(),
+        )
+        .unwrap();
+        assert!(service
+            .provision_cert_chain(
+                &mut db,
+                public_key.as_slice(),
+                get_fake_cert().as_slice(),
+                &[2],
+                0,
+                sec_level
+            )
+            .is_ok());
+
+        // make sure each caller gets a distinct key
+        assert_ne!(
+            service
+                .get_attestation_key(&mut db, first_caller, irpc_id)
+                .context("get first key")
+                .unwrap(),
+            service
+                .get_attestation_key(&mut db, second_caller, irpc_id)
+                .context("get second key")
+                .unwrap()
+        );
+
+        // repeated calls should return the same key for a given caller
+        assert_eq!(
+            service
+                .get_attestation_key(&mut db, first_caller, irpc_id)
+                .context("first caller a")
+                .unwrap(),
+            service
+                .get_attestation_key(&mut db, first_caller, irpc_id)
+                .context("first caller b")
+                .unwrap(),
+        );
+
+        assert_eq!(
+            service
+                .get_attestation_key(&mut db, second_caller, irpc_id)
+                .context("second caller a")
+                .unwrap(),
+            service
+                .get_attestation_key(&mut db, second_caller, irpc_id)
+                .context("second caller b")
+                .unwrap()
+        );
+    }
 }
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 9334930..83f0bee 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -160,6 +160,8 @@
                     let mut db = db.borrow_mut();
 
                     let (key_blob, mut blob_metadata) = SUPER_KEY
+                        .read()
+                        .unwrap()
                         .handle_super_encryption_on_key_init(
                             &mut db,
                             &LEGACY_MIGRATOR,
@@ -303,6 +305,8 @@
             .context("In create_operation.")?;
 
         let km_blob = SUPER_KEY
+            .read()
+            .unwrap()
             .unwrap_key_if_required(&blob_metadata, km_blob)
             .context("In create_operation. Failed to handle super encryption.")?;
 
@@ -736,8 +740,11 @@
             .ok_or_else(error::Error::sys)
             .context("No km_blob after successfully loading key. This should never happen.")?;
 
-        let wrapping_key_blob =
-            SUPER_KEY.unwrap_key_if_required(&wrapping_blob_metadata, &wrapping_key_blob).context(
+        let wrapping_key_blob = SUPER_KEY
+            .read()
+            .unwrap()
+            .unwrap_key_if_required(&wrapping_blob_metadata, &wrapping_key_blob)
+            .context(
                 "In import_wrapped_key. Failed to handle super encryption for wrapping key.",
             )?;
 
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index ca5e593..6862011 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -28,7 +28,6 @@
     legacy_blob::LegacyBlobLoader,
     legacy_migrator::LegacyMigrator,
     raw_device::KeyMintDevice,
-    try_insert::TryInsert,
     utils::watchdog as wd,
     utils::AID_KEYSTORE,
 };
@@ -50,7 +49,7 @@
 use std::{
     collections::HashMap,
     sync::Arc,
-    sync::{Mutex, Weak},
+    sync::{Mutex, RwLock, Weak},
 };
 use std::{convert::TryFrom, ops::Deref};
 
@@ -76,7 +75,7 @@
 /// different purpose, distinguished by alias. Each is associated with a static
 /// constant of this type.
 pub struct SuperKeyType<'a> {
-    /// Alias used to look the key up in the `persistent.keyentry` table.
+    /// Alias used to look up the key in the `persistent.keyentry` table.
     pub alias: &'a str,
     /// Encryption algorithm
     pub algorithm: SuperEncryptionAlgorithm,
@@ -256,7 +255,7 @@
 struct SkmState {
     user_keys: HashMap<UserId, UserSuperKeys>,
     key_index: HashMap<i64, Weak<SuperKey>>,
-    boot_level_key_cache: Option<BootLevelKeyCache>,
+    boot_level_key_cache: Option<Mutex<BootLevelKeyCache>>,
 }
 
 impl SkmState {
@@ -275,24 +274,24 @@
 
 #[derive(Default)]
 pub struct SuperKeyManager {
-    data: Mutex<SkmState>,
+    data: SkmState,
 }
 
 impl SuperKeyManager {
-    pub fn set_up_boot_level_cache(self: &Arc<Self>, db: &mut KeystoreDB) -> Result<()> {
-        let mut data = self.data.lock().unwrap();
-        if data.boot_level_key_cache.is_some() {
+    pub fn set_up_boot_level_cache(skm: &Arc<RwLock<Self>>, db: &mut KeystoreDB) -> Result<()> {
+        let mut skm_guard = skm.write().unwrap();
+        if skm_guard.data.boot_level_key_cache.is_some() {
             log::info!("In set_up_boot_level_cache: called for a second time");
             return Ok(());
         }
         let level_zero_key = get_level_zero_key(db)
             .context("In set_up_boot_level_cache: get_level_zero_key failed")?;
-        data.boot_level_key_cache = Some(BootLevelKeyCache::new(level_zero_key));
+        skm_guard.data.boot_level_key_cache =
+            Some(Mutex::new(BootLevelKeyCache::new(level_zero_key)));
         log::info!("Starting boot level watcher.");
-        let clone = self.clone();
+        let clone = skm.clone();
         std::thread::spawn(move || {
-            clone
-                .watch_boot_level()
+            Self::watch_boot_level(clone)
                 .unwrap_or_else(|e| log::error!("watch_boot_level failed:\n{:?}", e));
         });
         Ok(())
@@ -300,32 +299,40 @@
 
     /// Watch the `keystore.boot_level` system property, and keep boot level up to date.
     /// Blocks waiting for system property changes, so must be run in its own thread.
-    fn watch_boot_level(&self) -> Result<()> {
+    fn watch_boot_level(skm: Arc<RwLock<Self>>) -> Result<()> {
         let mut w = PropertyWatcher::new("keystore.boot_level")
             .context("In watch_boot_level: PropertyWatcher::new failed")?;
         loop {
             let level = w
                 .read(|_n, v| v.parse::<usize>().map_err(std::convert::Into::into))
                 .context("In watch_boot_level: read of property failed")?;
-            // watch_boot_level should only be called once data.boot_level_key_cache is Some,
-            // so it's safe to unwrap in the branches below.
-            if level < MAX_MAX_BOOT_LEVEL {
-                log::info!("Read keystore.boot_level value {}", level);
-                let mut data = self.data.lock().unwrap();
-                data.boot_level_key_cache
+
+            // This scope limits the skm_guard life, so we don't hold the skm_guard while
+            // waiting.
+            {
+                let mut skm_guard = skm.write().unwrap();
+                let boot_level_key_cache = skm_guard
+                    .data
+                    .boot_level_key_cache
                     .as_mut()
-                    .unwrap()
-                    .advance_boot_level(level)
-                    .context("In watch_boot_level: advance_boot_level failed")?;
-            } else {
-                log::info!(
-                    "keystore.boot_level {} hits maximum {}, finishing.",
-                    level,
-                    MAX_MAX_BOOT_LEVEL
-                );
-                let mut data = self.data.lock().unwrap();
-                data.boot_level_key_cache.as_mut().unwrap().finish();
-                break;
+                    .ok_or_else(Error::sys)
+                    .context("In watch_boot_level: Boot level cache not initialized")?
+                    .get_mut()
+                    .unwrap();
+                if level < MAX_MAX_BOOT_LEVEL {
+                    log::info!("Read keystore.boot_level value {}", level);
+                    boot_level_key_cache
+                        .advance_boot_level(level)
+                        .context("In watch_boot_level: advance_boot_level failed")?;
+                } else {
+                    log::info!(
+                        "keystore.boot_level {} hits maximum {}, finishing.",
+                        level,
+                        MAX_MAX_BOOT_LEVEL
+                    );
+                    boot_level_key_cache.finish();
+                    break;
+                }
             }
             w.wait().context("In watch_boot_level: property wait failed")?;
         }
@@ -334,34 +341,37 @@
 
     pub fn level_accessible(&self, boot_level: i32) -> bool {
         self.data
-            .lock()
-            .unwrap()
             .boot_level_key_cache
             .as_ref()
-            .map_or(false, |c| c.level_accessible(boot_level as usize))
+            .map_or(false, |c| c.lock().unwrap().level_accessible(boot_level as usize))
     }
 
-    pub fn forget_all_keys_for_user(&self, user: UserId) {
-        let mut data = self.data.lock().unwrap();
-        data.user_keys.remove(&user);
+    pub fn forget_all_keys_for_user(&mut self, user: UserId) {
+        self.data.user_keys.remove(&user);
     }
 
-    fn install_per_boot_key_for_user(&self, user: UserId, super_key: Arc<SuperKey>) -> Result<()> {
-        let mut data = self.data.lock().unwrap();
-        data.add_key_to_key_index(&super_key)
+    fn install_per_boot_key_for_user(
+        &mut self,
+        user: UserId,
+        super_key: Arc<SuperKey>,
+    ) -> Result<()> {
+        self.data
+            .add_key_to_key_index(&super_key)
             .context("In install_per_boot_key_for_user: add_key_to_key_index failed")?;
-        data.user_keys.entry(user).or_default().per_boot = Some(super_key);
+        self.data.user_keys.entry(user).or_default().per_boot = Some(super_key);
         Ok(())
     }
 
     fn lookup_key(&self, key_id: &SuperKeyIdentifier) -> Result<Option<Arc<SuperKey>>> {
-        let mut data = self.data.lock().unwrap();
         Ok(match key_id {
-            SuperKeyIdentifier::DatabaseId(id) => data.key_index.get(id).and_then(|k| k.upgrade()),
-            SuperKeyIdentifier::BootLevel(level) => data
+            SuperKeyIdentifier::DatabaseId(id) => {
+                self.data.key_index.get(id).and_then(|k| k.upgrade())
+            }
+            SuperKeyIdentifier::BootLevel(level) => self
+                .data
                 .boot_level_key_cache
-                .as_mut()
-                .map(|b| b.aes_key(*level as usize))
+                .as_ref()
+                .map(|b| b.lock().unwrap().aes_key(*level as usize))
                 .transpose()
                 .context("In lookup_key: aes_key failed")?
                 .flatten()
@@ -377,8 +387,7 @@
     }
 
     pub fn get_per_boot_key_by_user_id(&self, user_id: UserId) -> Option<Arc<SuperKey>> {
-        let data = self.data.lock().unwrap();
-        data.user_keys.get(&user_id).and_then(|e| e.per_boot.as_ref().cloned())
+        self.data.user_keys.get(&user_id).and_then(|e| e.per_boot.as_ref().cloned())
     }
 
     /// This function unlocks the super keys for a given user.
@@ -386,7 +395,7 @@
     /// super key cache. If there is no such key a new key is created, encrypted with
     /// a key derived from the given password and stored in the database.
     pub fn unlock_user_key(
-        &self,
+        &mut self,
         db: &mut KeystoreDB,
         user: UserId,
         pw: &Password,
@@ -493,7 +502,10 @@
     }
 
     /// Checks if user has setup LSKF, even when super key cache is empty for the user.
-    pub fn super_key_exists_in_db_for_user(
+    /// The reference to self is unused but it is required to prevent calling this function
+    /// concurrently with skm state database changes.
+    fn super_key_exists_in_db_for_user(
+        &self,
         db: &mut KeystoreDB,
         legacy_migrator: &LegacyMigrator,
         user_id: UserId,
@@ -515,7 +527,7 @@
     /// 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,
+        &mut self,
         db: &mut KeystoreDB,
         legacy_migrator: &LegacyMigrator,
         user_id: UserId,
@@ -544,24 +556,23 @@
     /// and return LskfUnlocked state.
     /// If the password is not provided, return Uninitialized state.
     pub fn check_and_initialize_super_key(
-        &self,
+        &mut self,
         db: &mut KeystoreDB,
         legacy_migrator: &LegacyMigrator,
         user_id: UserId,
         pw: Option<&Password>,
     ) -> Result<UserState> {
-        let super_key_exists_in_db =
-            Self::super_key_exists_in_db_for_user(db, legacy_migrator, user_id).context(
-                "In check_and_initialize_super_key. Failed to check if super key exists.",
-            )?;
+        let super_key_exists_in_db = self
+            .super_key_exists_in_db_for_user(db, legacy_migrator, 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 if let Some(pw) = pw {
-            //generate a new super key.
+            // Generate a new super key.
             let super_key = generate_aes256_key()
                 .context("In check_and_initialize_super_key: Failed to generate AES 256 key.")?;
-            //derive an AES256 key from the password and re-encrypt the super key
-            //before we insert it in the database.
+            // Derive an AES256 key from the password and re-encrypt the super key
+            // before we insert it in the database.
             let (encrypted_super_key, blob_metadata) = Self::encrypt_with_password(&super_key, pw)
                 .context("In check_and_initialize_super_key.")?;
 
@@ -589,9 +600,9 @@
         }
     }
 
-    //helper function to populate super key cache from the super key blob loaded from the database
+    // Helper function to populate super key cache from the super key blob loaded from the database.
     fn populate_cache_from_super_key_blob(
-        &self,
+        &mut self,
         user_id: UserId,
         algorithm: SuperEncryptionAlgorithm,
         entry: KeyEntry,
@@ -605,7 +616,7 @@
         Ok(super_key)
     }
 
-    /// Extracts super key from the entry loaded from the database
+    /// Extracts super key from the entry loaded from the database.
     pub fn extract_super_key_from_key_entry(
         algorithm: SuperEncryptionAlgorithm,
         entry: KeyEntry,
@@ -620,7 +631,7 @@
                 metadata.aead_tag(),
             ) {
                 (Some(&EncryptedBy::Password), Some(salt), Some(iv), Some(tag)) => {
-                    // Note that password encryption is AES no matter the value of algorithm
+                    // Note that password encryption is AES no matter the value of algorithm.
                     let key = pw.derive_key(Some(salt), AES_256_KEY_LENGTH).context(
                         "In extract_super_key_from_key_entry: Failed to generate key from password.",
                     )?;
@@ -684,7 +695,8 @@
         user_id: UserId,
         key_blob: &[u8],
     ) -> Result<(Vec<u8>, BlobMetaData)> {
-        match UserState::get(db, legacy_migrator, self, user_id)
+        match self
+            .get_user_state(db, legacy_migrator, user_id)
             .context("In super_encrypt. Failed to get user state.")?
         {
             UserState::LskfUnlocked(super_key) => {
@@ -699,9 +711,9 @@
         }
     }
 
-    //Helper function to encrypt a key with the given super key. Callers should select which super
-    //key to be used. This is called when a key is super encrypted at its creation as well as at its
-    //upgrade.
+    // Helper function to encrypt a key with the given super key. Callers should select which super
+    // key to be used. This is called when a key is super encrypted at its creation as well as at
+    // its upgrade.
     fn encrypt_with_aes_super_key(
         key_blob: &[u8],
         super_key: &SuperKey,
@@ -741,9 +753,13 @@
                     "Failed to super encrypt with LskfBound key."
                 )),
             SuperEncryptionType::ScreenLockBound => {
-                let mut data = self.data.lock().unwrap();
-                let entry = data.user_keys.entry(user_id).or_default();
-                if let Some(super_key) = entry.screen_lock_bound.as_ref() {
+                let entry = self
+                    .data
+                    .user_keys
+                    .get(&user_id)
+                    .map(|e| e.screen_lock_bound.as_ref())
+                    .flatten();
+                if let Some(super_key) = entry {
                     Self::encrypt_with_aes_super_key(key_blob, super_key).context(concat!(
                         "In handle_super_encryption_on_key_init. ",
                         "Failed to encrypt with ScreenLockBound key."
@@ -813,6 +829,7 @@
     /// When this is called, the caller must hold the lock on the SuperKeyManager.
     /// So it's OK that the check and creation are different DB transactions.
     fn get_or_create_super_key(
+        &mut self,
         db: &mut KeystoreDB,
         user_id: UserId,
         key_type: &SuperKeyType,
@@ -847,8 +864,8 @@
                     )
                 }
             };
-            //derive an AES256 key from the password and re-encrypt the super key
-            //before we insert it in the database.
+            // Derive an AES256 key from the password and re-encrypt the super key
+            // before we insert it in the database.
             let (encrypted_super_key, blob_metadata) =
                 Self::encrypt_with_password(&super_key, password)
                     .context("In get_or_create_super_key.")?;
@@ -876,52 +893,64 @@
 
     /// Decrypt the screen-lock bound keys for this user using the password and store in memory.
     pub fn unlock_screen_lock_bound_key(
-        &self,
+        &mut self,
         db: &mut KeystoreDB,
         user_id: UserId,
         password: &Password,
     ) -> Result<()> {
-        let mut data = self.data.lock().unwrap();
-        let entry = data.user_keys.entry(user_id).or_default();
-        let aes = entry
-            .screen_lock_bound
-            .get_or_try_to_insert_with(|| {
-                Self::get_or_create_super_key(
-                    db,
-                    user_id,
-                    &USER_SCREEN_LOCK_BOUND_KEY,
-                    password,
-                    None,
-                )
-            })?
-            .clone();
-        let ecdh = entry
-            .screen_lock_bound_private
-            .get_or_try_to_insert_with(|| {
-                Self::get_or_create_super_key(
-                    db,
-                    user_id,
-                    &USER_SCREEN_LOCK_BOUND_P521_KEY,
-                    password,
-                    Some(aes.clone()),
-                )
-            })?
-            .clone();
-        data.add_key_to_key_index(&aes)?;
-        data.add_key_to_key_index(&ecdh)?;
+        let (screen_lock_bound, screen_lock_bound_private) = self
+            .data
+            .user_keys
+            .get(&user_id)
+            .map(|e| (e.screen_lock_bound.clone(), e.screen_lock_bound_private.clone()))
+            .unwrap_or((None, None));
+
+        if screen_lock_bound.is_some() && screen_lock_bound_private.is_some() {
+            // Already unlocked.
+            return Ok(());
+        }
+
+        let aes = if let Some(screen_lock_bound) = screen_lock_bound {
+            // This is weird. If this point is reached only one of the screen locked keys was
+            // initialized. This should never happen.
+            screen_lock_bound
+        } else {
+            self.get_or_create_super_key(db, user_id, &USER_SCREEN_LOCK_BOUND_KEY, password, None)
+                .context("In unlock_screen_lock_bound_key: Trying to get or create symmetric key.")?
+        };
+
+        let ecdh = if let Some(screen_lock_bound_private) = screen_lock_bound_private {
+            // This is weird. If this point is reached only one of the screen locked keys was
+            // initialized. This should never happen.
+            screen_lock_bound_private
+        } else {
+            self.get_or_create_super_key(
+                db,
+                user_id,
+                &USER_SCREEN_LOCK_BOUND_P521_KEY,
+                password,
+                Some(aes.clone()),
+            )
+            .context("In unlock_screen_lock_bound_key: Trying to get or create asymmetric key.")?
+        };
+
+        self.data.add_key_to_key_index(&aes)?;
+        self.data.add_key_to_key_index(&ecdh)?;
+        let entry = self.data.user_keys.entry(user_id).or_default();
+        entry.screen_lock_bound = Some(aes);
+        entry.screen_lock_bound_private = Some(ecdh);
         Ok(())
     }
 
     /// Wipe the screen-lock bound keys for this user from memory.
     pub fn lock_screen_lock_bound_key(
-        &self,
+        &mut self,
         db: &mut KeystoreDB,
         user_id: UserId,
         unlocking_sids: &[i64],
     ) {
         log::info!("Locking screen bound for user {} sids {:?}", user_id, unlocking_sids);
-        let mut data = self.data.lock().unwrap();
-        let mut entry = data.user_keys.entry(user_id).or_default();
+        let mut entry = self.data.user_keys.entry(user_id).or_default();
         if !unlocking_sids.is_empty() {
             if let (Some(aes), Some(ecdh)) = (
                 entry.screen_lock_bound.as_ref().cloned(),
@@ -993,12 +1022,11 @@
     /// User has unlocked, not using a password. See if any of our stored auth tokens can be used
     /// to unlock the keys protecting UNLOCKED_DEVICE_REQUIRED keys.
     pub fn try_unlock_user_with_biometric(
-        &self,
+        &mut self,
         db: &mut KeystoreDB,
         user_id: UserId,
     ) -> Result<()> {
-        let mut data = self.data.lock().unwrap();
-        let mut entry = data.user_keys.entry(user_id).or_default();
+        let mut entry = self.data.user_keys.entry(user_id).or_default();
         if let Some(biometric) = entry.biometric_unlock.as_ref() {
             let (key_id_guard, key_entry) = db
                 .load_key_entry(
@@ -1038,8 +1066,8 @@
                         Ok((slb, slbp)) => {
                             entry.screen_lock_bound = Some(slb.clone());
                             entry.screen_lock_bound_private = Some(slbp.clone());
-                            data.add_key_to_key_index(&slb)?;
-                            data.add_key_to_key_index(&slbp)?;
+                            self.data.add_key_to_key_index(&slb)?;
+                            self.data.add_key_to_key_index(&slbp)?;
                             log::info!(concat!(
                                 "In try_unlock_user_with_biometric: ",
                                 "Successfully unlocked with biometric"
@@ -1055,6 +1083,122 @@
         }
         Ok(())
     }
+
+    /// Returns the keystore locked state of the given user. It requires the thread local
+    /// keystore database and a reference to the legacy migrator because it may need to
+    /// migrate the super key from the legacy blob database to the keystore database.
+    pub fn get_user_state(
+        &self,
+        db: &mut KeystoreDB,
+        legacy_migrator: &LegacyMigrator,
+        user_id: UserId,
+    ) -> Result<UserState> {
+        match self.get_per_boot_key_by_user_id(user_id) {
+            Some(super_key) => Ok(UserState::LskfUnlocked(super_key)),
+            None => {
+                // Check if a super key exists in the database or legacy database.
+                // If so, return locked user state.
+                if self
+                    .super_key_exists_in_db_for_user(db, legacy_migrator, user_id)
+                    .context("In get_user_state.")?
+                {
+                    Ok(UserState::LskfLocked)
+                } else {
+                    Ok(UserState::Uninitialized)
+                }
+            }
+        }
+    }
+
+    /// If the given user is unlocked:
+    /// * and `password` is None, the user is reset, all authentication bound keys are deleted and
+    ///   `Ok(UserState::Uninitialized)` is returned.
+    /// * and `password` is Some, `Ok(UserState::LskfUnlocked)` is returned.
+    /// If the given user is locked:
+    /// * and the user was initialized before, `Ok(UserState::Locked)` is returned.
+    /// * and the user was not initialized before:
+    ///   * and `password` is None, `Ok(Uninitialized)` is returned.
+    ///   * and `password` is Some, super keys are generated and `Ok(UserState::LskfUnlocked)` is
+    ///     returned.
+    pub fn reset_or_init_user_and_get_user_state(
+        &mut self,
+        db: &mut KeystoreDB,
+        legacy_migrator: &LegacyMigrator,
+        user_id: UserId,
+        password: Option<&Password>,
+    ) -> Result<UserState> {
+        match self.get_per_boot_key_by_user_id(user_id) {
+            Some(_) if password.is_none() => {
+                // Transitioning to swiping, delete only the super key in database and cache,
+                // and super-encrypted keys in database (and in KM).
+                self.reset_user(db, legacy_migrator, user_id, true).context(
+                    "In reset_or_init_user_and_get_user_state: Trying to delete keys from the db.",
+                )?;
+                // Lskf is now removed in Keystore.
+                Ok(UserState::Uninitialized)
+            }
+            Some(super_key) => {
+                // Keystore won't be notified when changing to a new password when LSKF is
+                // already setup. Therefore, ideally this path wouldn't be reached.
+                Ok(UserState::LskfUnlocked(super_key))
+            }
+            None => {
+                // Check if a super key exists in the database or legacy database.
+                // If so, return LskfLocked state.
+                // Otherwise, i) if the password is provided, initialize the super key and return
+                // LskfUnlocked state ii) if password is not provided, return Uninitialized state.
+                self.check_and_initialize_super_key(db, legacy_migrator, user_id, password)
+            }
+        }
+    }
+
+    /// Unlocks the given user with the given password. If the key was already unlocked or unlocking
+    /// was successful, `Ok(UserState::LskfUnlocked)` is returned.
+    /// If the user was never initialized `Ok(UserState::Uninitialized)` is returned.
+    pub fn unlock_and_get_user_state(
+        &mut self,
+        db: &mut KeystoreDB,
+        legacy_migrator: &LegacyMigrator,
+        user_id: UserId,
+        password: &Password,
+    ) -> Result<UserState> {
+        match self.get_per_boot_key_by_user_id(user_id) {
+            Some(super_key) => {
+                log::info!("In unlock_and_get_user_state. 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.
+                self.check_and_unlock_super_key(db, legacy_migrator, user_id, password)
+                    .context("In unlock_and_get_user_state. 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.
+    pub fn reset_user(
+        &mut self,
+        db: &mut KeystoreDB,
+        legacy_migrator: &LegacyMigrator,
+        user_id: UserId,
+        keep_non_super_encrypted_keys: bool,
+    ) -> Result<()> {
+        // Mark keys created on behalf of the user as unreferenced.
+        legacy_migrator
+            .bulk_delete_user(user_id, keep_non_super_encrypted_keys)
+            .context("In reset_user: Trying to delete legacy keys.")?;
+        db.unbind_keys_for_user(user_id, keep_non_super_encrypted_keys)
+            .context("In reset user. Error in unbinding keys.")?;
+
+        // Delete super key in cache, if exists.
+        self.forget_all_keys_for_user(user_id);
+        Ok(())
+    }
 }
 
 /// This enum represents different states of the user's life cycle in the device.
@@ -1072,110 +1216,6 @@
     Uninitialized,
 }
 
-impl UserState {
-    pub fn get(
-        db: &mut KeystoreDB,
-        legacy_migrator: &LegacyMigrator,
-        skm: &SuperKeyManager,
-        user_id: UserId,
-    ) -> Result<UserState> {
-        match skm.get_per_boot_key_by_user_id(user_id) {
-            Some(super_key) => Ok(UserState::LskfUnlocked(super_key)),
-            None => {
-                //Check if a super key exists in the database or legacy database.
-                //If so, return locked user state.
-                if SuperKeyManager::super_key_exists_in_db_for_user(db, legacy_migrator, user_id)
-                    .context("In get.")?
-                {
-                    Ok(UserState::LskfLocked)
-                } else {
-                    Ok(UserState::Uninitialized)
-                }
-            }
-        }
-    }
-
-    /// Queries user state when serving password change requests.
-    pub fn get_with_password_changed(
-        db: &mut KeystoreDB,
-        legacy_migrator: &LegacyMigrator,
-        skm: &SuperKeyManager,
-        user_id: UserId,
-        password: Option<&Password>,
-    ) -> Result<UserState> {
-        match skm.get_per_boot_key_by_user_id(user_id) {
-            Some(super_key) => {
-                if password.is_none() {
-                    //transitioning to swiping, delete only the super key in database and cache, and
-                    //super-encrypted keys in database (and in KM)
-                    Self::reset_user(db, skm, legacy_migrator, user_id, true).context(
-                        "In get_with_password_changed: Trying to delete keys from the db.",
-                    )?;
-                    //Lskf is now removed in Keystore
-                    Ok(UserState::Uninitialized)
-                } else {
-                    //Keystore won't be notified when changing to a new password when LSKF is
-                    //already setup. Therefore, ideally this path wouldn't be reached.
-                    Ok(UserState::LskfUnlocked(super_key))
-                }
-            }
-            None => {
-                //Check if a super key exists in the database or legacy database.
-                //If so, return LskfLocked state.
-                //Otherwise, i) if the password is provided, initialize the super key and return
-                //LskfUnlocked state ii) if password is not provided, return Uninitialized state.
-                skm.check_and_initialize_super_key(db, legacy_migrator, user_id, password)
-            }
-        }
-    }
-
-    /// Queries user state when serving password unlock requests.
-    pub fn get_with_password_unlock(
-        db: &mut KeystoreDB,
-        legacy_migrator: &LegacyMigrator,
-        skm: &SuperKeyManager,
-        user_id: UserId,
-        password: &Password,
-    ) -> 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, legacy_migrator, 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.
-    pub fn reset_user(
-        db: &mut KeystoreDB,
-        skm: &SuperKeyManager,
-        legacy_migrator: &LegacyMigrator,
-        user_id: UserId,
-        keep_non_super_encrypted_keys: bool,
-    ) -> Result<()> {
-        // mark keys created on behalf of the user as unreferenced.
-        legacy_migrator
-            .bulk_delete_user(user_id, keep_non_super_encrypted_keys)
-            .context("In reset_user: Trying to delete legacy keys.")?;
-        db.unbind_keys_for_user(user_id, keep_non_super_encrypted_keys)
-            .context("In reset user. Error in unbinding keys.")?;
-
-        //delete super key in cache, if exists
-        skm.forget_all_keys_for_user(user_id);
-        Ok(())
-    }
-}
-
 /// This enum represents three states a KeyMint Blob can be in, w.r.t super encryption.
 /// `Sensitive` holds the non encrypted key and a reference to its super key.
 /// `NonSensitive` holds a non encrypted key that is never supposed to be encrypted.
diff --git a/keystore2/src/try_insert.rs b/keystore2/src/try_insert.rs
deleted file mode 100644
index 6dd3962..0000000
--- a/keystore2/src/try_insert.rs
+++ /dev/null
@@ -1,100 +0,0 @@
-// Copyright 2021, The Android Open Source Project
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-//! The TryInsert trait adds to Option<T> the method
-//! get_or_try_to_insert_with, which is analogous to
-//! get_or_insert_with, but allows the called function to fail and propagates the failure.
-
-/// The TryInsert trait adds to Option<T> the method
-/// get_or_try_to_insert_with, which is analogous to
-/// get_or_insert_with, but allows the called function to fail and propagates the failure.
-pub trait TryInsert {
-    /// Type of the Ok branch of the Result
-    type Item;
-    /// Inserts a value computed from `f` into the option if it is [`None`],
-    /// then returns a mutable reference to the contained value. If `f`
-    /// returns Err, the Option is unchanged.
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// let mut x = None;
-    /// assert_eq!(x.get_or_try_to_insert_with(Err("oops".to_string())), Err("oops".to_string()))
-    /// {
-    ///     let y: &mut u32 = x.get_or_try_to_insert_with(|| Ok(5))?;
-    ///     assert_eq!(y, &5);
-    ///
-    ///     *y = 7;
-    /// }
-    ///
-    /// assert_eq!(x, Some(7));
-    /// ```
-    fn get_or_try_to_insert_with<E, F: FnOnce() -> Result<Self::Item, E>>(
-        &mut self,
-        f: F,
-    ) -> Result<&mut Self::Item, E>;
-}
-
-impl<T> TryInsert for Option<T> {
-    type Item = T;
-    fn get_or_try_to_insert_with<E, F: FnOnce() -> Result<Self::Item, E>>(
-        &mut self,
-        f: F,
-    ) -> Result<&mut Self::Item, E> {
-        if self.is_none() {
-            *self = Some(f()?);
-        }
-
-        match self {
-            Some(v) => Ok(v),
-            // SAFETY: a `None` variant for `self` would have been replaced by a `Some`
-            // variant in the code above.
-            None => unsafe { std::hint::unreachable_unchecked() },
-        }
-    }
-}
-
-#[cfg(test)]
-mod test {
-    use super::*;
-
-    fn fails() -> Result<i32, String> {
-        Err("fail".to_string())
-    }
-
-    fn succeeds() -> Result<i32, String> {
-        Ok(99)
-    }
-
-    #[test]
-    fn test() {
-        let mut x = None;
-        assert_eq!(x.get_or_try_to_insert_with(fails), Err("fail".to_string()));
-        assert_eq!(x, None);
-        assert_eq!(*x.get_or_try_to_insert_with(succeeds).unwrap(), 99);
-        assert_eq!(x, Some(99));
-        x = Some(42);
-        assert_eq!(*x.get_or_try_to_insert_with(fails).unwrap(), 42);
-        assert_eq!(x, Some(42));
-        assert_eq!(*x.get_or_try_to_insert_with(succeeds).unwrap(), 42);
-        assert_eq!(x, Some(42));
-        *x.get_or_try_to_insert_with(fails).unwrap() = 2;
-        assert_eq!(x, Some(2));
-        *x.get_or_try_to_insert_with(succeeds).unwrap() = 3;
-        assert_eq!(x, Some(3));
-        x = None;
-        *x.get_or_try_to_insert_with(succeeds).unwrap() = 5;
-        assert_eq!(x, Some(5));
-    }
-}