diff --git a/crypto/fscrypt/FsCrypt.cpp b/crypto/fscrypt/FsCrypt.cpp
index 418164e..3f388d3 100755
--- a/crypto/fscrypt/FsCrypt.cpp
+++ b/crypto/fscrypt/FsCrypt.cpp
@@ -201,7 +201,7 @@
     auto const paths = get_ce_key_paths(directory_path);
     for (auto const ce_key_path : paths) {
         LOG(INFO) << "Trying user CE key " << ce_key_path;
-        if (retrieveKey(ce_key_path, auth, ce_key)) {
+        if (retrieveKey(ce_key_path, auth, ce_key, true)) {
             LOG(INFO) << "Successfully retrieved key";
             fixate_user_ce_key(directory_path, ce_key_path, paths);
             return true;
@@ -409,7 +409,7 @@
         userid_t user_id = std::stoi(entry->d_name);
         auto key_path = de_dir + "/" + entry->d_name;
         KeyBuffer de_key;
-        if (!retrieveKey(key_path, kEmptyAuthentication, &de_key)) return false;
+        if (!retrieveKey(key_path, kEmptyAuthentication, &de_key, true)) return false;
         EncryptionPolicy de_policy;
         if (!install_storage_key(DATA_MNT_POINT, options, de_key, &de_policy)) return false;
         auto ret = s_de_policies.insert({user_id, de_policy});
@@ -707,12 +707,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 (retrieveKey(ce_key_current_path, retrieve_auth, &ce_key)) {
+    if (retrieveKey(ce_key_current_path, retrieve_auth, &ce_key, true)) {
         LOG(INFO) << "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 (retrieveKey(ce_key_current_path, kEmptyAuthentication, &ce_key)) {
+    } else if (retrieveKey(ce_key_current_path, kEmptyAuthentication, &ce_key, true)) {
         LOG(INFO) << "Successfully retrieved key with empty auth";
     } else {
         LOG(ERROR) << "Failed to retrieve key for user " << user_id;
diff --git a/crypto/fscrypt/KeyStorage.cpp b/crypto/fscrypt/KeyStorage.cpp
index ad9d491..83800ce 100755
--- a/crypto/fscrypt/KeyStorage.cpp
+++ b/crypto/fscrypt/KeyStorage.cpp
@@ -151,8 +151,20 @@
     if (!keymaster) return false;
     std::string key_temp;
 
-    if (!keymaster.exportKey(kmKey, &key_temp)) return false;
-    *key = KeyBuffer(key_temp.size());
+    auto ret = keymaster.exportKey(kmKey, &key_temp);
+    if (ret != km::ErrorCode::OK) {
+        if (ret == km::ErrorCode::KEY_REQUIRES_UPGRADE) {
+           std::string kmKeyStr(reinterpret_cast<const char*>(kmKey.data()), kmKey.size());
+           std::string Keystr;
+           if (!keymaster.upgradeKey(kmKeyStr, km::AuthorizationSet(), &Keystr)) return false;
+           KeyBuffer upgradedKey = KeyBuffer(Keystr.size());
+           memcpy(reinterpret_cast<void*>(upgradedKey.data()), Keystr.c_str(), upgradedKey.size());
+           ret = keymaster.exportKey(upgradedKey, &key_temp);
+           if (ret != km::ErrorCode::OK) return false;
+        } else {
+           return false;
+        }
+    }    *key = KeyBuffer(key_temp.size());
     memcpy(reinterpret_cast<void*>(key->data()), key_temp.c_str(), key->size());
     return true;
 }
@@ -164,7 +176,7 @@
                             .Authorization(km::TAG_APPLICATION_ID, km::support::blob2hidlVec(appId));
     km::HardwareAuthToken authToken;
     if (!auth.token.empty()) {
-        LOG(DEBUG) << "Supplying auth token to Keymaster";
+        LOG(INFO) << "Supplying auth token to Keymaster";
         authToken = km::support::hidlVec2AuthToken(km::support::blob2hidlVec(auth.token));
     }
     return {paramBuilder, authToken};
@@ -243,21 +255,21 @@
         LOG(DEBUG) << "Upgrading key: " << dir;
         std::string newKey;
         if (!keymaster.upgradeKey(kmKey, keyParams, &newKey)) return KeymasterOperation();
-        // auto newKeyPath = dir + "/" + kFn_keymaster_key_blob_upgraded;
-        // if (!writeStringToFile(newKey, newKeyPath)) return KeymasterOperation();
-        // if (!keepOld) {
-        //     if (rename(newKeyPath.c_str(), kmKeyPath.c_str()) != 0) {
-        //         PLOG(ERROR) << "Unable to move upgraded key to location: " << kmKeyPath;
-        //         return KeymasterOperation();
-        //     }
-        //     if (!::FsyncDirectory(dir)) {
-        //         LOG(ERROR) << "Key dir sync failed: " << dir;
-        //         return KeymasterOperation();
-        //     }
-        //     if (!kmDeleteKey(keymaster, kmKey)) {
-        //         LOG(ERROR) << "Key deletion failed during upgrade, continuing anyway: " << dir;
-        //     }
-        // }
+        auto newKeyPath = dir + "/" + kFn_keymaster_key_blob_upgraded;
+        if (!writeStringToFile(newKey, newKeyPath)) return KeymasterOperation();
+        if (!keepOld) {
+            if (rename(newKeyPath.c_str(), kmKeyPath.c_str()) != 0) {
+                PLOG(ERROR) << "Unable to move upgraded key to location: " << kmKeyPath;
+                return KeymasterOperation();
+            }
+            if (!::FsyncDirectory(dir)) {
+                LOG(ERROR) << "Key dir sync failed: " << dir;
+                return KeymasterOperation();
+            }
+            if (!kmDeleteKey(keymaster, kmKey)) {
+                LOG(ERROR) << "Key deletion failed during upgrade, continuing anyway: " << dir;
+            }
+        }
         kmKey = newKey;
         LOG(INFO) << "Key upgraded: " << dir;
     }
diff --git a/crypto/fscrypt/KeyStorage.h b/crypto/fscrypt/KeyStorage.h
index b88cdc6..2c88820 100755
--- a/crypto/fscrypt/KeyStorage.h
+++ b/crypto/fscrypt/KeyStorage.h
@@ -60,7 +60,7 @@
 
 // Retrieve the key from the named directory.
 bool retrieveKey(const std::string& dir, const KeyAuthentication& auth, KeyBuffer* key,
-                 bool keepOld = false);
+                 bool keepOld);
 
 // Securely destroy the key stored in the named directory and delete the directory.
 bool destroyKey(const std::string& dir);
diff --git a/crypto/fscrypt/Keymaster.cpp b/crypto/fscrypt/Keymaster.cpp
index 27353f1..c5b7bc3 100755
--- a/crypto/fscrypt/Keymaster.cpp
+++ b/crypto/fscrypt/Keymaster.cpp
@@ -136,7 +136,7 @@
     return true;
 }
 
-bool Keymaster::exportKey(const KeyBuffer& kmKey, std::string* key) {
+km::ErrorCode Keymaster::exportKey(const KeyBuffer& kmKey, std::string* key) {
     auto kmKeyBlob = km::support::blob2hidlVec(std::string(kmKey.data(), kmKey.size()));
     km::ErrorCode km_error;
     auto hidlCb = [&](km::ErrorCode ret, const hidl_vec<uint8_t>& exportedKeyBlob) {
@@ -148,13 +148,13 @@
     auto error = mDevice->exportKey(km::KeyFormat::RAW, kmKeyBlob, {}, {}, hidlCb);
     if (!error.isOk()) {
         LOG(ERROR) << "export_key failed: " << error.description();
-        return false;
+        return km::ErrorCode::UNKNOWN_ERROR;
     }
     if (km_error != km::ErrorCode::OK) {
         LOG(ERROR) << "export_key failed, code " << int32_t(km_error);
-        return false;
+        return km_error;
     }
-    return true;
+    return km::ErrorCode::OK;
 }
 
 bool Keymaster::deleteKey(const std::string& key) {
diff --git a/crypto/fscrypt/Keymaster.h b/crypto/fscrypt/Keymaster.h
index c3da4c8..6f74db4 100644
--- a/crypto/fscrypt/Keymaster.h
+++ b/crypto/fscrypt/Keymaster.h
@@ -112,7 +112,7 @@
     // Generate a key in the keymaster from the given params.
     bool generateKey(const km::AuthorizationSet& inParams, std::string* key);
     // Exports a keymaster key with STORAGE_KEY tag wrapped with a per-boot ephemeral key
-    bool exportKey(const KeyBuffer& kmKey, std::string* key);
+    km::ErrorCode exportKey(const KeyBuffer& kmKey, std::string* key);
     // If the keymaster supports it, permanently delete a key.
     bool deleteKey(const std::string& key);
     // Replace stored key blob in response to KM_ERROR_KEY_REQUIRES_UPGRADE.
diff --git a/crypto/fscrypt/MetadataCrypt.cpp b/crypto/fscrypt/MetadataCrypt.cpp
index d875ff0..1504367 100755
--- a/crypto/fscrypt/MetadataCrypt.cpp
+++ b/crypto/fscrypt/MetadataCrypt.cpp
@@ -160,6 +160,7 @@
      * or we rebooted before commiting the keys in a freak accident.
      * Either way, we can re-upgrade the key if we need to.
      */
+
     Keymaster keymaster;
     if (pathExists(newKeyPath)) {
         if (!android::base::ReadFileToString(newKeyPath, &sKey))
@@ -170,7 +171,7 @@
             unlink(newKeyPath.c_str());
     }
     bool needs_cp = cp_needsCheckpoint();
-    if (!retrieveOrGenerateKey(dir, temp, kEmptyAuthentication, gen, key, needs_cp)) return false;
+    if (!retrieveOrGenerateKey(dir, temp, kEmptyAuthentication, gen, key, true)) return false;
     if (needs_cp && pathExists(newKeyPath)) std::thread(commit_key, dir).detach();
     return true;
 }
diff --git a/partitionmanager.cpp b/partitionmanager.cpp
index 158d916..fe0e3b8 100755
--- a/partitionmanager.cpp
+++ b/partitionmanager.cpp
@@ -450,6 +450,7 @@
 				}
 			}
 		} else {
+			LOGINFO("FBE setup failed. Trying FDE...");
 			Set_FDE_Encrypt_Status();
 			int password_type = cryptfs_get_password_type();
 			if (password_type == CRYPT_TYPE_DEFAULT) {