Update user switching processes

Add a new callback for UserController. If the user is secure, initiate
a call to SystemUI to lockNow(), and wait for the callback. If the
callback does not arrive, crash the system after some time.

In SystemUI, add support for onBeforeUserSwitching. When this request
comes in, prepare the keyguard and show it before calling back to notify
UserController that it is ready to proceed.

Also, ensure that any switch is user that occurs after WM has been
notified that SystemUI is going away is stopped.

Bug: 322157041
Test: atest KeyguardViewMediatorTest
Flag: EXEMPT bugfix

Change-Id: I6e48ae4787680f0be80525b0228883523ef08a7d

5 files changed, 377 insertions, 87 deletions

diff --git a/core/java/android/app/KeyguardManager.java b/core/java/android/app/KeyguardManager.java
index b5ac4e78c7ad..d91838c4cc2b 100644
--- a/core/java/android/app/KeyguardManager.java
+++ b/core/java/android/app/KeyguardManager.java
@@ -205,6 +205,15 @@ public class KeyguardManager {
     public static final String EXTRA_DISALLOW_BIOMETRICS_IF_POLICY_EXISTS = "check_dpm";
 
     /**
+     * When switching to a secure user, system server will expect a callback when the UI has
+     * completed the switch.
+     *
+     * @hide
+     */
+    public static final String LOCK_ON_USER_SWITCH_CALLBACK = "onSwitchCallback";
+
+
+    /**
      *
      * Password lock type, see {@link #setLock}
      *
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
index 5baef915ea01..7fed7d253efe 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/KeyguardViewMediator.java
@@ -16,6 +16,7 @@
 
 package com.android.systemui.keyguard;
 
+import static android.app.KeyguardManager.LOCK_ON_USER_SWITCH_CALLBACK;
 import static android.app.StatusBarManager.SESSION_KEYGUARD;
 import static android.provider.Settings.Secure.LOCK_SCREEN_LOCK_AFTER_TIMEOUT;
 import static android.provider.Settings.System.LOCKSCREEN_SOUNDS_ENABLED;
@@ -75,6 +76,7 @@ import android.os.Bundle;
 import android.os.DeadObjectException;
 import android.os.Handler;
 import android.os.IBinder;
+import android.os.IRemoteCallback;
 import android.os.Looper;
 import android.os.Message;
 import android.os.PowerManager;
@@ -190,6 +192,8 @@ import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Iterator;
+import java.util.List;
 import java.util.Objects;
 import java.util.concurrent.Executor;
 import java.util.function.Consumer;
@@ -271,6 +275,9 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     private static final int SYSTEM_READY = 18;
     private static final int CANCEL_KEYGUARD_EXIT_ANIM = 19;
     private static final int BOOT_INTERACTOR = 20;
+    private static final int BEFORE_USER_SWITCHING = 21;
+    private static final int USER_SWITCHING = 22;
+    private static final int USER_SWITCH_COMPLETE = 23;
 
     /** Enum for reasons behind updating wakeAndUnlock state. */
     @Retention(RetentionPolicy.SOURCE)
@@ -288,6 +295,8 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         int WAKE_AND_UNLOCK = 3;
     }
 
+    private final List<LockNowCallback> mLockNowCallbacks = new ArrayList<>();
+
     /**
      * The default amount of time we stay awake (used for all key input)
      */
@@ -345,6 +354,11 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     private final ScreenOffAnimationController mScreenOffAnimationController;
     private final Lazy<NotificationShadeDepthController> mNotificationShadeDepthController;
     private final Lazy<ShadeController> mShadeController;
+    /*
+     * Records the user id on request to go away, for validation when WM calls back to start the
+     * exit animation.
+     */
+    private int mGoingAwayRequestedForUserId = -1;
 
     private boolean mSystemReady;
     private boolean mBootCompleted;
@@ -352,7 +366,6 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     private boolean mShuttingDown;
     private boolean mDozing;
     private boolean mAnimatingScreenOff;
-    private boolean mIgnoreDismiss;
     private final Context mContext;
     private final FalsingCollector mFalsingCollector;
 
@@ -615,41 +628,92 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
                 }
             };
 
