3 files changed, 27 insertions, 9 deletions

diff --git a/crypto/ext4crypt/KeyStorage4.cpp b/crypto/ext4crypt/KeyStorage4.cpp
index cab88a19d..b91d6e46b 100644
--- a/crypto/ext4crypt/KeyStorage4.cpp
+++ b/crypto/ext4crypt/KeyStorage4.cpp
@@ -165,10 +165,28 @@ bool getEphemeralWrappedKey(km::KeyFormat format, KeyBuffer& kmKey, KeyBuffer* k
     std::string key_temp;
     Keymaster keymaster;
     if (!keymaster) return false;
-    if (!keymaster.exportKey(format, kmKey, "!", "!", &key_temp)) return false;
-    *key = KeyBuffer(key_temp.size());
-    memcpy(reinterpret_cast<void*>(key->data()), key_temp.c_str(), key->size());
-    return true;
+
+    //Export once, if upgrade needed, upgrade and export again
+    bool export_again = true;
+    while (export_again) {
+        export_again = false;
+        auto ret = keymaster.exportKey(format, kmKey, "!", "!", &key_temp);
+        if (ret == km::ErrorCode::OK) {
+            *key = KeyBuffer(key_temp.size());
+            memcpy(reinterpret_cast<void*>(key->data()), key_temp.c_str(), key->size());
+            return true;
+        }
+        if (ret != km::ErrorCode::KEY_REQUIRES_UPGRADE) return false;
+        LOG(DEBUG) << "Upgrading key";
+        std::string kmKeyStr(reinterpret_cast<const char*>(kmKey.data()), kmKey.size());
+        std::string newKey;
+        if (!keymaster.upgradeKey(kmKeyStr, km::AuthorizationSet(), &newKey)) return false;
+        memcpy(reinterpret_cast<void*>(kmKey.data()), newKey.c_str(), kmKey.size());
+        LOG(INFO) << "Key upgraded";
+        export_again = true;
+    }
+    //Should never come here
+    return false;
 }
 
 static std::pair<km::AuthorizationSet, km::HardwareAuthToken> beginParams(
diff --git a/crypto/ext4crypt/Keymaster4.cpp b/crypto/ext4crypt/Keymaster4.cpp
index e5c059a61..6507bb180 100644
--- a/crypto/ext4crypt/Keymaster4.cpp
+++ b/crypto/ext4crypt/Keymaster4.cpp
@@ -142,7 +142,7 @@ bool Keymaster::generateKey(const km::AuthorizationSet& inParams, std::string* k
     return true;
 }
 
-bool Keymaster::exportKey(km::KeyFormat format, KeyBuffer& kmKey, const std::string& clientId,
+km::ErrorCode Keymaster::exportKey(km::KeyFormat format, KeyBuffer& kmKey, const std::string& clientId,
                           const std::string& appData, std::string* key) {
     auto kmKeyBlob = km::support::blob2hidlVec(std::string(kmKey.data(), kmKey.size()));
     auto emptyAssign = NULL;
@@ -159,13 +159,13 @@ bool Keymaster::exportKey(km::KeyFormat format, KeyBuffer& kmKey, const std::str
     auto error = mDevice->exportKey(format, kmKeyBlob, kmClientId, kmAppData, 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/ext4crypt/Keymaster4.h b/crypto/ext4crypt/Keymaster4.h
index 14c230bfd..704cc2c70 100644
--- a/crypto/ext4crypt/Keymaster4.h
+++ b/crypto/ext4crypt/Keymaster4.h
@@ -103,7 +103,7 @@ class Keymaster {
     // Generate a key in the keymaster from the given params.
     bool generateKey(const km::AuthorizationSet& inParams, std::string* key);
     // Export a key from keymaster.
-    bool exportKey(km::KeyFormat format, KeyBuffer& kmKey, const std::string& clientId,
+    km::ErrorCode exportKey(km::KeyFormat format, KeyBuffer& kmKey, const std::string& clientId,
                    const std::string& appData, std::string* key);
     // If the keymaster supports it, permanently delete a key.
     bool deleteKey(const std::string& key);