diff --git a/crypto/ext4crypt/KeyStorage4.h b/crypto/ext4crypt/KeyStorage4.h
index 35ff65e..f74865f 100644
--- a/crypto/ext4crypt/KeyStorage4.h
+++ b/crypto/ext4crypt/KeyStorage4.h
@@ -17,13 +17,17 @@
 #ifndef ANDROID_TWRP_KEYSTORAGE_H
 #define ANDROID_TWRP_KEYSTORAGE_H
 
+#include "Keymaster4.h"
 #include "KeyBuffer.h"
+#include <ext4_utils/ext4_crypt.h>
 
 #include <string>
 
 namespace android {
 namespace vold {
 
+namespace km = ::android::hardware::keymaster::V4_0;
+
 // Represents the information needed to decrypt a disk encryption key.
 // If "token" is nonempty, it is passed in as a required Gatekeeper auth token.
 // If "token" and "secret" are nonempty, "secret" is appended to the application-specific
@@ -39,6 +43,12 @@
     const std::string secret;
 };
 
+enum class KeyType {
+    DE_SYS,
+    DE_USER,
+    CE_USER
+};
+
 extern const KeyAuthentication kEmptyAuthentication;
 
 // Checks if path "path" exists.
@@ -67,6 +77,9 @@
 bool destroyKey(const std::string& dir);
 
 bool runSecdiscardSingle(const std::string& file);
+
+bool generateWrappedKey(userid_t user_id, KeyType key_type, KeyBuffer* key);
+bool getEphemeralWrappedKey(km::KeyFormat format, KeyBuffer& kmKey, KeyBuffer* key);
 }  // namespace vold
 }  // namespace android