-    KeyguardUpdateMonitorCallback mUpdateCallback = new KeyguardUpdateMonitorCallback() {
+    @VisibleForTesting
+    protected UserTracker.Callback mUserChangedCallback = new UserTracker.Callback() {
 
         @Override
-        public void onKeyguardVisibilityChanged(boolean visible) {
-            synchronized (KeyguardViewMediator.this) {
-                if (!visible && mPendingPinLock) {
-                    Log.i(TAG, "PIN lock requested, starting keyguard");
+        public void onBeforeUserSwitching(int newUser, @NonNull Runnable resultCallback) {
+            mHandler.sendMessage(mHandler.obtainMessage(BEFORE_USER_SWITCHING,
+                    newUser, 0, resultCallback));
+        }
 
-                    // Bring the keyguard back in order to show the PIN lock
-                    mPendingPinLock = false;
-                    doKeyguardLocked(null);
-                }
-            }
+        @Override
+        public void onUserChanging(int newUser, @NonNull Context userContext,
+                @NonNull Runnable resultCallback) {
+            mHandler.sendMessage(mHandler.obtainMessage(USER_SWITCHING,
+                    newUser, 0, resultCallback));
         }
 
         @Override
-        public void onUserSwitching(int userId) {
-            Log.d(TAG, String.format("onUserSwitching %d", userId));
-            synchronized (KeyguardViewMediator.this) {
-                mIgnoreDismiss = true;
-                notifyTrustedChangedLocked(mUpdateMonitor.getUserHasTrust(userId));
-                resetKeyguardDonePendingLocked();
+        public void onUserChanged(int newUser, Context userContext) {
+            mHandler.sendMessage(mHandler.obtainMessage(USER_SWITCH_COMPLETE,
+                    newUser, 0));
+        }
+    };
+
+    /**
+     * Handle {@link #BEFORE_USER_SWITCHING}
+     */
+    @VisibleForTesting
+    void handleBeforeUserSwitching(int userId, Runnable resultCallback) {
+        Log.d(TAG, String.format("onBeforeUserSwitching %d", userId));
+        synchronized (KeyguardViewMediator.this) {
+            mHandler.removeMessages(DISMISS);
+            notifyTrustedChangedLocked(mUpdateMonitor.getUserHasTrust(userId));
+            resetKeyguardDonePendingLocked();
+            adjustStatusBarLocked();
+            mKeyguardStateController.notifyKeyguardGoingAway(false);
+            if (mLockPatternUtils.isSecure(userId) && !mShowing) {
+                doKeyguardLocked(null);
+            } else {
                 resetStateLocked();
-                adjustStatusBarLocked();
             }
+            resultCallback.run();
         }
+    }
 
-        @Override
-        public void onUserSwitchComplete(int userId) {
-            mIgnoreDismiss = false;
-            Log.d(TAG, String.format("onUserSwitchComplete %d", userId));
+    /**
+     * Handle {@link #USER_SWITCHING}
+     */
+    @VisibleForTesting
+    void handleUserSwitching(int userId, Runnable resultCallback) {
+        Log.d(TAG, String.format("onUserSwitching %d", userId));
+        synchronized (KeyguardViewMediator.this) {
+            if (!mLockPatternUtils.isSecure(userId)) {
+                dismiss(null, null);
+            }
+            resultCallback.run();
+        }
+    }
+
+    /**
+     * Handle {@link #USER_SWITCH_COMPLETE}
+     */
+    @VisibleForTesting
+    void handleUserSwitchComplete(int userId) {
+        Log.d(TAG, String.format("onUserSwitchComplete %d", userId));
+        // Calling dismiss on a secure user will show the bouncer
+        if (mLockPatternUtils.isSecure(userId)) {
             // We are calling dismiss with a delay as there are race conditions in some scenarios
             // caused by async layout listeners
             mHandler.postDelayed(() -> dismiss(null /* callback */, null /* message */), 500);
         }
+    }
+
+    KeyguardUpdateMonitorCallback mUpdateCallback = new KeyguardUpdateMonitorCallback() {
+
+        @Override
+        public void onKeyguardVisibilityChanged(boolean visible) {
+            synchronized (KeyguardViewMediator.this) {
+                if (!visible && mPendingPinLock) {
+                    Log.i(TAG, "PIN lock requested, starting keyguard");
+
+                    // Bring the keyguard back in order to show the PIN lock
+                    mPendingPinLock = false;
+                    doKeyguardLocked(null);
+                }
+            }
+        }
 
         @Override
         public void onDeviceProvisioned() {
@@ -1658,7 +1722,13 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
                 com.android.internal.R.anim.lock_screen_behind_enter);
 
         mWorkLockController = new WorkLockActivityController(mContext, mUserTracker);
-
+        mUserTracker.addCallback(mUserChangedCallback, mContext.getMainExecutor());
+        // start() can be invoked in the middle of user switching, so check for this state and issue
+        // the call manually as that important event was missed.
+        if (mUserTracker.isUserSwitching()) {
+            handleBeforeUserSwitching(mUserTracker.getUserId(), () -> {});
+            handleUserSwitching(mUserTracker.getUserId(), () -> {});
+        }
         mJavaAdapter.alwaysCollectFlow(
                 mWallpaperRepository.getWallpaperSupportsAmbientMode(),
                 this::setWallpaperSupportsAmbientMode);
@@ -1707,7 +1777,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
             // System ready can be invoked in the middle of user switching, so check for this state
             // and issue the call manually as that important event was missed.
             if (mUserTracker.isUserSwitching()) {
-                mUpdateCallback.onUserSwitching(mUserTracker.getUserId());
+                mUserChangedCallback.onUserChanging(mUserTracker.getUserId(), mContext, () -> {});
             }
         }
         // Most services aren't available until the system reaches the ready state, so we
@@ -2341,12 +2411,23 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
      * Enable the keyguard if the settings are appropriate.
      */
     private void doKeyguardLocked(Bundle options) {
+        int currentUserId = mSelectedUserInteractor.getSelectedUserId();
+        if (options != null && options.getBinder(LOCK_ON_USER_SWITCH_CALLBACK) != null) {
+            LockNowCallback callback = new LockNowCallback(currentUserId,
+                    IRemoteCallback.Stub.asInterface(
+                            options.getBinder(LOCK_ON_USER_SWITCH_CALLBACK)));
+            synchronized (mLockNowCallbacks) {
+                mLockNowCallbacks.add(callback);
+            }
+            Log.d(TAG, "LockNowCallback required for user: " + callback.mUserId);
+        }
+
         // if another app is disabling us, don't show
         if (!mExternallyEnabled
                 && !mLockPatternUtils.isUserInLockdown(
                         mSelectedUserInteractor.getSelectedUserId())) {
             if (DEBUG) Log.d(TAG, "doKeyguard: not showing because externally disabled");
-
+            notifyLockNowCallback();
             mNeedToReshowWhenReenabled = true;
             return;
         }
@@ -2364,6 +2445,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
                     // We're removing "reset" in the refactor - "resetting" the views will happen
                     // as a reaction to the root cause of the "reset" signal.
                     if (KeyguardWmStateRefactor.isEnabled()) {
+                        notifyLockNowCallback();
                         return;
                     }
 
@@ -2376,6 +2458,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
                                     + "previously hiding. It should be safe to short-circuit "
                                     + "here.");
                     resetStateLocked(/* hideBouncer= */ false);
+                    notifyLockNowCallback();
                     return;
                 }
             } else {
@@ -2402,6 +2485,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
                 Log.d(TAG, "doKeyguard: not showing because device isn't provisioned and the sim is"
                         + " not locked or missing");
             }
+            notifyLockNowCallback();
             return;
         }
 
@@ -2409,6 +2493,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         if (mLockPatternUtils.isLockScreenDisabled(mSelectedUserInteractor.getSelectedUserId())
                 && !lockedOrMissing && !forceShow) {
             if (DEBUG) Log.d(TAG, "doKeyguard: not showing because lockscreen is off");
+            notifyLockNowCallback();
             return;
         }
 
@@ -2456,11 +2541,6 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     }
 
     public void dismiss(IKeyguardDismissCallback callback, CharSequence message) {
-        if (mIgnoreDismiss) {
-            android.util.Log.i(TAG, "Ignoring request to dismiss (user switch in progress?)");
-            return;
-        }
-
         if (mKeyguardStateController.isKeyguardGoingAway()) {
             Log.i(TAG, "Ignoring dismiss because we're already going away.");
             return;
@@ -2478,7 +2558,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
     }
 
     private void resetStateLocked(boolean hideBouncer) {
-        if (DEBUG) Log.e(TAG, "resetStateLocked");
+        if (DEBUG) Log.d(TAG, "resetStateLocked");
         Message msg = mHandler.obtainMessage(RESET, hideBouncer ? 1 : 0, 0);
         mHandler.sendMessage(msg);
     }
@@ -2726,6 +2806,18 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
                     message = "BOOT_INTERACTOR";
                     handleBootInteractor();
                     break;
+                case BEFORE_USER_SWITCHING:
+                    message = "BEFORE_USER_SWITCHING";
+                    handleBeforeUserSwitching(msg.arg1, (Runnable) msg.obj);
+                    break;
+                case USER_SWITCHING:
+                    message = "USER_SWITCHING";
+                    handleUserSwitching(msg.arg1, (Runnable) msg.obj);
+                    break;
+                case USER_SWITCH_COMPLETE:
+                    message = "USER_SWITCH_COMPLETE";
+                    handleUserSwitchComplete(msg.arg1);
+                    break;
             }
             Log.d(TAG, "KeyguardViewMediator queue processing message: " + message);
         }
