diff --git a/Ext4Crypt.cpp b/Ext4Crypt.cpp
index 90fc411..28140d2 100644
--- a/Ext4Crypt.cpp
+++ b/Ext4Crypt.cpp
@@ -381,7 +381,9 @@
 }
 
 static bool ensure_policy(const std::string& raw_ref, const std::string& path) {
-    if (e4crypt_policy_ensure(path.c_str(), raw_ref.data(), raw_ref.size()) != 0) {
+    if (e4crypt_policy_ensure(path.c_str(),
+                              raw_ref.data(), raw_ref.size(),
+                              cryptfs_get_file_encryption_mode()) != 0) {
         LOG(ERROR) << "Failed to set policy on: " << path;
         return false;
     }
@@ -440,6 +442,13 @@
         return true;
     }
 
+    std::string mode_filename = std::string("/data") + e4crypt_key_mode;
+    std::string mode = cryptfs_get_file_encryption_mode();
+    if (!android::base::WriteStringToFile(mode, mode_filename)) {
+        PLOG(ERROR) << "Cannot save type";
+        return false;
+    }
+
     std::string device_key;
     if (path_exists(device_key_path)) {
         if (!android::vold::retrieveKey(device_key_path,