diff --git a/FsCrypt.cpp b/FsCrypt.cpp
index 21495be..c0ec3eb 100644
--- a/FsCrypt.cpp
+++ b/FsCrypt.cpp
@@ -64,6 +64,9 @@
 using android::vold::BuildDataPath;
 using android::vold::kEmptyAuthentication;
 using android::vold::KeyBuffer;
+using android::vold::makeGen;
+using android::vold::retrieveKey;
+using android::vold::retrieveOrGenerateKey;
 using android::vold::writeStringToFile;
 using namespace android::fscrypt;
 
@@ -183,7 +186,7 @@
     auto const paths = get_ce_key_paths(directory_path);
     for (auto const ce_key_path : paths) {
         LOG(DEBUG) << "Trying user CE key " << ce_key_path;
-        if (android::vold::retrieveKey(ce_key_path, auth, ce_key)) {
+        if (retrieveKey(ce_key_path, auth, ce_key)) {
             LOG(DEBUG) << "Successfully retrieved key";
             fixate_user_ce_key(directory_path, ce_key_path, paths);
             return true;
@@ -274,8 +277,8 @@
     EncryptionOptions options;
     if (!get_data_file_encryption_options(&options)) return false;
     KeyBuffer de_key, ce_key;
-    if (!generateStorageKey(options, &de_key)) return false;
-    if (!generateStorageKey(options, &ce_key)) return false;
+    if (!generateStorageKey(makeGen(options), &de_key)) return false;
+    if (!generateStorageKey(makeGen(options), &ce_key)) return false;
     if (create_ephemeral) {
         // If the key should be created as ephemeral, don't store it.
         s_ephemeral_users.insert(user_id);
@@ -349,7 +352,7 @@
         if (s_de_policies.count(user_id) == 0) {
             auto key_path = de_dir + "/" + entry->d_name;
             KeyBuffer de_key;
-            if (!android::vold::retrieveKey(key_path, kEmptyAuthentication, &de_key)) return false;
+            if (!retrieveKey(key_path, kEmptyAuthentication, &de_key)) return false;
             EncryptionPolicy de_policy;
             if (!install_storage_key(DATA_MNT_POINT, options, de_key, &de_policy)) return false;
             s_de_policies[user_id] = de_policy;
@@ -372,8 +375,8 @@
     if (!get_data_file_encryption_options(&options)) return false;
 
     KeyBuffer device_key;
-    if (!android::vold::retrieveKey(true, kEmptyAuthentication, device_key_path, device_key_temp,
-                                    options, &device_key))
+    if (!retrieveOrGenerateKey(device_key_path, device_key_temp, kEmptyAuthentication,
+                               makeGen(options), &device_key))
         return false;
 
     EncryptionPolicy device_policy;
@@ -392,7 +395,7 @@
     LOG(INFO) << "Wrote system DE key reference to:" << ref_filename;
 
     KeyBuffer per_boot_key;
-    if (!generateStorageKey(options, &per_boot_key)) return false;
+    if (!generateStorageKey(makeGen(options), &per_boot_key)) return false;
     EncryptionPolicy per_boot_policy;
     if (!install_storage_key(DATA_MNT_POINT, options, per_boot_key, &per_boot_policy)) return false;
     std::string per_boot_ref_filename = std::string("/data") + fscrypt_key_per_boot_ref;
@@ -601,7 +604,7 @@
     EncryptionOptions options;
     if (!get_volume_file_encryption_options(&options)) return false;
     KeyBuffer key;
-    if (!android::vold::retrieveKey(true, auth, key_path, key_path + "_tmp", options, &key))
+    if (!retrieveOrGenerateKey(key_path, key_path + "_tmp", auth, makeGen(options), &key))
         return false;
     if (!install_storage_key(BuildDataPath(volume_uuid), options, key, policy)) return false;
     return true;
@@ -620,12 +623,12 @@
     auto const directory_path = get_ce_key_directory_path(user_id);
     KeyBuffer ce_key;
     std::string ce_key_current_path = get_ce_key_current_path(directory_path);
-    if (android::vold::retrieveKey(ce_key_current_path, retrieve_auth, &ce_key)) {
+    if (retrieveKey(ce_key_current_path, retrieve_auth, &ce_key)) {
         LOG(DEBUG) << "Successfully retrieved key";
         // TODO(147732812): Remove this once Locksettingservice is fixed.
         // Currently it calls fscrypt_clear_user_key_auth with a secret when lockscreen is
         // changed from swipe to none or vice-versa
-    } else if (android::vold::retrieveKey(ce_key_current_path, kEmptyAuthentication, &ce_key)) {
+    } else if (retrieveKey(ce_key_current_path, kEmptyAuthentication, &ce_key)) {
         LOG(DEBUG) << "Successfully retrieved key with empty auth";
     } else {
         LOG(ERROR) << "Failed to retrieve key for user " << user_id;
diff --git a/KeyUtil.cpp b/KeyUtil.cpp
index ae4d70b..2e810ff 100644
--- a/KeyUtil.cpp
+++ b/KeyUtil.cpp
@@ -36,8 +36,20 @@
 namespace android {
 namespace vold {
 
-bool randomKey(KeyBuffer* key) {
-    *key = KeyBuffer(FSCRYPT_MAX_KEY_SIZE);
+const KeyGeneration makeGen(const EncryptionOptions& options) {
+    return KeyGeneration{FSCRYPT_MAX_KEY_SIZE, true, options.use_hw_wrapped_key};
+}
+
+const KeyGeneration makeGen(const CryptoType& crypto) {
+    return KeyGeneration{crypto.get_keysize(), true, false};
+}
+
+const KeyGeneration neverGen() {
+    return KeyGeneration{0, false, false};
+}
+
+static bool randomKey(size_t size, KeyBuffer* key) {
+    *key = KeyBuffer(size);
     if (ReadRandomBytes(key->size(), key->data()) != 0) {
         // TODO status_t plays badly with PLOG, fix it.
         LOG(ERROR) << "Random read failed";
@@ -46,11 +58,17 @@
     return true;
 }
 
-bool generateStorageKey(const EncryptionOptions& options, KeyBuffer* key) {
-    if (options.use_hw_wrapped_key) {
+bool generateStorageKey(const KeyGeneration& gen, KeyBuffer* key) {
+    if (!gen.allow_gen) return false;
+    if (gen.use_hw_wrapped_key) {
+        if (gen.keysize != FSCRYPT_MAX_KEY_SIZE) {
+            LOG(ERROR) << "Cannot generate a wrapped key " << gen.keysize << " bytes long";
+            return false;
+        }
         return generateWrappedStorageKey(key);
+    } else {
+        return randomKey(gen.keysize, key);
     }
-    return randomKey(key);
 }
 
 // Return true if the kernel supports the ioctls to add/remove fscrypt keys
@@ -315,19 +333,19 @@
     return true;
 }
 
-bool retrieveKey(bool create_if_absent, const KeyAuthentication& key_authentication,
-                 const std::string& key_path, const std::string& tmp_path,
-                 const EncryptionOptions& options, KeyBuffer* key, bool keepOld) {
+bool retrieveOrGenerateKey(const std::string& key_path, const std::string& tmp_path,
+                           const KeyAuthentication& key_authentication, const KeyGeneration& gen,
+                           KeyBuffer* key, bool keepOld) {
     if (pathExists(key_path)) {
         LOG(DEBUG) << "Key exists, using: " << key_path;
         if (!retrieveKey(key_path, key_authentication, key, keepOld)) return false;
     } else {
-        if (!create_if_absent) {
+        if (!gen.allow_gen) {
             LOG(ERROR) << "No key found in " << key_path;
             return false;
         }
         LOG(INFO) << "Creating new key in " << key_path;
-        if (!generateStorageKey(options, key)) return false;
+        if (!generateStorageKey(gen, key)) return false;
         if (!storeKeyAtomically(key_path, tmp_path, key_authentication, *key)) return false;
     }
     return true;
diff --git a/KeyUtil.h b/KeyUtil.h
index 878b4ab..16aaf99 100644
--- a/KeyUtil.h
+++ b/KeyUtil.h
@@ -17,6 +17,7 @@
 #ifndef ANDROID_VOLD_KEYUTIL_H
 #define ANDROID_VOLD_KEYUTIL_H
 
+#include "CryptoType.h"
 #include "KeyBuffer.h"
 #include "KeyStorage.h"
 
@@ -30,9 +31,26 @@
 
 using namespace android::fscrypt;
 
-bool randomKey(KeyBuffer* key);
+// Description of how to generate a key when needed.
+struct KeyGeneration {
+    size_t keysize;
+    bool allow_gen;
+    bool use_hw_wrapped_key;
+};
 
-bool generateStorageKey(const EncryptionOptions& options, KeyBuffer* key);
+// Generate a key as specified in KeyGeneration
+bool generateStorageKey(const KeyGeneration& gen, KeyBuffer* key);
+
+// Returns KeyGeneration suitable for key as described in EncryptionOptions
+const KeyGeneration makeGen(const EncryptionOptions& options);
+
+// Returns KeyGeneration suitable for key as described in CryptoType
+const KeyGeneration makeGen(const CryptoType& crypto);
+
+// Returns a key with allow_gen false so generateStorageKey returns false;
+// this is used to indicate to retrieveOrGenerateKey that a key should not
+// be generated.
+const KeyGeneration neverGen();
 
 bool isFsKeyringSupported(void);
 
@@ -58,9 +76,9 @@
 // In the latter case, the caller is responsible for dropping caches.
 bool evictKey(const std::string& mountpoint, const EncryptionPolicy& policy);
 
-bool retrieveKey(bool create_if_absent, const KeyAuthentication& key_authentication,
-                 const std::string& key_path, const std::string& tmp_path,
-                 const EncryptionOptions& options, KeyBuffer* key, bool keepOld = true);
+bool retrieveOrGenerateKey(const std::string& key_path, const std::string& tmp_path,
+                           const KeyAuthentication& key_authentication, const KeyGeneration& gen,
+                           KeyBuffer* key, bool keepOld = true);
 
 }  // namespace vold
 }  // namespace android
diff --git a/MetadataCrypt.cpp b/MetadataCrypt.cpp
index 106978e..938ba34 100644
--- a/MetadataCrypt.cpp
+++ b/MetadataCrypt.cpp
@@ -129,24 +129,8 @@
     LOG(INFO) << "Old Key deleted: " << dir;
 }
 
-static bool retrieveMetadataKey(bool create_if_absent, const std::string& key_path,
-                                const std::string& tmp_path, KeyBuffer* key, bool keepOld) {
-    if (pathExists(key_path)) {
-        LOG(DEBUG) << "Key exists, using: " << key_path;
-        if (!retrieveKey(key_path, kEmptyAuthentication, key, keepOld)) return false;
-    } else {
-        if (!create_if_absent) {
-            LOG(ERROR) << "No key found in " << key_path;
-            return false;
-        }
-        LOG(INFO) << "Creating new key in " << key_path;
-        if (!randomKey(key)) return false;
-        if (!storeKeyAtomically(key_path, tmp_path, kEmptyAuthentication, *key)) return false;
-    }
-    return true;
-}
-
-static bool read_key(const std::string& metadata_key_dir, bool create_if_absent, KeyBuffer* key) {
+static bool read_key(const std::string& metadata_key_dir, const KeyGeneration& gen,
+                     KeyBuffer* key) {
     if (metadata_key_dir.empty()) {
         LOG(ERROR) << "Failed to get metadata_key_dir";
         return false;
@@ -168,14 +152,14 @@
     Keymaster keymaster;
     if (pathExists(newKeyPath)) {
         if (!android::base::ReadFileToString(newKeyPath, &sKey))
-            LOG(ERROR) << "Failed to read old key: " << dir;
+            LOG(ERROR) << "Failed to read incomplete key: " << dir;
         else if (!keymaster.deleteKey(sKey))
-            LOG(ERROR) << "Old key deletion failed, continuing anyway: " << dir;
+            LOG(ERROR) << "Incomplete key deletion failed, continuing anyway: " << dir;
         else
             unlink(newKeyPath.c_str());
     }
     bool needs_cp = cp_needsCheckpoint();
-    if (!retrieveMetadataKey(create_if_absent, dir, temp, key, needs_cp)) return false;
+    if (!retrieveOrGenerateKey(dir, temp, kEmptyAuthentication, gen, key, needs_cp)) return false;
     if (needs_cp && pathExists(newKeyPath)) std::thread(commit_key, dir).detach();
     return true;
 }
@@ -283,8 +267,9 @@
         return false;
     }
 
+    auto gen = needs_encrypt ? makeGen(cipher) : neverGen();
     KeyBuffer key;
-    if (!read_key(data_rec->metadata_key_dir, needs_encrypt, &key)) return false;
+    if (!read_key(data_rec->metadata_key_dir, gen, &key)) return false;
 
     std::string crypto_blkdev;
     if (!create_crypto_blk_dev(kDmNameUserdata, data_rec->blk_device, is_legacy,
diff --git a/cryptfs.cpp b/cryptfs.cpp
index 38661f3..04497b0 100644
--- a/cryptfs.cpp
+++ b/cryptfs.cpp
@@ -74,6 +74,7 @@
 using android::fs_mgr::GetEntryForMountPoint;
 using android::vold::CryptoType;
 using android::vold::KeyBuffer;
+using android::vold::KeyGeneration;
 using namespace android::dm;
 using namespace std::chrono_literals;
 
@@ -320,6 +321,10 @@
     return crypto_type;
 }
 
+const KeyGeneration cryptfs_get_keygen() {
+    return makeGen(get_crypto_type());
+}
+
 /* Should we use keymaster? */
 static int keymaster_check_compatibility() {
     return keymaster_compatibility_cryptfs_scrypt();
@@ -469,10 +474,6 @@
     ftr->p_factor = pf;
 }
 
-size_t cryptfs_get_keysize() {
-    return get_crypto_type().get_keysize();
-}
-
 static uint64_t get_fs_size(const char* dev) {
     int fd, block_size;
     struct ext4_super_block sb;
diff --git a/cryptfs.h b/cryptfs.h
index b34a8d9..872806e 100644
--- a/cryptfs.h
+++ b/cryptfs.h
@@ -26,6 +26,7 @@
 #include <cutils/properties.h>
 
 #include "KeyBuffer.h"
+#include "KeyUtil.h"
 
 #define CRYPT_FOOTER_OFFSET 0x4000
 
@@ -73,7 +74,6 @@
 const char* cryptfs_get_password(void);
 void cryptfs_clear_password(void);
 int cryptfs_isConvertibleToFBE(void);
-
-size_t cryptfs_get_keysize();
+const android::vold::KeyGeneration cryptfs_get_keygen();
 
 #endif /* ANDROID_VOLD_CRYPTFS_H */
diff --git a/model/Disk.cpp b/model/Disk.cpp
index bfaf2cd..f92435d 100644
--- a/model/Disk.cpp
+++ b/model/Disk.cpp
@@ -16,6 +16,7 @@
 
 #include "Disk.h"
 #include "FsCrypt.h"
+#include "KeyUtil.h"
 #include "PrivateVolume.h"
 #include "PublicVolume.h"
 #include "Utils.h"
@@ -505,11 +506,12 @@
         return -EIO;
     }
 
-    std::string keyRaw;
-    if (ReadRandomBytes(cryptfs_get_keysize(), keyRaw) != OK) {
+    KeyBuffer key;
+    if (!generateStorageKey(cryptfs_get_keygen(), &key)) {
         LOG(ERROR) << "Failed to generate key";
         return -EIO;
     }
+    std::string keyRaw(key.begin(), key.end());
 
     std::string partGuid;
     StrToHex(partGuidRaw, partGuid);