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;