diff --git a/keystore/Android.bp b/keystore/Android.bp
index cb736b3..68dde3d 100644
--- a/keystore/Android.bp
+++ b/keystore/Android.bp
@@ -84,6 +84,7 @@
     srcs: ["keystore_cli.cpp"],
     shared_libs: [
         "android.hardware.keymaster@3.0",
+        "android.hardware.keymaster@4.0",
         "libbinder",
         "libcrypto",
         "libcutils",
@@ -109,8 +110,8 @@
     srcs: ["keystore_cli_v2.cpp"],
     shared_libs: [
         "android.hardware.confirmationui@1.0",
-        "android.hardware.keymaster@3.0",
         "libbinder",
+        "android.hardware.keymaster@4.0",
         "libchrome",
         "libutils",
         "libhidlbase",
diff --git a/keystore/KeyStore.cpp b/keystore/KeyStore.cpp
index 144c1c7..0efc4a3 100644
--- a/keystore/KeyStore.cpp
+++ b/keystore/KeyStore.cpp
@@ -600,11 +600,6 @@
     }
 
     auto version = getDevice(SecurityLevel::TRUSTED_ENVIRONMENT)->halVersion();
-    if (version.error != ErrorCode::OK) {
-        ALOGE("Failed to get HAL version info");
-        return false;
-    }
-
     if (keyType == kRsaKeyType) return true;  // All versions support RSA
     return keyType == kEcKeyType && version.supportsEc;
 }
diff --git a/keystore/include/keystore/keymaster_types.h b/keystore/include/keystore/keymaster_types.h
index 62b43be..bd61294 100644
--- a/keystore/include/keystore/keymaster_types.h
+++ b/keystore/include/keystore/keymaster_types.h
@@ -83,6 +83,7 @@
 using keymaster::TAG_RSA_PUBLIC_EXPONENT;
 using keymaster::TAG_USAGE_EXPIRE_DATETIME;
 using keymaster::TAG_USER_AUTH_TYPE;
+using keymaster::TAG_USER_ID;
 using keymaster::TAG_USER_SECURE_ID;
 
 using keymaster::NullOr;
diff --git a/keystore/include/keystore/keystore_hidl_support.h b/keystore/include/keystore/keystore_hidl_support.h
index 7a3723e..781b153 100644
--- a/keystore/include/keystore/keystore_hidl_support.h
+++ b/keystore/include/keystore/keystore_hidl_support.h
@@ -26,11 +26,16 @@
 #include <android/hardware/keymaster/3.0/IKeymasterDevice.h>
 #include <hardware/hw_auth_token.h>
 #include <hidl/Status.h>
