Merge "Keystore 2.0: Implement shared secret negotiation."
diff --git a/OWNERS b/OWNERS
index fca66f8..7d1fd14 100644
--- a/OWNERS
+++ b/OWNERS
@@ -1,4 +1,5 @@
 swillden@google.com
 cbrubaker@google.com
 jdanis@google.com
-kroot@google.com
\ No newline at end of file
+kroot@google.com
+zeuthen@google.com
\ No newline at end of file
diff --git a/identity/main.cpp b/identity/main.cpp
index 08f2219..9add73c 100644
--- a/identity/main.cpp
+++ b/identity/main.cpp
@@ -53,5 +53,9 @@
     CHECK(ret == ::android::OK) << "Couldn't register binder service";
     LOG(ERROR) << "Registered binder service";
 
+    // Credstore is a single-threaded process. So devote the main thread
+    // to handling binder messages.
+    IPCThreadState::self()->joinThreadPool();
+
     return 0;
 }
diff --git a/keystore2/aidl/Android.bp b/keystore2/aidl/Android.bp
index 69ba0b4..183096c 100644
--- a/keystore2/aidl/Android.bp
+++ b/keystore2/aidl/Android.bp
@@ -28,7 +28,8 @@
     unstable: true,
     backend: {
         java: {
-            sdk_version: "module_current",
+            platform_apis: true,
+            srcs_available: true,
         },
         rust: {
             enabled: true,
@@ -46,7 +47,8 @@
     unstable: true,
     backend: {
         java: {
-            sdk_version: "module_current",
+            platform_apis: true,
+            srcs_available: true,
         },
         rust: {
             enabled: true,
@@ -64,6 +66,7 @@
     backend: {
         java: {
             enabled: true,
+            srcs_available: true,
         },
         rust: {
             enabled: true,
@@ -82,7 +85,8 @@
     unstable: true,
     backend: {
         java: {
-            sdk_version: "module_current",
+            platform_apis: true,
+            srcs_available: true,
         },
         rust: {
             enabled: true,
@@ -102,9 +106,8 @@
     unstable: true,
     backend: {
         java: {
-            enabled: true,
-            sdk_version: "module_current",
             platform_apis: true,
+            srcs_available: true,
         },
         ndk: {
             enabled: true,
@@ -124,7 +127,8 @@
     unstable: true,
     backend: {
         java: {
-            sdk_version: "module_current",
+            platform_apis: true,
+            srcs_available: true,
         },
         rust: {
             enabled: true,
@@ -141,7 +145,8 @@
     unstable: true,
     backend: {
         java: {
-            sdk_version: "module_current",
+            platform_apis: true,
+            srcs_available: true,
         },
         rust: {
             enabled: true,
diff --git a/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl b/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl
index f47d7f5..277b9dd 100644
--- a/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl
+++ b/keystore2/aidl/android/security/apc/IConfirmationCallback.aidl
@@ -21,6 +21,7 @@
 /**
  * This callback interface must be implemented by the client to receive the result of the user
  * confirmation.
+ * @hide
  */
 interface IConfirmationCallback {
     /**
diff --git a/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl b/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl
index 26ccf0f..3162224 100644
--- a/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl
+++ b/keystore2/aidl/android/security/apc/IProtectedConfirmation.aidl
@@ -18,6 +18,7 @@
 
 import android.security.apc.IConfirmationCallback;
 
+/** @hide */
 interface IProtectedConfirmation {
 
     /**
diff --git a/keystore2/aidl/android/security/apc/ResponseCode.aidl b/keystore2/aidl/android/security/apc/ResponseCode.aidl
index 7ae3e1c..9a3619f 100644
--- a/keystore2/aidl/android/security/apc/ResponseCode.aidl
+++ b/keystore2/aidl/android/security/apc/ResponseCode.aidl
@@ -19,6 +19,7 @@
 /**
  * Used as service specific exception code by IProtectedConfirmation and as result
  * code by IConfirmationCallback
+ * @hide
  */
 @Backing(type="int")
 enum ResponseCode {
diff --git a/keystore2/aidl/android/security/attestationmanager/ByteArray.aidl b/keystore2/aidl/android/security/attestationmanager/ByteArray.aidl
index a1592ec..dc37b1b 100644
--- a/keystore2/aidl/android/security/attestationmanager/ByteArray.aidl
+++ b/keystore2/aidl/android/security/attestationmanager/ByteArray.aidl
@@ -18,7 +18,6 @@
 
 /**
  * Simple data holder for a byte array, allowing for multidimensional arrays in AIDL.
- *
  * @hide
  */
 parcelable ByteArray {
diff --git a/keystore2/aidl/android/security/attestationmanager/IAttestationManager.aidl b/keystore2/aidl/android/security/attestationmanager/IAttestationManager.aidl
index 85eee57..e77a21e 100644
--- a/keystore2/aidl/android/security/attestationmanager/IAttestationManager.aidl
+++ b/keystore2/aidl/android/security/attestationmanager/IAttestationManager.aidl
@@ -21,7 +21,6 @@
 
 /**
  * Internal interface for performing device attestation.
- *
  * @hide
  */
 interface IAttestationManager {
diff --git a/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl b/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
index 6dc172e..86472eb 100644
--- a/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
+++ b/keystore2/aidl/android/security/authorization/IKeystoreAuthorization.aidl
@@ -23,6 +23,7 @@
 /**
  * IKeystoreAuthorization interface exposes the methods for other system components to
  * provide keystore with the information required to enforce authorizations on key usage.
+ * @hide
  */
 interface IKeystoreAuthorization {
 
diff --git a/keystore2/aidl/android/security/authorization/LockScreenEvent.aidl b/keystore2/aidl/android/security/authorization/LockScreenEvent.aidl
index 877a916..c7553a2 100644
--- a/keystore2/aidl/android/security/authorization/LockScreenEvent.aidl
+++ b/keystore2/aidl/android/security/authorization/LockScreenEvent.aidl
@@ -14,6 +14,7 @@
 
 package android.security.authorization;
 
+/** @hide */
 @Backing(type="int")
 enum LockScreenEvent {
     UNLOCK = 0,
diff --git a/keystore2/aidl/android/security/authorization/ResponseCode.aidl b/keystore2/aidl/android/security/authorization/ResponseCode.aidl
index 94f1120..169dc7b 100644
--- a/keystore2/aidl/android/security/authorization/ResponseCode.aidl
+++ b/keystore2/aidl/android/security/authorization/ResponseCode.aidl
@@ -16,6 +16,7 @@
 
 /**
  * Used as exception codes by IKeystoreAuthorization.
+ * @hide
  */
 @Backing(type="int")
 enum ResponseCode {
diff --git a/keystore2/aidl/android/security/compat/IKeystoreCompatService.aidl b/keystore2/aidl/android/security/compat/IKeystoreCompatService.aidl
index 4b6a93b..50bfa19 100644
--- a/keystore2/aidl/android/security/compat/IKeystoreCompatService.aidl
+++ b/keystore2/aidl/android/security/compat/IKeystoreCompatService.aidl
@@ -25,6 +25,7 @@
  * The compatibility service allows Keystore 2.0 to connect to legacy wrapper implementations that
  * it hosts itself without registering them as a service. Keystore 2.0 would not be allowed to
  * register a HAL service, so instead it registers this service which it can then connect to.
+ * @hide
  */
 interface IKeystoreCompatService {
     /**
diff --git a/keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl b/keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl
index 3115e92..50e674d 100644
--- a/keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl
+++ b/keystore2/aidl/android/security/maintenance/IKeystoreMaintenance.aidl
@@ -35,7 +35,6 @@
      * user id.
      *
      * @param userId - Android user id
-     * @hide
      */
     void onUserAdded(in int userId);
 
@@ -47,7 +46,6 @@
      * `ResponseCode::SYSTEM_ERROR` - if failed to delete the keys of the user being deleted.
      *
      * @param userId - Android user id
-     * @hide
      */
     void onUserRemoved(in int userId);
 
@@ -62,7 +60,6 @@
      *
      * @param userId - Android user id
      * @param password - a secret derived from the synthetic password of the user
-     * @hide
      */
     void onUserPasswordChanged(in int userId, in @nullable byte[] password);
 
@@ -73,7 +70,6 @@
      * @param domain - One of Domain.APP or Domain.SELINUX.
      * @param nspace - The UID of the app that is to be cleared if domain is Domain.APP or
      *                 the SEPolicy namespace if domain is Domain.SELINUX.
-     * @hide
      */
      void clearNamespace(Domain domain, long nspace);
 
@@ -86,7 +82,6 @@
      * `ResponseCode::SYSTEM_ERROR` - if an error occurred when querying the user state.
      *
      * @param userId - Android user id
-     * @hide
      */
     UserState getState(in int userId);
 }
diff --git a/keystore2/aidl/android/security/maintenance/UserState.aidl b/keystore2/aidl/android/security/maintenance/UserState.aidl
index b6fe278..376f4fb 100644
--- a/keystore2/aidl/android/security/maintenance/UserState.aidl
+++ b/keystore2/aidl/android/security/maintenance/UserState.aidl
@@ -14,6 +14,7 @@
 
 package android.security.maintenance;
 
+/** @hide */
 @Backing(type="int")
 enum UserState {
     UNINITIALIZED = 0,
diff --git a/keystore2/aidl/android/security/vpnprofilestore/IVpnProfileStore.aidl b/keystore2/aidl/android/security/vpnprofilestore/IVpnProfileStore.aidl
index 054a4d7..8375b7b 100644
--- a/keystore2/aidl/android/security/vpnprofilestore/IVpnProfileStore.aidl
+++ b/keystore2/aidl/android/security/vpnprofilestore/IVpnProfileStore.aidl
@@ -18,7 +18,6 @@
 
 /**
  * Internal interface for accessing and storing VPN profiles.
- *
  * @hide
  */
 interface IVpnProfileStore {
diff --git a/keystore2/src/attestation_key_utils.rs b/keystore2/src/attestation_key_utils.rs
new file mode 100644
index 0000000..425eec6
--- /dev/null
+++ b/keystore2/src/attestation_key_utils.rs
@@ -0,0 +1,126 @@
+// 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.
+
+//! Implements get_attestation_key_info which loads remote provisioned or user
+//! generated attestation keys.
+
+use crate::database::{BlobMetaData, KeyEntryLoadBits, KeyType};
+use crate::database::{KeyIdGuard, KeystoreDB};
+use crate::error::{Error, ErrorCode};
+use crate::permission::KeyPerm;
+use crate::remote_provisioning::RemProvState;
+use crate::utils::check_key_permission;
+use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
+    AttestationKey::AttestationKey, Certificate::Certificate, KeyParameter::KeyParameter,
+};
+use android_system_keystore2::aidl::android::system::keystore2::{
+    Domain::Domain, KeyDescriptor::KeyDescriptor,
+};
+use anyhow::{Context, Result};
+use keystore2_crypto::parse_subject_from_certificate;
+
+/// KeyMint takes two different kinds of attestation keys. Remote provisioned keys
+/// and those that have been generated by the user. Unfortunately, they need to be
+/// handled quite differently, thus the different representations.
+pub enum AttestationKeyInfo {
+    RemoteProvisioned {
+        attestation_key: AttestationKey,
+        attestation_certs: Certificate,
+    },
+    UserGenerated {
+        key_id_guard: KeyIdGuard,
+        blob: Vec<u8>,
+        blob_metadata: BlobMetaData,
+        issuer_subject: Vec<u8>,
+    },
+}
+
+/// This function loads and, optionally, assigns the caller's remote provisioned
+/// attestation key or, if `attest_key_descriptor` is given, it loads the user
+/// generated attestation key from the database.
+pub fn get_attest_key_info(
+    key: &KeyDescriptor,
+    caller_uid: u32,
+    attest_key_descriptor: Option<&KeyDescriptor>,
+    params: &[KeyParameter],
+    rem_prov_state: &RemProvState,
+    db: &mut KeystoreDB,
+) -> Result<Option<AttestationKeyInfo>> {
+    match attest_key_descriptor {
+        None => rem_prov_state
+            .get_remotely_provisioned_attestation_key_and_certs(&key, caller_uid, params, db)
+            .context(concat!(
+                "In get_attest_key_and_cert_chain: ",
+                "Trying to get remotely provisioned attestation key."
+            ))
+            .map(|result| {
+                result.map(|(attestation_key, attestation_certs)| {
+                    AttestationKeyInfo::RemoteProvisioned { attestation_key, attestation_certs }
+                })
+            }),
+        Some(attest_key) => get_user_generated_attestation_key(&attest_key, caller_uid, db)
+            .context("In get_attest_key_and_cert_chain: Trying to load attest key")
+            .map(Some),
+    }
+}
+
+fn get_user_generated_attestation_key(
+    key: &KeyDescriptor,
+    caller_uid: u32,
+    db: &mut KeystoreDB,
+) -> Result<AttestationKeyInfo> {
+    let (key_id_guard, blob, cert, blob_metadata) =
+        load_attest_key_blob_and_cert(&key, caller_uid, db)
+            .context("In get_user_generated_attestation_key: Failed to load blob and cert")?;
+
+    let issuer_subject: Vec<u8> = parse_subject_from_certificate(&cert).context(
+        "In get_user_generated_attestation_key: Failed to parse subject from certificate.",
+    )?;
+
+    Ok(AttestationKeyInfo::UserGenerated { key_id_guard, blob, issuer_subject, blob_metadata })
+}
+
+fn load_attest_key_blob_and_cert(
+    key: &KeyDescriptor,
+    caller_uid: u32,
+    db: &mut KeystoreDB,
+) -> Result<(KeyIdGuard, Vec<u8>, Vec<u8>, BlobMetaData)> {
+    match key.domain {
+        Domain::BLOB => Err(Error::Km(ErrorCode::INVALID_ARGUMENT)).context(
+            "In load_attest_key_blob_and_cert: Domain::BLOB attestation keys not supported",
+        ),
+        _ => {
+            let (key_id_guard, mut key_entry) = db
+                .load_key_entry(
+                    &key,
+                    KeyType::Client,
+                    KeyEntryLoadBits::BOTH,
+                    caller_uid,
+                    |k, av| check_key_permission(KeyPerm::use_(), k, &av),
+                )
+                .context("In load_attest_key_blob_and_cert: Failed to load key.")?;
+
+            let (blob, blob_metadata) =
+                key_entry.take_key_blob_info().ok_or_else(Error::sys).context(concat!(
+                    "In load_attest_key_blob_and_cert: Successfully loaded key entry,",
+                    " but KM blob was missing."
+                ))?;
+            let cert = key_entry.take_cert().ok_or_else(Error::sys).context(concat!(
+                "In load_attest_key_blob_and_cert: Successfully loaded key entry,",
+                " but cert was missing."
+            ))?;
+            Ok((key_id_guard, blob, cert, blob_metadata))
+        }
+    }
+}
diff --git a/keystore2/src/authorization.rs b/keystore2/src/authorization.rs
index 5abb426..553746a 100644
--- a/keystore2/src/authorization.rs
+++ b/keystore2/src/authorization.rs
@@ -33,6 +33,7 @@
     ResponseCode::ResponseCode as KsResponseCode };
 use anyhow::{Context, Result};
 use binder::IBinderInternal;
+use keystore2_crypto::Password;
 use keystore2_selinux as selinux;
 
 /// This is the Authorization error type, it wraps binder exceptions and the
@@ -128,10 +129,10 @@
         &self,
         lock_screen_event: LockScreenEvent,
         user_id: i32,
-        password: Option<&[u8]>,
+        password: Option<Password>,
     ) -> Result<()> {
         match (lock_screen_event, password) {
-            (LockScreenEvent::UNLOCK, Some(user_password)) => {
+            (LockScreenEvent::UNLOCK, Some(password)) => {
                 //This corresponds to the unlock() method in legacy keystore API.
                 //check permission
                 check_keystore_permission(KeystorePerm::unlock())
@@ -145,7 +146,7 @@
                             &LEGACY_MIGRATOR,
                             &SUPER_KEY,
                             user_id as u32,
-                            user_password,
+                            &password,
                         )
                     })
                     .context("In on_lock_screen_event: Unlock with password.")?
@@ -213,7 +214,10 @@
         user_id: i32,
         password: Option<&[u8]>,
     ) -> BinderResult<()> {
-        map_or_log_err(self.on_lock_screen_event(lock_screen_event, user_id, password), Ok)
+        map_or_log_err(
+            self.on_lock_screen_event(lock_screen_event, user_id, password.map(|pw| pw.into())),
+            Ok,
+        )
     }
 
     fn getAuthTokensForCredStore(
diff --git a/keystore2/src/crypto/lib.rs b/keystore2/src/crypto/lib.rs
index bd5906c..98e6eef 100644
--- a/keystore2/src/crypto/lib.rs
+++ b/keystore2/src/crypto/lib.rs
@@ -149,42 +149,68 @@
     }
 }
 
-/// Generates a key from the given password and salt.
-/// The salt must be exactly 16 bytes long.
-/// Two key sizes are accepted: 16 and 32 bytes.
-pub fn derive_key_from_password(
-    pw: &[u8],
-    salt: Option<&[u8]>,
-    key_length: usize,
-) -> Result<ZVec, Error> {
-    let salt: *const u8 = match salt {
-        Some(s) => {
-            if s.len() != SALT_LENGTH {
-                return Err(Error::InvalidSaltLength);
-            }
-            s.as_ptr()
-        }
-        None => std::ptr::null(),
-    };
+/// Represents a "password" that can be used to key the PBKDF2 algorithm.
+pub enum Password<'a> {
+    /// Borrow an existing byte array
+    Ref(&'a [u8]),
+    /// Use an owned ZVec to store the key
+    Owned(ZVec),
+}
 
-    match key_length {
-        AES_128_KEY_LENGTH | AES_256_KEY_LENGTH => {}
-        _ => return Err(Error::InvalidKeyLength),
+impl<'a> From<&'a [u8]> for Password<'a> {
+    fn from(pw: &'a [u8]) -> Self {
+        Self::Ref(pw)
+    }
+}
+
+impl<'a> Password<'a> {
+    fn get_key(&'a self) -> &'a [u8] {
+        match self {
+            Self::Ref(b) => b,
+            Self::Owned(z) => &*z,
+        }
     }
 
-    let mut result = ZVec::new(key_length)?;
+    /// Generate a key from the given password and salt.
+    /// The salt must be exactly 16 bytes long.
+    /// Two key sizes are accepted: 16 and 32 bytes.
+    pub fn derive_key(&self, salt: Option<&[u8]>, key_length: usize) -> Result<ZVec, Error> {
+        let pw = self.get_key();
 
-    unsafe {
-        generateKeyFromPassword(
-            result.as_mut_ptr(),
-            result.len(),
-            pw.as_ptr() as *const std::os::raw::c_char,
-            pw.len(),
-            salt,
-        )
-    };
+        let salt: *const u8 = match salt {
+            Some(s) => {
+                if s.len() != SALT_LENGTH {
+                    return Err(Error::InvalidSaltLength);
+                }
+                s.as_ptr()
+            }
+            None => std::ptr::null(),
+        };
 
-    Ok(result)
+        match key_length {
+            AES_128_KEY_LENGTH | AES_256_KEY_LENGTH => {}
+            _ => return Err(Error::InvalidKeyLength),
+        }
+
+        let mut result = ZVec::new(key_length)?;
+
+        unsafe {
+            generateKeyFromPassword(
+                result.as_mut_ptr(),
+                result.len(),
+                pw.as_ptr() as *const std::os::raw::c_char,
+                pw.len(),
+                salt,
+            )
+        };
+
+        Ok(result)
+    }
+
+    /// Try to make another Password object with the same data.
+    pub fn try_clone(&self) -> Result<Password<'static>, Error> {
+        Ok(Password::Owned(ZVec::try_from(self.get_key())?))
+    }
 }
 
 /// Calls the boringssl HKDF_extract function.
diff --git a/keystore2/src/database.rs b/keystore2/src/database.rs
index 22ab02e..0e8b3d7 100644
--- a/keystore2/src/database.rs
+++ b/keystore2/src/database.rs
@@ -1357,15 +1357,11 @@
     }
 
     fn is_locked_error(e: &anyhow::Error) -> bool {
-        matches!(e.root_cause().downcast_ref::<rusqlite::ffi::Error>(),
-        Some(rusqlite::ffi::Error {
-            code: rusqlite::ErrorCode::DatabaseBusy,
-            ..
-        })
-        | Some(rusqlite::ffi::Error {
-            code: rusqlite::ErrorCode::DatabaseLocked,
-            ..
-        }))
+        matches!(
+            e.root_cause().downcast_ref::<rusqlite::ffi::Error>(),
+            Some(rusqlite::ffi::Error { code: rusqlite::ErrorCode::DatabaseBusy, .. })
+                | Some(rusqlite::ffi::Error { code: rusqlite::ErrorCode::DatabaseLocked, .. })
+        )
     }
 
     /// Creates a new key entry and allocates a new randomized id for the new key.
@@ -4889,7 +4885,7 @@
     #[test]
     fn test_store_super_key() -> Result<()> {
         let mut db = new_test_db()?;
-        let pw = "xyzabc".as_bytes();
+        let pw: keystore2_crypto::Password = (&b"xyzabc"[..]).into();
         let super_key = keystore2_crypto::generate_aes256_key()?;
         let secret = String::from("keystore2 is great.");
         let secret_bytes = secret.into_bytes();
diff --git a/keystore2/src/enforcements.rs b/keystore2/src/enforcements.rs
index 06648f1..2cc704b 100644
--- a/keystore2/src/enforcements.rs
+++ b/keystore2/src/enforcements.rs
@@ -804,7 +804,7 @@
                     let token_valid = now_in_millis
                         .checked_sub(auth_token_entry.time_received().milli_seconds())
                         .map_or(false, |token_age_in_millis| {
-                            token_age_in_millis > auth_token_max_age_millis
+                            auth_token_max_age_millis > token_age_in_millis
                         });
                     token_valid && auth_token_entry.satisfies(&sids, auth_type)
                 })
diff --git a/keystore2/src/legacy_blob.rs b/keystore2/src/legacy_blob.rs
index 3fc77b7..5f40ece 100644
--- a/keystore2/src/legacy_blob.rs
+++ b/keystore2/src/legacy_blob.rs
@@ -26,7 +26,7 @@
     SecurityLevel::SecurityLevel, Tag::Tag, TagType::TagType,
 };
 use anyhow::{Context, Result};
-use keystore2_crypto::{aes_gcm_decrypt, derive_key_from_password, ZVec};
+use keystore2_crypto::{aes_gcm_decrypt, Password, ZVec};
 use std::collections::{HashMap, HashSet};
 use std::{convert::TryInto, fs::File, path::Path, path::PathBuf};
 use std::{
@@ -1036,7 +1036,7 @@
     }
 
     /// Load and decrypt legacy super key blob.
-    pub fn load_super_key(&self, user_id: u32, pw: &[u8]) -> Result<Option<ZVec>> {
+    pub fn load_super_key(&self, user_id: u32, pw: &Password) -> Result<Option<ZVec>> {
         let path = self.make_super_key_filename(user_id);
         let blob = Self::read_generic_blob(&path)
             .context("In load_super_key: While loading super key.")?;
@@ -1046,7 +1046,8 @@
                 Blob {
                     value: BlobValue::PwEncrypted { iv, tag, data, salt, key_size }, ..
                 } => {
-                    let key = derive_key_from_password(pw, Some(&salt), key_size)
+                    let key = pw
+                        .derive_key(Some(&salt), key_size)
                         .context("In load_super_key: Failed to derive key from password.")?;
                     let blob = aes_gcm_decrypt(&data, &iv, &tag, &key).context(
                         "In load_super_key: while trying to decrypt legacy super key blob.",
@@ -1294,7 +1295,7 @@
             Some(&error::Error::Rc(ResponseCode::LOCKED))
         );
 
-        key_manager.unlock_user_key(&mut db, 0, PASSWORD, &legacy_blob_loader)?;
+        key_manager.unlock_user_key(&mut db, 0, &(PASSWORD.into()), &legacy_blob_loader)?;
 
         if let (Some((Blob { flags, value: _ }, _params)), Some(cert), Some(chain)) =
             legacy_blob_loader.load_by_uid_alias(10223, "authbound", Some(&key_manager))?
diff --git a/keystore2/src/legacy_migrator.rs b/keystore2/src/legacy_migrator.rs
index 1ae8719..7567070 100644
--- a/keystore2/src/legacy_migrator.rs
+++ b/keystore2/src/legacy_migrator.rs
@@ -29,9 +29,8 @@
 };
 use anyhow::{Context, Result};
 use core::ops::Deref;
-use keystore2_crypto::ZVec;
+use keystore2_crypto::{Password, ZVec};
 use std::collections::{HashMap, HashSet};
-use std::convert::TryInto;
 use std::sync::atomic::{AtomicU8, Ordering};
 use std::sync::mpsc::channel;
 use std::sync::{Arc, Mutex};
@@ -334,7 +333,7 @@
     pub fn with_try_migrate_super_key<F, T>(
         &self,
         user_id: u32,
-        pw: &[u8],
+        pw: &Password,
         mut key_accessor: F,
     ) -> Result<Option<T>>
     where
@@ -345,12 +344,9 @@
             Ok(None) => {}
             Err(e) => return Err(e),
         }
-
-        let pw: ZVec = pw
-            .try_into()
-            .context("In with_try_migrate_super_key: copying the password into a zvec.")?;
+        let pw = pw.try_clone().context("In with_try_migrate_super_key: Cloning password.")?;
         let result = self.do_serialized(move |migrator_state| {
-            migrator_state.check_and_migrate_super_key(user_id, pw)
+            migrator_state.check_and_migrate_super_key(user_id, &pw)
         });
 
         if let Some(result) = result {
@@ -550,7 +546,7 @@
         }
     }
 
-    fn check_and_migrate_super_key(&mut self, user_id: u32, pw: ZVec) -> Result<()> {
+    fn check_and_migrate_super_key(&mut self, user_id: u32, pw: &Password) -> Result<()> {
         if self.recently_migrated_super_key.contains(&user_id) {
             return Ok(());
         }
@@ -561,7 +557,7 @@
             .context("In check_and_migrate_super_key: Trying to load legacy super key.")?
         {
             let (blob, blob_metadata) =
-                crate::super_key::SuperKeyManager::encrypt_with_password(&super_key, &pw)
+                crate::super_key::SuperKeyManager::encrypt_with_password(&super_key, pw)
                     .context("In check_and_migrate_super_key: Trying to encrypt super key.")?;
 
             self.db.store_super_key(user_id, &(&blob, &blob_metadata)).context(concat!(
diff --git a/keystore2/src/lib.rs b/keystore2/src/lib.rs
index 2a00ae3..2e8ced6 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -36,6 +36,7 @@
 pub mod user_manager;
 pub mod utils;
 
+mod attestation_key_utils;
 mod db_utils;
 mod gc;
 mod super_key;
diff --git a/keystore2/src/operation.rs b/keystore2/src/operation.rs
index 5cc38cc..0f5bcaf 100644
--- a/keystore2/src/operation.rs
+++ b/keystore2/src/operation.rs
@@ -167,12 +167,14 @@
     outcome: Mutex<Outcome>,
     owner: u32, // Uid of the operation's owner.
     auth_info: Mutex<AuthInfo>,
+    forced: bool,
 }
 
 struct PruningInfo {
     last_usage: Instant,
     owner: u32,
     index: usize,
+    forced: bool,
 }
 
 // We don't except more than 32KiB of data in `update`, `updateAad`, and `finish`.
@@ -185,6 +187,7 @@
         km_op: binder::Strong<dyn IKeyMintOperation>,
         owner: u32,
         auth_info: AuthInfo,
+        forced: bool,
     ) -> Self {
         Self {
             index,
@@ -193,6 +196,7 @@
             outcome: Mutex::new(Outcome::Unknown),
             owner,
             auth_info: Mutex::new(auth_info),
+            forced,
         }
     }
 
@@ -218,6 +222,7 @@
             last_usage: *self.last_usage.lock().expect("In get_pruning_info."),
             owner: self.owner,
             index: self.index,
+            forced: self.forced,
         })
     }
 
@@ -465,6 +470,7 @@
         km_op: binder::public_api::Strong<dyn IKeyMintOperation>,
         owner: u32,
         auth_info: AuthInfo,
+        forced: bool,
     ) -> Arc<Operation> {
         // We use unwrap because we don't allow code that can panic while locked.
         let mut operations = self.operations.lock().expect("In create_operation.");
@@ -477,12 +483,13 @@
             s.upgrade().is_none()
         }) {
             Some(free_slot) => {
-                let new_op = Arc::new(Operation::new(index - 1, km_op, owner, auth_info));
+                let new_op = Arc::new(Operation::new(index - 1, km_op, owner, auth_info, forced));
                 *free_slot = Arc::downgrade(&new_op);
                 new_op
             }
             None => {
-                let new_op = Arc::new(Operation::new(operations.len(), km_op, owner, auth_info));
+                let new_op =
+                    Arc::new(Operation::new(operations.len(), km_op, owner, auth_info, forced));
                 operations.push(Arc::downgrade(&new_op));
                 new_op
             }
@@ -565,7 +572,7 @@
     /// ## Update
     /// We also allow callers to cannibalize their own sibling operations if no other
     /// slot can be found. In this case the least recently used sibling is pruned.
-    pub fn prune(&self, caller: u32) -> Result<(), Error> {
+    pub fn prune(&self, caller: u32, forced: bool) -> Result<(), Error> {
         loop {
             // Maps the uid of the owner to the number of operations that owner has
             // (running_siblings). More operations per owner lowers the pruning
@@ -590,7 +597,8 @@
                     }
                 });
 
-            let caller_malus = 1u64 + *owners.entry(caller).or_default();
+            // If the operation is forced, the caller has a malus of 0.
+            let caller_malus = if forced { 0 } else { 1u64 + *owners.entry(caller).or_default() };
 
             // We iterate through all operations computing the malus and finding
             // the candidate with the highest malus which must also be higher
@@ -604,7 +612,7 @@
             let mut oldest_caller_op: Option<CandidateInfo> = None;
             let candidate = pruning_info.iter().fold(
                 None,
-                |acc: Option<CandidateInfo>, &PruningInfo { last_usage, owner, index }| {
+                |acc: Option<CandidateInfo>, &PruningInfo { last_usage, owner, index, forced }| {
                     // Compute the age of the current operation.
                     let age = now
                         .checked_duration_since(last_usage)
@@ -624,12 +632,17 @@
                     }
 
                     // Compute the malus of the current operation.
-                    // Expect safety: Every owner in pruning_info was counted in
-                    // the owners map. So this unwrap cannot panic.
-                    let malus = *owners
-                        .get(&owner)
-                        .expect("This is odd. We should have counted every owner in pruning_info.")
-                        + ((age.as_secs() + 1) as f64).log(6.0).floor() as u64;
+                    let malus = if forced {
+                        // Forced operations have a malus of 0. And cannot even be pruned
+                        // by other forced operations.
+                        0
+                    } else {
+                        // Expect safety: Every owner in pruning_info was counted in
+                        // the owners map. So this unwrap cannot panic.
+                        *owners.get(&owner).expect(
+                            "This is odd. We should have counted every owner in pruning_info.",
+                        ) + ((age.as_secs() + 1) as f64).log(6.0).floor() as u64
+                    };
 
                     // Now check if the current operation is a viable/better candidate
                     // the one currently stored in the accumulator.
diff --git a/keystore2/src/remote_provisioning.rs b/keystore2/src/remote_provisioning.rs
index cc97573..8c04088 100644
--- a/keystore2/src/remote_provisioning.rs
+++ b/keystore2/src/remote_provisioning.rs
@@ -165,26 +165,26 @@
     ///
     /// It returns the ResponseCode `OUT_OF_KEYS` if there is not one key currently assigned to the
     /// `caller_uid` and there are none available to assign.
-    pub fn get_remote_provisioning_key_and_certs(
+    pub fn get_remotely_provisioned_attestation_key_and_certs(
         &self,
         key: &KeyDescriptor,
         caller_uid: u32,
         params: &[KeyParameter],
         db: &mut KeystoreDB,
-    ) -> Result<(Option<AttestationKey>, Option<Certificate>)> {
+    ) -> Result<Option<(AttestationKey, Certificate)>> {
         if !self.is_asymmetric_key(params) || !self.check_rem_prov_enabled(db)? {
             // There is no remote provisioning component for this security level on the
             // device. Return None so the underlying KM instance knows to use its
             // factory provisioned key instead. Alternatively, it's not an asymmetric key
             // and therefore will not be attested.
-            Ok((None, None))
+            Ok(None)
         } else {
             match self.get_rem_prov_attest_key(&key, caller_uid, db).context(concat!(
                 "In get_remote_provisioning_key_and_certs: Failed to get ",
                 "attestation key"
             ))? {
-                Some(cert_chain) => Ok((
-                    Some(AttestationKey {
+                Some(cert_chain) => Ok(Some((
+                    AttestationKey {
                         keyBlob: cert_chain.private_key.to_vec(),
                         attestKeyParams: vec![],
                         issuerSubjectName: parse_subject_from_certificate(&cert_chain.batch_cert)
@@ -192,10 +192,10 @@
                             "In get_remote_provisioning_key_and_certs: Failed to ",
                             "parse subject."
                         ))?,
-                    }),
-                    Some(Certificate { encodedCertificate: cert_chain.cert_chain }),
-                )),
-                None => Ok((None, None)),
+                    },
+                    Certificate { encodedCertificate: cert_chain.cert_chain },
+                ))),
+                None => Ok(None),
             }
         }
     }
diff --git a/keystore2/src/security_level.rs b/keystore2/src/security_level.rs
index 641b77a..5a776fb 100644
--- a/keystore2/src/security_level.rs
+++ b/keystore2/src/security_level.rs
@@ -18,7 +18,7 @@
 
 use crate::globals::get_keymint_device;
 use android_hardware_security_keymint::aidl::android::hardware::security::keymint::{
-    Algorithm::Algorithm, AttestationKey::AttestationKey, Certificate::Certificate,
+    Algorithm::Algorithm, AttestationKey::AttestationKey,
     HardwareAuthenticatorType::HardwareAuthenticatorType, IKeyMintDevice::IKeyMintDevice,
     KeyCreationResult::KeyCreationResult, KeyFormat::KeyFormat,
     KeyMintHardwareInfo::KeyMintHardwareInfo, KeyParameter::KeyParameter,
@@ -32,7 +32,8 @@
     KeyMetadata::KeyMetadata, KeyParameters::KeyParameters,
 };
 
-use crate::database::{CertificateInfo, KeyIdGuard, KeystoreDB};
+use crate::attestation_key_utils::{get_attest_key_info, AttestationKeyInfo};
+use crate::database::{CertificateInfo, KeyIdGuard};
 use crate::globals::{DB, ENFORCEMENTS, LEGACY_MIGRATOR, SUPER_KEY};
 use crate::key_parameter::KeyParameter as KsKeyParam;
 use crate::key_parameter::KeyParameterValue as KsKeyParamValue;
@@ -57,7 +58,6 @@
 };
 use anyhow::{anyhow, Context, Result};
 use binder::{IBinderInternal, Strong, ThreadState};
-use keystore2_crypto::parse_subject_from_certificate;
 
 /// Implementation of the IKeystoreSecurityLevel Interface.
 pub struct KeystoreSecurityLevel {
@@ -209,6 +209,11 @@
             Domain::BLOB => {
                 check_key_permission(KeyPerm::use_(), key, &None)
                     .context("In create_operation: checking use permission for Domain::BLOB.")?;
+                if forced {
+                    check_key_permission(KeyPerm::req_forced_op(), key, &None).context(
+                        "In create_operation: checking forced permission for Domain::BLOB.",
+                    )?;
+                }
                 (
                     match &key.blob {
                         Some(blob) => blob,
@@ -233,7 +238,13 @@
                                 KeyType::Client,
                                 KeyEntryLoadBits::KM,
                                 caller_uid,
-                                |k, av| check_key_permission(KeyPerm::use_(), k, &av),
+                                |k, av| {
+                                    check_key_permission(KeyPerm::use_(), k, &av)?;
+                                    if forced {
+                                        check_key_permission(KeyPerm::req_forced_op(), k, &av)?;
+                                    }
+                                    Ok(())
+                                },
                             )
                         })
                     })
@@ -270,17 +281,12 @@
                 purpose,
                 key_properties.as_ref(),
                 operation_parameters.as_ref(),
-                // TODO b/178222844 Replace this with the configuration returned by
-                //      KeyMintDevice::getHardwareInfo.
-                //      For now we assume that strongbox implementations need secure timestamps.
-                self.security_level == SecurityLevel::STRONGBOX,
+                self.hw_info.timestampTokenRequired,
             )
             .context("In create_operation.")?;
 
         let immediate_hat = immediate_hat.unwrap_or_default();
 
-        let user_id = uid_to_android_user(caller_uid);
-
         let km_blob = SUPER_KEY
             .unwrap_key_if_required(&blob_metadata, km_blob)
             .context("In create_operation. Failed to handle super encryption.")?;
@@ -304,7 +310,7 @@
                         &immediate_hat,
                     )) {
                         Err(Error::Km(ErrorCode::TOO_MANY_OPERATIONS)) => {
-                            self.operation_db.prune(caller_uid)?;
+                            self.operation_db.prune(caller_uid, forced)?;
                             continue;
                         }
                         v => return v,
@@ -317,7 +323,7 @@
 
         let operation = match begin_result.operation {
             Some(km_op) => {
-                self.operation_db.create_operation(km_op, caller_uid, auth_info)
+                self.operation_db.create_operation(km_op, caller_uid, auth_info, forced)
             },
             None => return Err(Error::sys()).context("In create_operation: Begin operation returned successfully, but did not return a valid operation."),
         };
@@ -335,6 +341,10 @@
                 0 => None,
                 _ => Some(KeyParameters { keyParameter: begin_result.params }),
             },
+            // An upgraded blob should only be returned if the caller has permission
+            // to use Domain::BLOB keys. If we got to this point, we already checked
+            // that the caller had that permission.
+            upgradedBlob: if key.domain == Domain::BLOB { upgraded_blob } else { None },
         })
     }
 
@@ -419,16 +429,19 @@
         };
 
         // generate_key requires the rebind permission.
+        // Must return on error for security reasons.
         check_key_permission(KeyPerm::rebind(), &key, &None).context("In generate_key.")?;
-        let (attest_key, cert_chain) = match (key.domain, attest_key_descriptor) {
-            (Domain::BLOB, None) => (None, None),
+
+        let attestation_key_info = match (key.domain, attest_key_descriptor) {
+            (Domain::BLOB, _) => None,
             _ => DB
-                .with::<_, Result<(Option<AttestationKey>, Option<Certificate>)>>(|db| {
-                    self.get_attest_key_and_cert_chain(
+                .with(|db| {
+                    get_attest_key_info(
                         &key,
                         caller_uid,
                         attest_key_descriptor,
                         params,
+                        &self.rem_prov_state,
                         &mut db.borrow_mut(),
                     )
                 })
@@ -438,92 +451,47 @@
             .context("In generate_key: Trying to get aaid.")?;
 
         let km_dev: Strong<dyn IKeyMintDevice> = self.keymint.get_interface()?;
-        map_km_error(km_dev.addRngEntropy(entropy))
-            .context("In generate_key: Trying to add entropy.")?;
-        let mut creation_result = map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
-            .context("In generate_key: While generating Key")?;
-        // The certificate chain ultimately gets flattened into a big DER encoded byte array,
-        // so providing that blob upfront in a single certificate entry should be fine.
-        if let Some(cert) = cert_chain {
-            creation_result.certificateChain.push(cert);
+
+        let creation_result = match attestation_key_info {
+            Some(AttestationKeyInfo::UserGenerated {
+                key_id_guard,
+                blob,
+                blob_metadata,
+                issuer_subject,
+            }) => self
+                .upgrade_keyblob_if_required_with(
+                    &*km_dev,
+                    Some(key_id_guard),
+                    &(&KeyBlob::Ref(&blob), &blob_metadata),
+                    &params,
+                    |blob| {
+                        let attest_key = Some(AttestationKey {
+                            keyBlob: blob.to_vec(),
+                            attestKeyParams: vec![],
+                            issuerSubjectName: issuer_subject.clone(),
+                        });
+                        map_km_error(km_dev.generateKey(&params, attest_key.as_ref()))
+                    },
+                )
+                .context("In generate_key: Using user generated attestation key.")
+                .map(|(result, _)| result),
+            Some(AttestationKeyInfo::RemoteProvisioned { attestation_key, attestation_certs }) => {
+                map_km_error(km_dev.generateKey(&params, Some(&attestation_key)))
+                    .context("While generating Key with remote provisioned attestation key.")
+                    .map(|mut creation_result| {
+                        creation_result.certificateChain.push(attestation_certs);
+                        creation_result
+                    })
+            }
+            None => map_km_error(km_dev.generateKey(&params, None))
+                .context("While generating Key without explicit attestation key."),
         }
+        .context("In generate_key.")?;
+
         let user_id = uid_to_android_user(caller_uid);
         self.store_new_key(key, creation_result, user_id, Some(flags)).context("In generate_key.")
     }
 
-    fn get_attest_key_and_cert_chain(
-        &self,
-        key: &KeyDescriptor,
-        caller_uid: u32,
-        attest_key_descriptor: Option<&KeyDescriptor>,
-        params: &[KeyParameter],
-        db: &mut KeystoreDB,
-    ) -> Result<(Option<AttestationKey>, Option<Certificate>)> {
-        match attest_key_descriptor {
-            None => self
-                .rem_prov_state
-                .get_remote_provisioning_key_and_certs(&key, caller_uid, params, db),
-            Some(attest_key) => Ok((
-                Some(
-                    self.get_attest_key(&attest_key, caller_uid)
-                        .context("In generate_key: Trying to load attest key")?,
-                ),
-                None,
-            )),
-        }
-    }
-
-    fn get_attest_key(&self, key: &KeyDescriptor, caller_uid: u32) -> Result<AttestationKey> {
-        let (km_blob, cert) = self
-            .load_attest_key_blob_and_cert(&key, caller_uid)
-            .context("In get_attest_key: Failed to load blob and cert")?;
-
-        let issuer_subject: Vec<u8> = parse_subject_from_certificate(&cert)
-            .context("In get_attest_key: Failed to parse subject from certificate.")?;
-
-        Ok(AttestationKey {
-            keyBlob: km_blob.to_vec(),
-            attestKeyParams: [].to_vec(),
-            issuerSubjectName: issuer_subject,
-        })
-    }
-
-    fn load_attest_key_blob_and_cert(
-        &self,
-        key: &KeyDescriptor,
-        caller_uid: u32,
-    ) -> Result<(Vec<u8>, Vec<u8>)> {
-        match key.domain {
-            Domain::BLOB => Err(error::Error::Km(ErrorCode::INVALID_ARGUMENT)).context(
-                "In load_attest_key_blob_and_cert: Domain::BLOB attestation keys not supported",
-            ),
-            _ => {
-                let (key_id_guard, mut key_entry) = DB
-                    .with::<_, Result<(KeyIdGuard, KeyEntry)>>(|db| {
-                        db.borrow_mut().load_key_entry(
-                            &key,
-                            KeyType::Client,
-                            KeyEntryLoadBits::BOTH,
-                            caller_uid,
-                            |k, av| check_key_permission(KeyPerm::use_(), k, &av),
-                        )
-                    })
-                    .context("In load_attest_key_blob_and_cert: Failed to load key.")?;
-
-                let (blob, _) =
-                    key_entry.take_key_blob_info().ok_or_else(Error::sys).context(concat!(
-                        "In load_attest_key_blob_and_cert: Successfully loaded key entry,",
-                        " but KM blob was missing."
-                    ))?;
-                let cert = key_entry.take_cert().ok_or_else(Error::sys).context(concat!(
-                    "In load_attest_key_blob_and_cert: Successfully loaded key entry,",
-                    " but cert was missing."
-                ))?;
-                Ok((blob, cert))
-            }
-        }
-    }
-
     fn import_key(
         &self,
         key: &KeyDescriptor,
diff --git a/keystore2/src/super_key.rs b/keystore2/src/super_key.rs
index 5ee685a..aa434d6 100644
--- a/keystore2/src/super_key.rs
+++ b/keystore2/src/super_key.rs
@@ -23,8 +23,8 @@
 use android_system_keystore2::aidl::android::system::keystore2::Domain::Domain;
 use anyhow::{Context, Result};
 use keystore2_crypto::{
-    aes_gcm_decrypt, aes_gcm_encrypt, derive_key_from_password, generate_aes256_key, generate_salt,
-    ZVec, AES_256_KEY_LENGTH,
+    aes_gcm_decrypt, aes_gcm_encrypt, generate_aes256_key, generate_salt, Password, ZVec,
+    AES_256_KEY_LENGTH,
 };
 use std::ops::Deref;
 use std::{
@@ -130,7 +130,7 @@
         &self,
         db: &mut KeystoreDB,
         user: UserId,
-        pw: &[u8],
+        pw: &Password,
         legacy_blob_loader: &LegacyBlobLoader,
     ) -> Result<()> {
         let (_, entry) = db
@@ -233,7 +233,7 @@
         db: &mut KeystoreDB,
         legacy_migrator: &LegacyMigrator,
         user_id: u32,
-        pw: &[u8],
+        pw: &Password,
     ) -> Result<UserState> {
         let result = legacy_migrator
             .with_try_migrate_super_key(user_id, pw, || db.load_super_key(user_id))
@@ -261,7 +261,7 @@
         db: &mut KeystoreDB,
         legacy_migrator: &LegacyMigrator,
         user_id: u32,
-        pw: Option<&[u8]>,
+        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(
@@ -296,7 +296,7 @@
         &self,
         user_id: u32,
         entry: KeyEntry,
-        pw: &[u8],
+        pw: &Password,
     ) -> Result<SuperKey> {
         let super_key = Self::extract_super_key_from_key_entry(entry, pw).context(
             "In populate_cache_from_super_key_blob. Failed to extract super key from key entry",
@@ -306,7 +306,7 @@
     }
 
     /// Extracts super key from the entry loaded from the database
-    pub fn extract_super_key_from_key_entry(entry: KeyEntry, pw: &[u8]) -> Result<SuperKey> {
+    pub fn extract_super_key_from_key_entry(entry: KeyEntry, pw: &Password) -> Result<SuperKey> {
         if let Some((blob, metadata)) = entry.key_blob_info() {
             let key = match (
                 metadata.encrypted_by(),
@@ -315,9 +315,9 @@
                 metadata.aead_tag(),
             ) {
                 (Some(&EncryptedBy::Password), Some(salt), Some(iv), Some(tag)) => {
-                    let key = derive_key_from_password(pw, Some(salt), AES_256_KEY_LENGTH).context(
-                    "In extract_super_key_from_key_entry: Failed to generate key from password.",
-                )?;
+                    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.",
+                    )?;
 
                     aes_gcm_decrypt(blob, iv, tag, &key).context(
                         "In extract_super_key_from_key_entry: Failed to decrypt key blob.",
@@ -344,9 +344,13 @@
     }
 
     /// Encrypts the super key from a key derived from the password, before storing in the database.
-    pub fn encrypt_with_password(super_key: &[u8], pw: &[u8]) -> Result<(Vec<u8>, BlobMetaData)> {
+    pub fn encrypt_with_password(
+        super_key: &[u8],
+        pw: &Password,
+    ) -> Result<(Vec<u8>, BlobMetaData)> {
         let salt = generate_salt().context("In encrypt_with_password: Failed to generate salt.")?;
-        let derived_key = derive_key_from_password(pw, Some(&salt), AES_256_KEY_LENGTH)
+        let derived_key = pw
+            .derive_key(Some(&salt), AES_256_KEY_LENGTH)
             .context("In encrypt_with_password: Failed to derive password.")?;
         let mut metadata = BlobMetaData::new();
         metadata.add(BlobMetaEntry::EncryptedBy(EncryptedBy::Password));
@@ -514,7 +518,7 @@
         legacy_migrator: &LegacyMigrator,
         skm: &SuperKeyManager,
         user_id: u32,
-        password: Option<&[u8]>,
+        password: Option<&Password>,
     ) -> Result<UserState> {
         match skm.get_per_boot_key_by_user_id(user_id) {
             Some(super_key) => {
@@ -548,7 +552,7 @@
         legacy_migrator: &LegacyMigrator,
         skm: &SuperKeyManager,
         user_id: u32,
-        password: &[u8],
+        password: &Password,
     ) -> Result<UserState> {
         match skm.get_per_boot_key_by_user_id(user_id) {
             Some(super_key) => {
diff --git a/keystore2/src/user_manager.rs b/keystore2/src/user_manager.rs
index 123f3a1..0cc2e92 100644
--- a/keystore2/src/user_manager.rs
+++ b/keystore2/src/user_manager.rs
@@ -29,6 +29,7 @@
 use android_system_keystore2::aidl::android::system::keystore2::ResponseCode::ResponseCode;
 use anyhow::{Context, Result};
 use binder::{IBinderInternal, Strong};
+use keystore2_crypto::Password;
 
 /// This struct is defined to implement the aforementioned AIDL interface.
 /// As of now, it is an empty struct.
@@ -42,7 +43,7 @@
         Ok(result)
     }
 
-    fn on_user_password_changed(user_id: i32, password: Option<&[u8]>) -> Result<()> {
+    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_keystore_permission(KeystorePerm::change_password())
@@ -55,7 +56,7 @@
                     &LEGACY_MIGRATOR,
                     &SUPER_KEY,
                     user_id as u32,
-                    password,
+                    password.as_ref(),
                 )
             })
             .context("In on_user_password_changed.")?
@@ -121,7 +122,7 @@
 
 impl IKeystoreMaintenance for Maintenance {
     fn onUserPasswordChanged(&self, user_id: i32, password: Option<&[u8]>) -> BinderResult<()> {
-        map_or_log_err(Self::on_user_password_changed(user_id, password), Ok)
+        map_or_log_err(Self::on_user_password_changed(user_id, password.map(|pw| pw.into())), Ok)
     }
 
     fn onUserAdded(&self, user_id: i32) -> BinderResult<()> {
diff --git a/ondevice-signing/Android.bp b/ondevice-signing/Android.bp
index 1c3706d..2e5e02e 100644
--- a/ondevice-signing/Android.bp
+++ b/ondevice-signing/Android.bp
@@ -90,6 +90,8 @@
     "VerityUtils.cpp",
   ],
 
+  header_libs: ["odrefresh_headers"],
+
   static_libs: [
     "libmini_keyctl_static", // TODO need static?
     "libc++fs",
diff --git a/ondevice-signing/odsign_main.cpp b/ondevice-signing/odsign_main.cpp
index 58e49ec..f6d663b 100644
--- a/ondevice-signing/odsign_main.cpp
+++ b/ondevice-signing/odsign_main.cpp
@@ -29,6 +29,7 @@
 #include <android-base/properties.h>
 #include <android-base/scopeguard.h>
 #include <logwrap/logwrap.h>
+#include <odrefresh/odrefresh.h>
 
 #include "CertUtils.h"
 #include "KeymasterSigningKey.h"
@@ -56,7 +57,7 @@
 static const char* kFsVerityProcPath = "/proc/sys/fs/verity";
 
 static const bool kForceCompilation = false;
-static const bool kUseKeystore = false;
+static const bool kUseKeystore = true;
 
 static const char* kOdsignVerificationDoneProp = "odsign.verification.done";
 static const char* kOdsignKeyDoneProp = "odsign.key.done";
@@ -101,18 +102,11 @@
     return {};
 }
 
-bool compileArtifacts(bool force) {
+art::odrefresh::ExitCode compileArtifacts(bool force) {
     const char* const argv[] = {kOdrefreshPath, force ? "--force-compile" : "--compile"};
-
-    return logwrap_fork_execvp(arraysize(argv), argv, nullptr, false, LOG_ALOG, false, nullptr) ==
-           0;
-}
-
-bool validateArtifacts() {
-    const char* const argv[] = {kOdrefreshPath, "--check"};
-
-    return logwrap_fork_execvp(arraysize(argv), argv, nullptr, false, LOG_ALOG, false, nullptr) ==
-           0;
+    const int exit_code =
+        logwrap_fork_execvp(arraysize(argv), argv, nullptr, false, LOG_ALOG, false, nullptr);
+    return static_cast<art::odrefresh::ExitCode>(exit_code);
 }
 
 static std::string toHex(const std::vector<uint8_t>& digest) {
@@ -349,30 +343,27 @@
         }
     }
 
-    // Ask ART whether it considers the artifacts valid
-    LOG(INFO) << "Asking odrefresh to verify artifacts (if present)...";
-    bool artifactsValid = validateArtifacts();
-    LOG(INFO) << "odrefresh said they are " << (artifactsValid ? "VALID" : "INVALID");
-
-    // A post-condition of validating artifacts is that if the ones on /system
-    // are used, kArtArtifactsDir is removed. Conversely, if kArtArtifactsDir
-    // exists, those are artifacts that will be used, and we should verify them.
-    int err = access(kArtArtifactsDir.c_str(), F_OK);
-    // If we receive any error other than ENOENT, be suspicious
-    bool artifactsPresent = (err == 0) || (err < 0 && errno != ENOENT);
-    if (artifactsPresent) {
-        auto verificationResult = verifyArtifacts(*key, supportsFsVerity);
-        if (!verificationResult.ok()) {
-            LOG(ERROR) << verificationResult.error().message();
-            return -1;
+    art::odrefresh::ExitCode odrefresh_status = compileArtifacts(kForceCompilation);
+    if (odrefresh_status == art::odrefresh::ExitCode::kOkay) {
+        LOG(INFO) << "odrefresh said artifacts are VALID";
+        // A post-condition of validating artifacts is that if the ones on /system
+        // are used, kArtArtifactsDir is removed. Conversely, if kArtArtifactsDir
+        // exists, those are artifacts that will be used, and we should verify them.
+        int err = access(kArtArtifactsDir.c_str(), F_OK);
+        // If we receive any error other than ENOENT, be suspicious
+        bool artifactsPresent = (err == 0) || (err < 0 && errno != ENOENT);
+        if (artifactsPresent) {
+            auto verificationResult = verifyArtifacts(*key, supportsFsVerity);
+            if (!verificationResult.ok()) {
+                LOG(ERROR) << verificationResult.error().message();
+                return -1;
+            }
         }
-    }
-
-    if (!artifactsValid || kForceCompilation) {
-        LOG(INFO) << "Starting compilation... ";
-        bool ret = compileArtifacts(kForceCompilation);
-        LOG(INFO) << "Compilation done, returned " << ret;
-
+    } else if (odrefresh_status == art::odrefresh::ExitCode::kCompilationSuccess ||
+               odrefresh_status == art::odrefresh::ExitCode::kCompilationFailed) {
+        const bool compiled_all = odrefresh_status == art::odrefresh::ExitCode::kCompilationSuccess;
+        LOG(INFO) << "odrefresh compiled " << (compiled_all ? "all" : "partial")
+                  << " artifacts, returned " << odrefresh_status;
         Result<std::map<std::string, std::string>> digests;
         if (supportsFsVerity) {
             digests = addFilesToVerityRecursive(kArtArtifactsDir, *key);
@@ -385,12 +376,17 @@
             LOG(ERROR) << digests.error().message();
             return -1;
         }
-
         auto persistStatus = persistDigests(*digests, *key);
         if (!persistStatus.ok()) {
             LOG(ERROR) << persistStatus.error().message();
             return -1;
         }
+    } else if (odrefresh_status == art::odrefresh::ExitCode::kCleanupFailed) {
+        LOG(ERROR) << "odrefresh failed cleaning up existing artifacts";
+        return -1;
+    } else {
+        LOG(ERROR) << "odrefresh exited unexpectedly, returned " << odrefresh_status;
+        return -1;
     }
 
     LOG(INFO) << "On-device signing done.";