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 641b77a..2af3544 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 {
@@ -419,16 +419,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 +441,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,