RESTRICT AUTOMERGE
Update keyguard locked state from TrustManagerService

TrustManagerService holds the ground truth about whether a user is
locked or not, so update keystore using the information there,
instead of doing it from KeyguardStateMonitor. This fixes the issue
of work profile locked state not being correctly pushed to keystore.

Note: since this change is likely to be backported as a security
patch, I'm refraining from doing major refactoring right now.

Bug: 141329041
Bug: 144430870
Test: manually with KeyPairSampleApp
Change-Id: I3472ece73d573a775345ebcceeeb2cc460374c9b
(cherry picked from commit f9418dbb2c2469dd271e4aebefda5b6a4b485f3a)
diff --git a/keystore/java/android/security/KeyStore.java b/keystore/java/android/security/KeyStore.java
index 646aa13..9866c30 100644
--- a/keystore/java/android/security/KeyStore.java
+++ b/keystore/java/android/security/KeyStore.java
@@ -1067,6 +1067,17 @@
         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 f78d263..add0b01 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 @@
 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 @@
     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 @@
         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 @@
         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 7408dd4..5f5cd3c 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.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 @@
     @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 @@
         }
     }
 
+    /**
+     * 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 @@
         }
         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 @@
             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 @@
                     synchronized (mDeviceLockedForUser) {
                         mDeviceLockedForUser.put(userId, locked);
                     }
+
+                    KeyStore.getInstance().onUserLockedStateChanged(userId, locked);
+
                     if (locked) {
                         try {
                             ActivityManager.getService().notifyLockedProfile(userId);