+#include <keymasterV4_0/keymaster_utils.h>
 
 #include <keystore/keymaster_types.h>
 
 namespace keystore {
 
+using android::hardware::keymaster::V4_0::support::blob2hidlVec;
+using android::hardware::keymaster::V4_0::support::hidlVec2AuthToken;
+using android::hardware::keymaster::V4_0::support::authToken2HidlVec;
+
 inline static std::ostream& formatArgs(std::ostream& out) {
     return out;
 }
@@ -69,32 +74,6 @@
 #define KS_HANDLE_HIDL_ERROR(rc)                                                                   \
     ::keystore::ksHandleHidlError(rc, __FILE__, ":", __LINE__, ":", __PRETTY_FUNCTION__)
 
-inline static hidl_vec<uint8_t> blob2hidlVec(const uint8_t* data, const size_t length,
-                                             bool inPlace = true) {
-    hidl_vec<uint8_t> result;
-    if (inPlace)
-        result.setToExternal(const_cast<unsigned char*>(data), length);
-    else {
-        result.resize(length);
-        memcpy(&result[0], data, length);
-    }
-    return result;
-}
-
-inline static hidl_vec<uint8_t> blob2hidlVec(const std::string& value) {
-    hidl_vec<uint8_t> result;
-    result.setToExternal(
-        reinterpret_cast<uint8_t*>(const_cast<std::string::value_type*>(value.data())),
-        static_cast<size_t>(value.size()));
-    return result;
-}
-
-inline static hidl_vec<uint8_t> blob2hidlVec(const std::vector<uint8_t>& blob) {
-    hidl_vec<uint8_t> result;
-    result.setToExternal(const_cast<uint8_t*>(blob.data()), static_cast<size_t>(blob.size()));
-    return result;
-}
-
 template <typename T, typename OutIter>
 inline static OutIter copy_bytes_to_iterator(const T& value, OutIter dest) {
     const uint8_t* value_ptr = reinterpret_cast<const uint8_t*>(&value);
@@ -159,59 +138,6 @@
     return token;
 }
 
-inline static hidl_vec<uint8_t> authToken2HidlVec(const HardwareAuthToken& token) {
-    static_assert(1 /* version size */ + sizeof(token.challenge) + sizeof(token.userId) +
-                          sizeof(token.authenticatorId) + sizeof(token.authenticatorType) +
-                          sizeof(token.timestamp) + kHmacSize ==
-                      sizeof(hw_auth_token_t),
-                  "HardwareAuthToken content size does not match hw_auth_token_t size");
-
-    hidl_vec<uint8_t> result;
-    result.resize(sizeof(hw_auth_token_t));
-    auto pos = result.begin();
-    *pos++ = 0;  // Version byte
-    pos = copy_bytes_to_iterator(token.challenge, pos);
-    pos = copy_bytes_to_iterator(token.userId, pos);
-    pos = copy_bytes_to_iterator(token.authenticatorId, pos);
-    auto auth_type = htonl(static_cast<uint32_t>(token.authenticatorType));
-    pos = copy_bytes_to_iterator(auth_type, pos);
-    auto timestamp = htonq(token.timestamp);
-    pos = copy_bytes_to_iterator(timestamp, pos);
-    if (token.mac.size() != kHmacSize) {
-        std::fill(pos, pos + kHmacSize, 0);
-    } else {
-        std::copy(token.mac.begin(), token.mac.end(), pos);
-    }
-
-    return result;
-}
-
-inline static HardwareAuthToken hidlVec2AuthToken(const hidl_vec<uint8_t>& buffer) {
-    HardwareAuthToken token;
-    static_assert(1 /* version size */ + sizeof(token.challenge) + sizeof(token.userId) +
-                          sizeof(token.authenticatorId) + sizeof(token.authenticatorType) +
-                          sizeof(token.timestamp) + kHmacSize ==
-                      sizeof(hw_auth_token_t),
-                  "HardwareAuthToken content size does not match hw_auth_token_t size");
-
-    if (buffer.size() != sizeof(hw_auth_token_t)) return {};
-
-    auto pos = buffer.begin();
-    ++pos;  // skip first byte
-    pos = copy_bytes_from_iterator(&token.challenge, pos);
-    pos = copy_bytes_from_iterator(&token.userId, pos);
-    pos = copy_bytes_from_iterator(&token.authenticatorId, pos);
-    pos = copy_bytes_from_iterator(&token.authenticatorType, pos);
-    token.authenticatorType = static_cast<HardwareAuthenticatorType>(
-        ntohl(static_cast<uint32_t>(token.authenticatorType)));
-    pos = copy_bytes_from_iterator(&token.timestamp, pos);
-    token.timestamp = ntohq(token.timestamp);
-    token.mac.resize(kHmacSize);
-    std::copy(pos, pos + kHmacSize, token.mac.data());
-
-    return token;
-}
-
 inline std::string hidlVec2String(const hidl_vec<uint8_t>& value) {
     return std::string(reinterpret_cast<const std::string::value_type*>(&value[0]), value.size());
 }
diff --git a/keystore/key_store_service.cpp b/keystore/key_store_service.cpp
index 89c31a5..c3f661c 100644
--- a/keystore/key_store_service.cpp
+++ b/keystore/key_store_service.cpp
@@ -372,6 +372,7 @@
         return Status::ok();
     }
 
