diff --git a/keystore/user_state.cpp b/keystore/user_state.cpp
index bc3f6d9..8d993e2 100644
--- a/keystore/user_state.cpp
+++ b/keystore/user_state.cpp
@@ -140,10 +140,13 @@
 }
 
 ResponseCode UserState::writeMasterKey(const android::String8& pw) {
-    std::vector<uint8_t> passwordKey(MASTER_KEY_SIZE_BYTES);
+    std::vector<uint8_t> passwordKey(mMasterKey.size());
     generateKeyFromPassword(passwordKey, pw, mSalt);
-    Blob masterKeyBlob(mMasterKey.data(), mMasterKey.size(), mSalt, sizeof(mSalt),
-                       TYPE_MASTER_KEY_AES256);
+    auto blobType = TYPE_MASTER_KEY_AES256;
+    if (mMasterKey.size() == kAes128KeySizeBytes) {
+        blobType = TYPE_MASTER_KEY;
+    }
+    Blob masterKeyBlob(mMasterKey.data(), mMasterKey.size(), mSalt, sizeof(mSalt), blobType);
     auto lockedEntry = LockedKeyBlobEntry::get(mMasterKeyEntry);
     return lockedEntry.writeBlobs(masterKeyBlob, {}, passwordKey, STATE_NO_ERROR);
 }
@@ -174,7 +177,7 @@
 
     size_t masterKeySize = MASTER_KEY_SIZE_BYTES;
     if (rawBlob.type == TYPE_MASTER_KEY) {
-        masterKeySize = SHA1_DIGEST_SIZE_BYTES;
+        masterKeySize = kAes128KeySizeBytes;
     }
 
     std::vector<uint8_t> passwordKey(masterKeySize);
@@ -263,7 +266,7 @@
     const EVP_MD* digest = EVP_sha256();
 
     // SHA1 was used prior to increasing the key size
-    if (key.size() == SHA1_DIGEST_SIZE_BYTES) {
+    if (key.size() == kAes128KeySizeBytes) {
         digest = EVP_sha1();
     }