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/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/lib.rs b/keystore2/src/lib.rs
index 0e51eff..bdb599d 100644
--- a/keystore2/src/lib.rs
+++ b/keystore2/src/lib.rs
@@ -35,6 +35,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/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 2ad8fe1..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 {
@@ -341,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 },
         })
     }
 
@@ -425,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(),
                     )
                 })
@@ -444,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/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 33d04ca..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"
@@ -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.";