+    enforcement_policy.set_device_locked(true, userId);
     mKeyStore->lock(userId);
     *aidl_return = static_cast<int32_t>(ResponseCode::NO_ERROR);
     return Status::ok();
@@ -400,6 +401,7 @@
         return Status::ok();
     }
 
+    enforcement_policy.set_device_locked(false, userId);
     const String8 password8(pw);
     // read master key, decrypt with password, initialize mMasterKey*.
     *aidl_return = static_cast<int32_t>(mKeyStore->readMasterKey(password8, userId));
@@ -1466,7 +1468,7 @@
 }
 
 Status KeyStoreService::addAuthToken(const ::std::vector<uint8_t>& authTokenAsVector,
-                                     int32_t* aidl_return) {
+                                     int32_t android_uid, int32_t* aidl_return) {
 
     // TODO(swillden): When gatekeeper and fingerprint are ready, this should be updated to
     // receive a HardwareAuthToken, rather than an opaque byte array.
@@ -1488,6 +1490,8 @@
         return Status::ok();
     }
 
+    enforcement_policy.set_device_locked(false, android_uid);
+
     mAuthTokenTable.AddAuthenticationToken(hidlVec2AuthToken(hidl_vec<uint8_t>(authTokenAsVector)));
     *aidl_return = static_cast<int32_t>(ResponseCode::NO_ERROR);
     return Status::ok();
diff --git a/keystore/key_store_service.h b/keystore/key_store_service.h
index ce809f8..ea4a089 100644
--- a/keystore/key_store_service.h
+++ b/keystore/key_store_service.h
@@ -146,7 +146,7 @@
     ::android::binder::Status isOperationAuthorized(const ::android::sp<::android::IBinder>& token,
                                                     bool* _aidl_return) override;
     ::android::binder::Status addAuthToken(const ::std::vector<uint8_t>& authToken,
-                                           int32_t* _aidl_return) override;
+                                           int32_t android_uid, int32_t* _aidl_return) override;
     ::android::binder::Status onUserAdded(int32_t userId, int32_t parentId,
                                           int32_t* _aidl_return) override;
     ::android::binder::Status onUserRemoved(int32_t userId, int32_t* _aidl_return) override;
diff --git a/keystore/keymaster_enforcement.cpp b/keystore/keymaster_enforcement.cpp
index d78a5a6..690927c 100644
--- a/keystore/keymaster_enforcement.cpp
+++ b/keystore/keymaster_enforcement.cpp
@@ -223,6 +223,8 @@
     bool caller_nonce_authorized_by_key = false;
     bool authentication_required = false;
     bool auth_token_matched = false;
+    bool unlocked_device_required = false;
+    int32_t user_id = -1;
 
     for (auto& param : auth_set) {
 
@@ -282,10 +284,18 @@
             }
             break;
 
+        case Tag::USER_ID:
+            user_id = authorizationValue(TAG_USER_ID, param).value();
+            break;
+
         case Tag::CALLER_NONCE:
             caller_nonce_authorized_by_key = true;
             break;
 
+        case Tag::UNLOCKED_DEVICE_REQUIRED:
+            unlocked_device_required = true;
+            break;
+
         /* Tags should never be in key auths. */
         case Tag::INVALID:
         case Tag::ROOT_OF_TRUST:
@@ -356,6 +366,10 @@
         }
     }
 
