diff --git a/keystore/java/android/security/KeyStore.java b/keystore/java/android/security/KeyStore.java
index 646aa13664c48..9866c30537144 100644
--- a/keystore/java/android/security/KeyStore.java
+++ b/keystore/java/android/security/KeyStore.java
@@ -1067,6 +1067,17 @@ public class KeyStore {
         return onUserPasswordChanged(UserHandle.getUserId(Process.myUid()), newPassword);
     }
 
+    /**
+     * Notify keystore about the latest user locked state. This is to support keyguard-bound key.
+     */
+    public void onUserLockedStateChanged(int userHandle, boolean locked) {
+        try {
+            mBinder.onKeyguardVisibilityChanged(locked, userHandle);
+        } catch (RemoteException e) {
+            Log.w(TAG, "Failed to update user locked state " + userHandle, e);
+        }
+    }
+
     private class KeyAttestationCallbackResult {
         private KeystoreResponse keystoreResponse;
         private KeymasterCertificateChain certificateChain;
diff --git a/services/core/java/com/android/server/policy/keyguard/KeyguardStateMonitor.java b/services/core/java/com/android/server/policy/keyguard/KeyguardStateMonitor.java
index f78d2639df1a8..add0b01f18799 100644
--- a/services/core/java/com/android/server/policy/keyguard/KeyguardStateMonitor.java
+++ b/services/core/java/com/android/server/policy/keyguard/KeyguardStateMonitor.java
@@ -19,8 +19,6 @@ package com.android.server.policy.keyguard;
 import android.app.ActivityManager;
 import android.content.Context;
 import android.os.RemoteException;
-import android.os.ServiceManager;
-import android.security.keystore.IKeystoreService;
 import android.util.Slog;
 
 import com.android.internal.policy.IKeyguardService;
@@ -53,16 +51,11 @@ public class KeyguardStateMonitor extends IKeyguardStateCallback.Stub {
     private final LockPatternUtils mLockPatternUtils;
     private final StateCallback mCallback;
 
-    IKeystoreService mKeystoreService;
-
     public KeyguardStateMonitor(Context context, IKeyguardService service, StateCallback callback) {
         mLockPatternUtils = new LockPatternUtils(context);
         mCurrentUserId = ActivityManager.getCurrentUser();
         mCallback = callback;
 
-        mKeystoreService = IKeystoreService.Stub.asInterface(ServiceManager
-                .getService("android.security.keystore"));
-
         try {
             service.addStateMonitorCallback(this);
         } catch (RemoteException e) {
@@ -95,23 +88,6 @@ public class KeyguardStateMonitor extends IKeyguardStateCallback.Stub {
         mIsShowing = showing;
 
         mCallback.onShowingChanged();
-        int retry = 2;
-        while (retry > 0) {
-            try {
-                mKeystoreService.onKeyguardVisibilityChanged(showing, mCurrentUserId);
-                break;
-            } catch (RemoteException e) {
-                if (retry == 2) {
-                    Slog.w(TAG, "Error informing keystore of screen lock. Keystore may have died"
-                            + " -> refreshing service token and retrying");
-                    mKeystoreService = IKeystoreService.Stub.asInterface(ServiceManager
-                            .getService("android.security.keystore"));
-                } else {
-                    Slog.e(TAG, "Error informing keystore of screen lock after retrying once", e);
-                }
-                --retry;
-            }
-        }
     }
 
     @Override // Binder interface
@@ -123,10 +99,6 @@ public class KeyguardStateMonitor extends IKeyguardStateCallback.Stub {
         mCurrentUserId = userId;
     }
 
-    private synchronized int getCurrentUser() {
-        return mCurrentUserId;
-    }
-
     @Override // Binder interface
     public void onInputRestrictedStateChanged(boolean inputRestricted) {
         mInputRestricted = inputRestricted;
diff --git a/services/core/java/com/android/server/trust/TrustManagerService.java b/services/core/java/com/android/server/trust/TrustManagerService.java
index 7408dd40b5cad..5f5cd3c461170 100644
--- a/services/core/java/com/android/server/trust/TrustManagerService.java
+++ b/services/core/java/com/android/server/trust/TrustManagerService.java
@@ -53,6 +53,7 @@ import android.os.SystemClock;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.provider.Settings;
+import android.security.KeyStore;
 import android.service.trust.TrustAgentService;
 import android.text.TextUtils;
 import android.util.ArrayMap;
@@ -135,6 +136,33 @@ public class TrustManagerService extends SystemService {
     @GuardedBy("mUserIsTrusted")
     private final SparseBooleanArray mUserIsTrusted = new SparseBooleanArray();
 
+    /**
+     * Stores the locked state for users on the device. There are three different type of users
+     * which are handled slightly differently:
+     * <ul>
+     *  <li> Users with real keyguard
+     *  These are users who can be switched to ({@link UserInfo#supportsSwitchToByUser()}). Their
+     *  locked state is derived by a combination of user secure state, keyguard state, trust agent
+     *  decision and biometric authentication result. These are updated via
+     *  {@link #refreshDeviceLockedForUser(int)} and result stored in {@link #mDeviceLockedForUser}.
+     *  <li> Managed profiles with unified challenge
+     *  Managed profile with unified challenge always shares the same locked state as their parent,
+     *  so their locked state is not recorded in  {@link #mDeviceLockedForUser}. Instead,
+     *  {@link ITrustManager#isDeviceLocked(int)} always resolves their parent user handle and
+     *  queries its locked state instead.
+     *  <li> Managed profiles with separate challenge
+     *  Locked state for profile with separate challenge is determined by other parts of the
+     *  framework (mostly PowerManager) and pushed to TrustManagerService via
+     *  {@link ITrustManager#setDeviceLockedForUser(int, boolean)}. Although in a corner case when
+     *  the profile has a separate but empty challenge, setting its {@link #mDeviceLockedForUser} to
+     *  {@code false} is actually done by {@link #refreshDeviceLockedForUser(int)}.
+     * </ul>
+     * TODO: Rename {@link ITrustManager#setDeviceLockedForUser(int, boolean)} to
+     * {@code setDeviceLockedForProfile} to better reflect its purpose. Unifying
+     * {@code setDeviceLockedForProfile} and {@link #setDeviceLockedForUser} would also be nice.
+     * At the moment they both update {@link #mDeviceLockedForUser} but have slightly different
+     * side-effects: one notifies trust agents while the other sends out a broadcast.
+     */
     @GuardedBy("mDeviceLockedForUser")
     private final SparseBooleanArray mDeviceLockedForUser = new SparseBooleanArray();
 
@@ -601,6 +629,10 @@ public class TrustManagerService extends SystemService {
         }
     }
 
+    /**
+     * Update the user's locked state. Only applicable to users with a real keyguard
+     * ({@link UserInfo#supportsSwitchToByUser}) and unsecured managed profiles.
+     */
     private void refreshDeviceLockedForUser(int userId) {
         if (userId != UserHandle.USER_ALL && userId < UserHandle.USER_SYSTEM) {
             Log.e(TAG, "refreshDeviceLockedForUser(userId=" + userId + "): Invalid user handle,"
@@ -661,6 +693,15 @@ public class TrustManagerService extends SystemService {
         }
         if (changed) {
             dispatchDeviceLocked(userId, locked);
+
+            KeyStore.getInstance().onUserLockedStateChanged(userId, locked);
+            // Also update the user's profiles who have unified challenge, since they
+            // share the same unlocked state (see {@link #isDeviceLocked(int)})
+            for (int profileHandle : mUserManager.getEnabledProfileIds(userId)) {
+                if (mLockPatternUtils.isManagedProfileWithUnifiedChallenge(profileHandle)) {
+                    KeyStore.getInstance().onUserLockedStateChanged(profileHandle, locked);
+                }
+            }
         }
     }
 
@@ -1194,6 +1235,10 @@ public class TrustManagerService extends SystemService {
             return "0x" + Integer.toHexString(i);
         }
 
+        /**
+         * Changes the lock status for the given user. This is only applicable to managed profiles,
+         * other users should be handled by Keyguard.
+         */
         @Override
         public void setDeviceLockedForUser(int userId, boolean locked) {
             enforceReportPermission();
@@ -1204,6 +1249,9 @@ public class TrustManagerService extends SystemService {
                     synchronized (mDeviceLockedForUser) {
                         mDeviceLockedForUser.put(userId, locked);
                     }
+
+                    KeyStore.getInstance().onUserLockedStateChanged(userId, locked);
+
                     if (locked) {
                         try {
                             ActivityManager.getService().notifyLockedProfile(userId);