@@ -2867,6 +2959,9 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         mUiBgExecutor.execute(() -> {
             Log.d(TAG, "updateActivityLockScreenState(" + showing + ", " + aodShowing + ", "
                     + reason + ")");
+            if (showing) {
+                notifyLockNowCallback();
+            }
 
             if (KeyguardWmStateRefactor.isEnabled()) {
                 // Handled in WmLockscreenVisibilityManager if flag is enabled.
@@ -2911,6 +3006,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         synchronized (KeyguardViewMediator.this) {
             if (!mSystemReady) {
                 if (DEBUG) Log.d(TAG, "ignoring handleShow because system is not ready.");
+                notifyLockNowCallback();
                 return;
             }
             if (DEBUG) Log.d(TAG, "handleShow");
@@ -2969,12 +3065,11 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         }
     }
 
-    private final Runnable mKeyguardGoingAwayRunnable = new Runnable() {
+    final Runnable mKeyguardGoingAwayRunnable = new Runnable() {
         @SuppressLint("MissingPermission")
         @Override
         public void run() {
             Trace.beginSection("KeyguardViewMediator.mKeyGuardGoingAwayRunnable");
-            Log.d(TAG, "keyguardGoingAwayRunnable");
             mKeyguardViewControllerLazy.get().keyguardGoingAway();
 
             int flags = 0;
@@ -3011,6 +3106,10 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
 
             // Handled in WmLockscreenVisibilityManager if flag is enabled.
             if (!KeyguardWmStateRefactor.isEnabled()) {
+                mGoingAwayRequestedForUserId = mSelectedUserInteractor.getSelectedUserId();
+                Log.d(TAG, "keyguardGoingAway requested for userId: "
+                        + mGoingAwayRequestedForUserId);
+
                 // Don't actually hide the Keyguard at the moment, wait for window manager
                 // until it tells us it's safe to do so with startKeyguardExitAnimation.
                 // Posting to mUiOffloadThread to ensure that calls to ActivityTaskManager
@@ -3149,6 +3248,30 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
             RemoteAnimationTarget[] nonApps, IRemoteAnimationFinishedCallback finishedCallback) {
         Log.d(TAG, "handleStartKeyguardExitAnimation startTime=" + startTime
                 + " fadeoutDuration=" + fadeoutDuration);
+        int currentUserId = mSelectedUserInteractor.getSelectedUserId();
+        if (mGoingAwayRequestedForUserId != currentUserId) {
+            Log.e(TAG, "Not executing handleStartKeyguardExitAnimationInner() due to userId "
+                    + "mismatch. Requested: " + mGoingAwayRequestedForUserId + ", current: "
+                    + currentUserId);
+            if (finishedCallback != null) {
+                // There will not execute animation, send a finish callback to ensure the remote
+                // animation won't hang there.
+                try {
+                    finishedCallback.onAnimationFinished();
+                } catch (RemoteException e) {
+                    Slog.w(TAG, "Failed to call onAnimationFinished", e);
+                }
+            }
+            mHiding = false;
+            if (mLockPatternUtils.isSecure(currentUserId)) {
+                doKeyguardLocked(null);
+            } else {
+                resetStateLocked();
+                dismiss(null, null);
+            }
+            return;
+        }
+
         synchronized (KeyguardViewMediator.this) {
             mIsKeyguardExitAnimationCanceled = false;
             // Tell ActivityManager that we canceled the keyguard animation if
@@ -3393,6 +3516,12 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
      * app transition before finishing the current RemoteAnimation, or the keyguard being re-shown).
      */
     private void handleCancelKeyguardExitAnimation() {
+        if (mGoingAwayRequestedForUserId != mSelectedUserInteractor.getSelectedUserId()) {
+            Log.e(TAG, "Setting pendingLock = true due to userId mismatch. Requested: "
+                    + mGoingAwayRequestedForUserId + ", current: "
+                    + mSelectedUserInteractor.getSelectedUserId());
+            setPendingLock(true);
+        }
         if (mPendingLock) {
             Log.d(TAG, "#handleCancelKeyguardExitAnimation: keyguard exit animation cancelled. "
                     + "There's a pending lock, so we were cancelled because the device was locked "
@@ -3493,6 +3622,7 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         mSurfaceBehindRemoteAnimationRequested = true;
 
         if (ENABLE_NEW_KEYGUARD_SHELL_TRANSITIONS && !KeyguardWmStateRefactor.isEnabled()) {
+            mGoingAwayRequestedForUserId = mSelectedUserInteractor.getSelectedUserId();
             startKeyguardTransition(false /* keyguardShowing */, false /* aodShowing */);
             return;
         }
@@ -3513,6 +3643,9 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
 
             if (!KeyguardWmStateRefactor.isEnabled()) {
                 // Handled in WmLockscreenVisibilityManager.
+                mGoingAwayRequestedForUserId = mSelectedUserInteractor.getSelectedUserId();
+                Log.d(TAG, "keyguardGoingAway requested for userId: "
+                        + mGoingAwayRequestedForUserId);
                 mActivityTaskManagerService.keyguardGoingAway(flags);
             }
         } catch (RemoteException e) {
@@ -3966,6 +4099,29 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
         mUiBgExecutor.execute(mTrustManager::reportKeyguardShowingChanged);
     }
 
+    private void notifyLockNowCallback() {
+        List<LockNowCallback> callbacks;
+        synchronized (mLockNowCallbacks) {
+            callbacks = new ArrayList<LockNowCallback>(mLockNowCallbacks);
+            mLockNowCallbacks.clear();
+        }
+        Iterator<LockNowCallback> iter = callbacks.listIterator();
+        while (iter.hasNext()) {
+            LockNowCallback callback = iter.next();
+            iter.remove();
+            if (callback.mUserId != mSelectedUserInteractor.getSelectedUserId()) {
+                Log.i(TAG, "Not notifying lockNowCallback due to user mismatch");
+                continue;
+            }
+            Log.i(TAG, "Notifying lockNowCallback");
+            try {
+                callback.mRemoteCallback.sendResult(null);
+            } catch (RemoteException e) {
+                Log.e(TAG, "Could not issue LockNowCallback sendResult", e);
+            }
+        }
+    }
+
     private void notifyTrustedChangedLocked(boolean trusted) {
         int size = mKeyguardStateCallbacks.size();
         for (int i = size - 1; i >= 0; i--) {
@@ -4130,4 +4286,14 @@ public class KeyguardViewMediator implements CoreStartable, Dumpable,
             }
         };
     }
+
+    private class LockNowCallback {
+        final int mUserId;
+        final IRemoteCallback mRemoteCallback;
+
+        LockNowCallback(int userId, IRemoteCallback remoteCallback) {
+            mUserId = userId;
+            mRemoteCallback = remoteCallback;
+        }
+    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
index 38acd23d282c..0c9213c3a722 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/KeyguardViewMediatorTest.java
@@ -206,6 +206,7 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
     private @Mock ShadeInteractor mShadeInteractor;
     private @Mock ShadeWindowLogger mShadeWindowLogger;
     private @Mock SelectedUserInteractor mSelectedUserInteractor;
+    private @Mock UserTracker.Callback mUserTrackerCallback;
     private @Mock KeyguardInteractor mKeyguardInteractor;
     private @Mock KeyguardTransitionBootInteractor mKeyguardTransitionBootInteractor;
     private @Captor ArgumentCaptor<KeyguardStateController.Callback>
@@ -280,7 +281,7 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
                 () -> mShadeInteractor,
                 mShadeWindowLogger,
                 () -> mSelectedUserInteractor,
-                mUserTracker,
+                mock(UserTracker.class),
                 mKosmos.getNotificationShadeWindowModel(),
                 mSecureSettings,
                 mKosmos::getCommunalInteractor,
@@ -318,7 +319,7 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
 
         } catch (Exception e) {
             // Just so we don't have to add the exception signature to every test.
-            fail();
+            fail(e.getMessage());
         }
     }
 
@@ -330,18 +331,156 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
 
         /* First test the default behavior: handleUserSwitching() is not invoked */
         when(mUserTracker.isUserSwitching()).thenReturn(false);
-        mViewMediator.mUpdateCallback = mock(KeyguardUpdateMonitorCallback.class);
         mViewMediator.onSystemReady();
         TestableLooper.get(this).processAllMessages();
 
-        verify(mViewMediator.mUpdateCallback, never()).onUserSwitching(userId);
+        verify(mUserTrackerCallback, never()).onUserChanging(eq(userId), eq(mContext),
+                any(Runnable.class));
 
         /* Next test user switching is already in progress when started */
         when(mUserTracker.isUserSwitching()).thenReturn(true);
         mViewMediator.onSystemReady();
         TestableLooper.get(this).processAllMessages();
 
-        verify(mViewMediator.mUpdateCallback).onUserSwitching(userId);
+        verify(mUserTrackerCallback).onUserChanging(eq(userId), eq(mContext),
+                any(Runnable.class));
+    }
+
+    @Test
+    @TestableLooper.RunWithLooper(setAsMainLooper = true)
+    public void testGoingAwayFollowedByBeforeUserSwitchDoesNotHideKeyguard() {
+        setCurrentUser(/* userId= */1099, /* isSecure= */false);
+
+        // Setup keyguard
+        mViewMediator.onSystemReady();
+        processAllMessagesAndBgExecutorMessages();
+        mViewMediator.setShowingLocked(true, "");
+
+        // Request keyguard going away
+        when(mKeyguardStateController.isKeyguardGoingAway()).thenReturn(true);
+        mViewMediator.mKeyguardGoingAwayRunnable.run();
+
+        // After the request, begin a switch to a new secure user
+        int nextUserId = 500;
+        setCurrentUser(nextUserId, /* isSecure= */true);
+        Runnable result = mock(Runnable.class);
+        mViewMediator.handleBeforeUserSwitching(nextUserId, result);
+        processAllMessagesAndBgExecutorMessages();
+        verify(result).run();
+
+        // After that request has begun, have WM tell us to exit keyguard
+        RemoteAnimationTarget[] apps = new RemoteAnimationTarget[]{
+                mock(RemoteAnimationTarget.class)
+        };
+        RemoteAnimationTarget[] wallpapers = new RemoteAnimationTarget[]{
+                mock(RemoteAnimationTarget.class)
+        };
+        IRemoteAnimationFinishedCallback callback = mock(IRemoteAnimationFinishedCallback.class);
+        mViewMediator.startKeyguardExitAnimation(TRANSIT_OLD_KEYGUARD_GOING_AWAY, apps, wallpapers,
+                null, callback);
+        processAllMessagesAndBgExecutorMessages();
+
+        // The call to exit should be rejected, and keyguard should still be visible
+        verify(mKeyguardUnlockAnimationController, never()).notifyStartSurfaceBehindRemoteAnimation(
+                any(), any(), any(), anyLong(), anyBoolean());
+        try {
+            assertATMSLockScreenShowing(true);
+        } catch (Exception e) {
+            fail(e.getMessage());
+        }
+        assertTrue(mViewMediator.isShowingAndNotOccluded());
+    }
+
+    @Test
+    @TestableLooper.RunWithLooper(setAsMainLooper = true)
+    public void testUserSwitchToSecureUserShowsBouncer() {
+        setCurrentUser(/* userId= */1099, /* isSecure= */true);
+
+        // Setup keyguard
+        mViewMediator.onSystemReady();
+        processAllMessagesAndBgExecutorMessages();
+        mViewMediator.setShowingLocked(true, "");
+
+        // After the request, begin a switch to a new secure user
+        int nextUserId = 500;
+        setCurrentUser(nextUserId, /* isSecure= */true);
+
+        Runnable beforeResult = mock(Runnable.class);
+        mViewMediator.handleBeforeUserSwitching(nextUserId, beforeResult);
+        processAllMessagesAndBgExecutorMessages();
+        verify(beforeResult).run();
+
+        // Dismiss should not be called while user switch is in progress
+        Runnable onSwitchResult = mock(Runnable.class);
+        mViewMediator.handleUserSwitching(nextUserId, onSwitchResult);
+        processAllMessagesAndBgExecutorMessages();
+        verify(onSwitchResult).run();
+        verify(mStatusBarKeyguardViewManager, never()).dismissAndCollapse();
+
+        // The attempt to dismiss only comes on user switch complete, which will trigger a call to
+        // show the bouncer in StatusBarKeyguardViewManager
+        mViewMediator.handleUserSwitchComplete(nextUserId);
+        TestableLooper.get(this).moveTimeForward(600);
+        processAllMessagesAndBgExecutorMessages();
+
+        verify(mStatusBarKeyguardViewManager).dismissAndCollapse();
+    }
+
+    @Test
+    @TestableLooper.RunWithLooper(setAsMainLooper = true)
+    public void testUserSwitchToInsecureUserDismissesKeyguard() {
+        int userId = 1099;
+        when(mUserTracker.getUserId()).thenReturn(userId);
+
+        // Setup keyguard
+        mViewMediator.onSystemReady();
+        processAllMessagesAndBgExecutorMessages();
+        mViewMediator.setShowingLocked(true, "");
+
+        // After the request, begin a switch to an insecure user
+        int nextUserId = 500;
+        when(mLockPatternUtils.isSecure(nextUserId)).thenReturn(false);
+
+        Runnable beforeResult = mock(Runnable.class);
+        mViewMediator.handleBeforeUserSwitching(nextUserId, beforeResult);
+        processAllMessagesAndBgExecutorMessages();
+        verify(beforeResult).run();
+
+        // The call to dismiss comes during the user switch
+        Runnable onSwitchResult = mock(Runnable.class);
+        mViewMediator.handleUserSwitching(nextUserId, onSwitchResult);
+        processAllMessagesAndBgExecutorMessages();
+        verify(onSwitchResult).run();
+
+        verify(mStatusBarKeyguardViewManager).dismissAndCollapse();
+    }
+
+    @Test
+    @TestableLooper.RunWithLooper(setAsMainLooper = true)
+    public void testUserSwitchToSecureUserWhileKeyguardNotVisibleShowsKeyguard() {
+        setCurrentUser(/* userId= */1099, /* isSecure= */true);
+
+        // Setup keyguard as not visible
+        mViewMediator.onSystemReady();
+        processAllMessagesAndBgExecutorMessages();
+        mViewMediator.setShowingLocked(false, "");
+        processAllMessagesAndBgExecutorMessages();
+
+        // Begin a switch to a new secure user
+        int nextUserId = 500;
+        setCurrentUser(nextUserId, /* isSecure= */true);
+
+        Runnable beforeResult = mock(Runnable.class);
+        mViewMediator.handleBeforeUserSwitching(nextUserId, beforeResult);
+        processAllMessagesAndBgExecutorMessages();
+        verify(beforeResult).run();
+
+        try {
+            assertATMSLockScreenShowing(true);
+        } catch (Exception e) {
+            fail();
+        }
+        assertTrue(mViewMediator.isShowingAndNotOccluded());
     }
 
     @Test
@@ -1105,7 +1244,7 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
         processAllMessagesAndBgExecutorMessages();
 
         verify(mStatusBarKeyguardViewManager, never()).reset(anyBoolean());
-        assertATMSAndKeyguardViewMediatorStatesMatch();
+
     }
 
     @Test
@@ -1149,6 +1288,7 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
         IRemoteAnimationFinishedCallback callback = mock(IRemoteAnimationFinishedCallback.class);
 
         when(mKeyguardStateController.isKeyguardGoingAway()).thenReturn(true);
+        mViewMediator.mKeyguardGoingAwayRunnable.run();
         mViewMediator.startKeyguardExitAnimation(TRANSIT_OLD_KEYGUARD_GOING_AWAY, apps, wallpapers,
                 null, callback);
         processAllMessagesAndBgExecutorMessages();
@@ -1203,13 +1343,6 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
 
         // The captor will have the most recent setLockScreenShown call's value.
         assertEquals(showing, showingCaptor.getValue());
-
-        // We're now just after the last setLockScreenShown call. If we expect the lockscreen to be
-        // showing, ensure that we didn't subsequently ask for it to go away.
-        if (showing) {
-            orderedSetLockScreenShownCalls.verify(mActivityTaskManagerService, never())
-                    .keyguardGoingAway(anyInt());
-        }
     }
 
     /**
@@ -1370,6 +1503,7 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
                 mKeyguardInteractor,
                 mKeyguardTransitionBootInteractor,
                 mock(WindowManagerOcclusionManager.class));
+        mViewMediator.mUserChangedCallback = mUserTrackerCallback;
         mViewMediator.start();
 
         mViewMediator.registerCentralSurfaces(mCentralSurfaces, null, null, null, null);
@@ -1383,4 +1517,10 @@ public class KeyguardViewMediatorTest extends SysuiTestCase {
     private void captureKeyguardUpdateMonitorCallback() {
         verify(mUpdateMonitor).registerCallback(mKeyguardUpdateMonitorCallbackCaptor.capture());
     }
+
+    private void setCurrentUser(int userId, boolean isSecure) {
+        when(mUserTracker.getUserId()).thenReturn(userId);
+        when(mSelectedUserInteractor.getSelectedUserId()).thenReturn(userId);
+        when(mLockPatternUtils.isSecure(userId)).thenReturn(isSecure);
+    }
 }
diff --git a/services/core/java/com/android/server/am/UserController.java b/services/core/java/com/android/server/am/UserController.java
index 27e9e44f1090..e0fbaf43ea43 100644
--- a/services/core/java/com/android/server/am/UserController.java
+++ b/services/core/java/com/android/server/am/UserController.java
@@ -31,6 +31,7 @@ import static android.app.ActivityManagerInternal.ALLOW_FULL_ONLY;
 import static android.app.ActivityManagerInternal.ALLOW_NON_FULL;
 import static android.app.ActivityManagerInternal.ALLOW_NON_FULL_IN_PROFILE;
 import static android.app.ActivityManagerInternal.ALLOW_PROFILES_OR_NON_FULL;
+import static android.app.KeyguardManager.LOCK_ON_USER_SWITCH_CALLBACK;
 import static android.os.PowerWhitelistManager.REASON_BOOT_COMPLETED;
 import static android.os.PowerWhitelistManager.REASON_LOCKED_BOOT_COMPLETED;
 import static android.os.PowerWhitelistManager.TEMPORARY_ALLOWLIST_TYPE_FOREGROUND_SERVICE_ALLOWED;
@@ -3904,10 +3905,6 @@ class UserController implements Handler.Callback {
             return mService.mWindowManager;
         }
 
-        ActivityTaskManagerInternal getActivityTaskManagerInternal() {
-            return mService.mAtmInternal;
-        }
-
         void activityManagerOnUserStopped(@UserIdInt int userId) {
             LocalServices.getService(ActivityTaskManagerInternal.class).onUserStopped(userId);
         }
@@ -4122,40 +4119,25 @@ class UserController implements Handler.Callback {
         }
 
         void lockDeviceNowAndWaitForKeyguardShown() {
-            if (getWindowManager().isKeyguardLocked()) {
-                Slogf.w(TAG, "Not locking the device since the keyguard is already locked");
-                return;
-            }
-
             final TimingsTraceAndSlog t = new TimingsTraceAndSlog();
             t.traceBegin("lockDeviceNowAndWaitForKeyguardShown");
 
             final CountDownLatch latch = new CountDownLatch(1);
-            ActivityTaskManagerInternal.ScreenObserver screenObserver =
-                    new ActivityTaskManagerInternal.ScreenObserver() {
-                        @Override
-                        public void onAwakeStateChanged(boolean isAwake) {
-
-                        }
-
-                        @Override
-                        public void onKeyguardStateChanged(boolean isShowing) {
-                            if (isShowing) {
-                                latch.countDown();
-                            }
-                        }
-                    };
-
-            getActivityTaskManagerInternal().registerScreenObserver(screenObserver);
-            getWindowManager().lockDeviceNow();
+            Bundle bundle = new Bundle();
+            bundle.putBinder(LOCK_ON_USER_SWITCH_CALLBACK, new IRemoteCallback.Stub() {
+                public void sendResult(Bundle data) {
+                    latch.countDown();
+                }
+            });
+            getWindowManager().lockNow(bundle);
             try {
                 if (!latch.await(20, TimeUnit.SECONDS)) {
-                    throw new RuntimeException("Keyguard is not shown in 20 seconds");
+                    throw new RuntimeException("User controller expected a callback while waiting "
+                            + "to show the keyguard. Timed out after 20 seconds.");
                 }
             } catch (InterruptedException e) {
                 throw new RuntimeException(e);
             } finally {
-                getActivityTaskManagerInternal().unregisterScreenObserver(screenObserver);
                 t.traceEnd();
             }
         }
diff --git a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
index 06958b81d846..1627f683cd3e 100644
--- a/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
+++ b/services/tests/servicestests/src/com/android/server/am/UserControllerTest.java
@@ -25,6 +25,7 @@ import static android.app.ActivityManagerInternal.ALLOW_FULL_ONLY;
 import static android.app.ActivityManagerInternal.ALLOW_NON_FULL;
 import static android.app.ActivityManagerInternal.ALLOW_NON_FULL_IN_PROFILE;
 import static android.app.ActivityManagerInternal.ALLOW_PROFILES_OR_NON_FULL;
+import static android.app.KeyguardManager.LOCK_ON_USER_SWITCH_CALLBACK;
 import static android.content.pm.PackageManager.PERMISSION_GRANTED;
 import static android.testing.DexmakerShareClassLoaderRule.runWithDexmakerShareClassLoader;
 
@@ -115,7 +116,6 @@ import com.android.server.pm.UserManagerInternal;
 import com.android.server.pm.UserManagerService;
 import com.android.server.pm.UserTypeDetails;
 import com.android.server.pm.UserTypeFactory;
-import com.android.server.wm.ActivityTaskManagerInternal;
 import com.android.server.wm.WindowManagerService;
 
 import com.google.common.collect.Range;
@@ -1563,11 +1563,11 @@ public class UserControllerTest {
         // and the thread is still alive
         assertTrue(threadStartUser.isAlive());
 
-        // mock send the keyguard shown event
-        ArgumentCaptor<ActivityTaskManagerInternal.ScreenObserver> captor = ArgumentCaptor.forClass(
-                ActivityTaskManagerInternal.ScreenObserver.class);
-        verify(mInjector.mActivityTaskManagerInternal).registerScreenObserver(captor.capture());
-        captor.getValue().onKeyguardStateChanged(true);
+        // mock the binder response for the user switch completion
+        ArgumentCaptor<Bundle> captor = ArgumentCaptor.forClass(Bundle.class);
+        verify(mInjector.mWindowManagerMock).lockNow(captor.capture());
+        IRemoteCallback.Stub.asInterface(captor.getValue().getBinder(
+                LOCK_ON_USER_SWITCH_CALLBACK)).sendResult(null);
 
         // verify the switch now moves on...
         Thread.sleep(1000);
@@ -1757,7 +1757,6 @@ public class UserControllerTest {
         private final IStorageManager mStorageManagerMock;
         private final UserManagerInternal mUserManagerInternalMock;
         private final WindowManagerService mWindowManagerMock;
-        private final ActivityTaskManagerInternal mActivityTaskManagerInternal;
         private final PowerManagerInternal mPowerManagerInternal;
         private final AlarmManagerInternal mAlarmManagerInternal;
         private final KeyguardManager mKeyguardManagerMock;
@@ -1779,7 +1778,6 @@ public class UserControllerTest {
             mUserManagerMock = mock(UserManagerService.class);
             mUserManagerInternalMock = mock(UserManagerInternal.class);
             mWindowManagerMock = mock(WindowManagerService.class);
-            mActivityTaskManagerInternal = mock(ActivityTaskManagerInternal.class);
             mStorageManagerMock = mock(IStorageManager.class);
             mPowerManagerInternal = mock(PowerManagerInternal.class);
             mAlarmManagerInternal = mock(AlarmManagerInternal.class);
@@ -1843,11 +1841,6 @@ public class UserControllerTest {
         }
 
         @Override
-        ActivityTaskManagerInternal getActivityTaskManagerInternal() {
-            return mActivityTaskManagerInternal;
-        }
-
-        @Override
         PowerManagerInternal getPowerManagerInternal() {
             return mPowerManagerInternal;
         }