+    if (unlocked_device_required && is_device_locked(user_id)) {
+        return ErrorCode::DEVICE_LOCKED;
+    }
+
     if (authentication_required && !auth_token_matched) {
         ALOGE("Auth required but no matching auth token found");
         return ErrorCode::KEY_USER_NOT_AUTHENTICATED;
diff --git a/keystore/keymaster_enforcement.h b/keystore/keymaster_enforcement.h
index d7b27fc..6e6c54f 100644
--- a/keystore/keymaster_enforcement.h
+++ b/keystore/keymaster_enforcement.h
@@ -142,6 +142,11 @@
      */
     virtual bool ValidateTokenSignature(const HardwareAuthToken& token) const = 0;
 
+    /*
+     * Returns true if the device screen is currently locked for the specified user.
+     */
+    virtual bool is_device_locked(int32_t userId) const = 0;
+
   private:
     ErrorCode AuthorizeUpdateOrFinish(const AuthorizationSet& auth_set,
                                       const HardwareAuthToken& auth_token, uint64_t op_handle);
diff --git a/keystore/keystore_keymaster_enforcement.h b/keystore/keystore_keymaster_enforcement.h
index 3cdf649..e114ea9 100644
--- a/keystore/keystore_keymaster_enforcement.h
+++ b/keystore/keystore_keymaster_enforcement.h
@@ -84,6 +84,19 @@
         // signing key. Assume the token is good.
         return true;
     }
+
+    bool is_device_locked(int32_t userId) const override {
+        // If we haven't had a set call for this user yet, assume the device is locked.
+        if (mIsDeviceLockedForUser.count(userId) == 0) return true;
+        return mIsDeviceLockedForUser.find(userId)->second;
+    }
+
+    void set_device_locked(bool isLocked, int32_t userId) {
+        mIsDeviceLockedForUser[userId] = isLocked;
+    }
+
+  private:
+    std::map<int32_t, bool> mIsDeviceLockedForUser;
 };
 
 } // namespace keystore
diff --git a/keystore/keystore_main.cpp b/keystore/keystore_main.cpp
index e1fdd3f..1ec32dd 100644
--- a/keystore/keystore_main.cpp
+++ b/keystore/keystore_main.cpp
@@ -72,7 +72,7 @@
                               << Wrapper::WrappedIKeymasterDevice::descriptor
                               << "\" with interface name \"" << name << "\"";
 
-                sp<Keymaster> kmDevice(new Wrapper(device));
+                sp<Keymaster> kmDevice(new Wrapper(device, name));
                 auto halVersion = kmDevice->halVersion();
                 SecurityLevel securityLevel = halVersion.securityLevel;
                 LOG(INFO) << "found " << Wrapper::WrappedIKeymasterDevice::descriptor
@@ -175,7 +175,7 @@
     if (!result[SecurityLevel::SOFTWARE]) {
         auto fbdev = android::keystore::makeSoftwareKeymasterDevice();
         CHECK(fbdev.get()) << "Unable to create Software Keymaster Device";
-        result[SecurityLevel::SOFTWARE] = new Keymaster3(fbdev);
+        result[SecurityLevel::SOFTWARE] = new Keymaster3(fbdev, "Software");
     }
     return result;
 }
@@ -197,8 +197,6 @@
     CHECK(configure_selinux() != -1) << "Failed to configure SELinux.";
 
     auto halVersion = kmDevices[SecurityLevel::TRUSTED_ENVIRONMENT]->halVersion();
-    CHECK(halVersion.error == keystore::ErrorCode::OK)
-        << "Error " << toString(halVersion.error) << " getting HAL version";
 
     // If the hardware is keymaster 2.0 or higher we will not allow the fallback device for import
     // or generation of keys. The fallback device is only used for legacy keys present on the
diff --git a/keystore/tests/auth_token_formatting_test.cpp b/keystore/tests/auth_token_formatting_test.cpp
index 2677718..0ecc4cc 100644
--- a/keystore/tests/auth_token_formatting_test.cpp
+++ b/keystore/tests/auth_token_formatting_test.cpp
@@ -19,6 +19,7 @@
 #include <endian.h>
 #include <hidl/HidlSupport.h>
 #include <keymaster/logger.h>
+#include <keymasterV4_0/keymaster_utils.h>
 
 #include <keystore/keymaster_types.h>
 #include <keystore/keystore_hidl_support.h>