AM: make LockTask truly multi-task-aware.

LockTask was first designed to be single-task only. It was later
expanded to multi-task use cases in a rather flawed manner. Many aspects
of the implementation overlook the implications of more than one locked
tasks, and sometimes completely ignore the LockTask-unaware apps.

For example, LockTask-unaware apps are not registered in
LockTaskController's list of locked tasks, although they can be launched
in LockTask mode as long as they are whitelisted. As a consequence,
these unregistered tasks are not finished when their whitelist
authorization is removed.

In this patch, we make sure LockTask-unaware apps are also registered,
and introduce the notion of a root task, thus providing a genuine
multi-task experience with LockTask.

Bug: 66130096
Bug: 66132035
Test: bit FrameworksServicesTests:com.android.server.am.LockTaskControllerTest
Test: cts-tradefed run cts-dev -m DevicePolicyManager --test com.android.cts.devicepolicy.DeviceOwnerTest#testLockTask_deviceOwnerUser
Test: manual, by using TestDPC > Start kiosk mode
Change-Id: I68a26d2198fa3991f85ae1e7a6acae5981c34db5
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index f2e0493..0246263 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -10806,7 +10806,7 @@
         }
     }
 
-    private void startLockTaskModeLocked(@Nullable TaskRecord task, boolean isAppPinning) {
+    private void startLockTaskModeLocked(@Nullable TaskRecord task, boolean isSystemCaller) {
         if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK, "startLockTaskModeLocked: " + task);
         if (task == null || task.mLockTaskAuth == LOCK_TASK_AUTH_DONT_LOCK) {
             return;
@@ -10820,13 +10820,16 @@
         // When a task is locked, dismiss the pinned stack if it exists
         mStackSupervisor.removeStacksInWindowingModes(WINDOWING_MODE_PINNED);
 
-        // isAppPinning is used to distinguish between locked and pinned mode, as pinned mode
-        // is initiated by system after the pinning request was shown and locked mode is initiated
-        // by an authorized app directly
+        // {@code isSystemCaller} is used to distinguish whether this request is initiated by the
+        // system or a specific app.
+        // * System-initiated requests will only start the pinned mode (screen pinning)
+        // * App-initiated requests
+        //   - will put the device in fully locked mode (LockTask), if the app is whitelisted
+        //   - will start the pinned mode, otherwise
         final int callingUid = Binder.getCallingUid();
         long ident = Binder.clearCallingIdentity();
         try {
-            mLockTaskController.startLockTaskMode(task, isAppPinning, callingUid);
+            mLockTaskController.startLockTaskMode(task, isSystemCaller, callingUid);
         } finally {
             Binder.restoreCallingIdentity(ident);
         }
@@ -10839,7 +10842,7 @@
             if (r == null) {
                 return;
             }
-            startLockTaskModeLocked(r.getTask(), false /* not system initiated */);
+            startLockTaskModeLocked(r.getTask(), false /* isSystemCaller */);
         }
     }
 
@@ -10851,7 +10854,7 @@
         try {
             synchronized (this) {
                 startLockTaskModeLocked(mStackSupervisor.anyTaskForIdLocked(taskId),
-                        true /* system initiated */);
+                        true /* isSystemCaller */);
             }
         } finally {
             Binder.restoreCallingIdentity(ident);
@@ -10859,8 +10862,14 @@
     }
 
     @Override
-    public void stopLockTaskMode() {
-        stopLockTaskModeInternal(false /* not system initiated */);
+    public void stopLockTaskModeByToken(IBinder token) {
+        synchronized (this) {
+            final ActivityRecord r = ActivityRecord.forTokenLocked(token);
+            if (r == null) {
+                return;
+            }
+            stopLockTaskModeInternal(r.getTask(), false /* isSystemCaller */);
+        }
     }
 
     /**
@@ -10870,15 +10879,15 @@
     @Override
     public void stopSystemLockTaskMode() throws RemoteException {
         enforceCallingPermission(MANAGE_ACTIVITY_STACKS, "stopSystemLockTaskMode");
-        stopLockTaskModeInternal(true /* system initiated */);
+        stopLockTaskModeInternal(null, true /* isSystemCaller */);
     }
 
-    private void stopLockTaskModeInternal(boolean isSystemRequest) {
+    private void stopLockTaskModeInternal(@Nullable TaskRecord task, boolean isSystemCaller) {
         final int callingUid = Binder.getCallingUid();
         long ident = Binder.clearCallingIdentity();
         try {
             synchronized (this) {
-                mLockTaskController.stopLockTaskMode(isSystemRequest, callingUid);
+                mLockTaskController.stopLockTaskMode(task, isSystemCaller, callingUid);
             }
             // Launch in-call UI if a call is ongoing. This is necessary to allow stopping the lock
             // task and jumping straight into a call in the case of emergency call back.
diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index ba41bd4..88403fc 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -3474,7 +3474,7 @@
                 }
 
                 if (endTask) {
-                    mService.mLockTaskController.removeLockedTask(task);
+                    mService.mLockTaskController.clearLockedTask(task);
                 }
             } else if (r.state != ActivityState.PAUSING) {
                 // If the activity is PAUSING, we will complete the finish once
@@ -4334,8 +4334,9 @@
         }
         Slog.i(TAG, "moveTaskToBack: " + tr);
 
-        // If the task is locked, then show the lock task toast
-        if (mService.mLockTaskController.checkLockedTask(tr)) {
+        // In LockTask mode, moving a locked task to the back of the stack may expose unlocked
+        // ones. Therefore we need to check if this operation is allowed.
+        if (!mService.mLockTaskController.canMoveTaskToBack(tr)) {
             return false;
         }
 
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index 6ec158e..1ee946b 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -20,6 +20,7 @@
 import static android.Manifest.permission.INTERNAL_SYSTEM_WINDOW;
 import static android.Manifest.permission.START_ANY_ACTIVITY;
 import static android.Manifest.permission.START_TASKS_FROM_RECENTS;
+import static android.app.ActivityManager.LOCK_TASK_MODE_LOCKED;
 import static android.app.ActivityManager.START_TASK_TO_FRONT;
 import static android.app.ActivityManager.StackId.INVALID_STACK_ID;
 import static android.app.ITaskStackListener.FORCED_RESIZEABLE_REASON_SECONDARY_DISPLAY;
@@ -83,6 +84,7 @@
 import static com.android.server.am.TaskRecord.INVALID_TASK_ID;
 import static com.android.server.am.TaskRecord.LOCK_TASK_AUTH_LAUNCHABLE;
 import static com.android.server.am.TaskRecord.LOCK_TASK_AUTH_LAUNCHABLE_PRIV;
+import static com.android.server.am.TaskRecord.LOCK_TASK_AUTH_WHITELISTED;
 import static com.android.server.am.TaskRecord.REPARENT_KEEP_STACK_AT_FRONT;
 import static com.android.server.am.TaskRecord.REPARENT_LEAVE_STACK_IN_PLACE;
 import static com.android.server.am.TaskRecord.REPARENT_MOVE_STACK_TO_FRONT;
@@ -1306,8 +1308,11 @@
             mService.updateLruProcessLocked(app, true, null);
             mService.updateOomAdjLocked();
 
-            if (task.mLockTaskAuth == LOCK_TASK_AUTH_LAUNCHABLE ||
-                    task.mLockTaskAuth == LOCK_TASK_AUTH_LAUNCHABLE_PRIV) {
+            if (task.mLockTaskAuth == LOCK_TASK_AUTH_LAUNCHABLE
+                    || task.mLockTaskAuth == LOCK_TASK_AUTH_LAUNCHABLE_PRIV
+                    || (task.mLockTaskAuth == LOCK_TASK_AUTH_WHITELISTED
+                            && mService.mLockTaskController.getLockTaskModeState()
+                            == LOCK_TASK_MODE_LOCKED)) {
                 mService.mLockTaskController.startLockTaskMode(task, false, 0 /* blank UID */);
             }
 
@@ -2769,6 +2774,7 @@
         if (tr != null) {
             tr.removeTaskActivitiesLocked(pauseImmediately);
             cleanUpRemovedTaskLocked(tr, killProcess, removeFromRecents);
+            mService.mLockTaskController.clearLockedTask(tr);
             if (tr.isPersistable) {
                 mService.notifyTaskPersisterLocked(null, true);
             }
diff --git a/services/core/java/com/android/server/am/LockTaskController.java b/services/core/java/com/android/server/am/LockTaskController.java
index 4b2a084..e87b4e6 100644
--- a/services/core/java/com/android/server/am/LockTaskController.java
+++ b/services/core/java/com/android/server/am/LockTaskController.java
@@ -143,9 +143,17 @@
     LockTaskNotify mLockTaskNotify;
 
     /**
-     * The chain of tasks in lockTask mode. The current frontmost task is at the top, and tasks
-     * may be finished until there is only one entry left. If this is empty the system is not
-     * in lockTask mode.
+     * The chain of tasks in LockTask mode, in the order of when they first entered LockTask mode.
+     *
+     * The first task in the list, which started the current LockTask session, is called the root
+     * task. It coincides with the Home task in a typical multi-app kiosk deployment. When there are
+     * more than one locked tasks, the root task can't be finished. Nor can it be moved to the back
+     * of the stack by {@link ActivityStack#moveTaskToBackLocked(int)};
+     *
+     * Calling {@link Activity#stopLockTask()} on the root task will finish all tasks but itself in
+     * this list, and the device will exit LockTask mode.
+     *
+     * The list is empty if LockTask is inactive.
      */
     private final ArrayList<TaskRecord> mLockTaskModeTasks = new ArrayList<>();
 
@@ -164,7 +172,7 @@
      * {@link ActivityManager#LOCK_TASK_MODE_NONE}, {@link ActivityManager#LOCK_TASK_MODE_LOCKED},
      * {@link ActivityManager#LOCK_TASK_MODE_PINNED}
      */
-    private int mLockTaskModeState;
+    private int mLockTaskModeState = LOCK_TASK_MODE_NONE;
 
     /**
      * This is ActivityStackSupervisor's Handler.
@@ -199,8 +207,29 @@
      * @return whether the given task is locked at the moment. Locked tasks cannot be moved to the
      * back of the stack.
      */
-    boolean checkLockedTask(TaskRecord task) {
-        if (mLockTaskModeTasks.contains(task)) {
+    @VisibleForTesting
+    boolean isTaskLocked(TaskRecord task) {
+        return mLockTaskModeTasks.contains(task);
+    }
+
+    /**
+     * @return {@code true} whether this task first started the current LockTask session.
+     */
+    private boolean isRootTask(TaskRecord task) {
+        return mLockTaskModeTasks.indexOf(task) == 0;
+    }
+
+    /**
+     * @return whether the given activity is blocked from finishing, because it is the only activity
+     * of the last locked task and finishing it would mean that lock task mode is ended illegally.
+     */
+    boolean activityBlockedFromFinish(ActivityRecord activity) {
+        final TaskRecord task = activity.getTask();
+        if (activity == task.getRootActivity()
+                && activity == task.getTopActivity()
+                && task.mLockTaskAuth != LOCK_TASK_AUTH_LAUNCHABLE_PRIV
+                && isRootTask(task)) {
+            Slog.i(TAG, "Not finishing task in lock task mode");
             showLockTaskToast();
             return true;
         }
@@ -208,20 +237,16 @@
     }
 
     /**
-     * @return whether the given activity is blocked from finishing, because it is the root activity
-     * of the last locked task and finishing it would mean that lock task mode is ended illegally.
+     * @return whether the given task can be moved to the back of the stack with
+     * {@link ActivityStack#moveTaskToBackLocked(int)}
+     * @see #mLockTaskModeTasks
      */
-    boolean activityBlockedFromFinish(ActivityRecord activity) {
-        TaskRecord task = activity.getTask();
-        if (activity == task.getRootActivity()
-                && task.mLockTaskAuth != LOCK_TASK_AUTH_LAUNCHABLE_PRIV
-                && mLockTaskModeTasks.size() == 1
-                && mLockTaskModeTasks.contains(task)) {
-            Slog.i(TAG, "Not finishing task in lock task mode");
+    boolean canMoveTaskToBack(TaskRecord task) {
+        if (isRootTask(task)) {
             showLockTaskToast();
-            return true;
+            return false;
         }
-        return false;
+        return true;
     }
 
     /**
@@ -246,7 +271,7 @@
     private boolean isLockTaskModeViolationInternal(TaskRecord task, boolean isNewClearTask) {
         // TODO: Double check what's going on here. If the task is already in lock task mode, it's
         // likely whitelisted, so will return false below.
-        if (getLockedTask() == task && !isNewClearTask) {
+        if (isTaskLocked(task) && !isNewClearTask) {
             // If the task is already at the top and won't be cleared, then allow the operation
             return false;
         }
@@ -270,80 +295,116 @@
     /**
      * Stop the current lock task mode.
      *
-     * @param isSystemInitiated indicates whether this request was initiated by the system via
-     *                          {@link ActivityManagerService#stopSystemLockTaskMode()}.
+     * This is called by {@link ActivityManagerService} and performs various checks before actually
+     * finishing the locked task.
+     *
+     * @param task the task that requested the end of lock task mode ({@code null} for quitting app
+     *             pinning mode)
+     * @param isSystemCaller indicates whether this request comes from the system via
+     *                       {@link ActivityManagerService#stopSystemLockTaskMode()}. If
+     *                       {@code true}, it means the user intends to stop pinned mode through UI;
+     *                       otherwise, it's called by an app and we need to stop locked or pinned
+     *                       mode, subject to checks.
      * @param callingUid the caller that requested the end of lock task mode.
+     * @throws IllegalArgumentException if the calling task is invalid (e.g., {@code null} or not in
+     *                                  foreground)
      * @throws SecurityException if the caller is not authorized to stop the lock task mode, i.e. if
      *                           they differ from the one that launched lock task mode.
      */
-    void stopLockTaskMode(boolean isSystemInitiated, int callingUid) {
-        final TaskRecord lockTask = getLockedTask();
-        if (lockTask == null || mLockTaskModeState == LOCK_TASK_MODE_NONE) {
-            // Our work here is done.
+    void stopLockTaskMode(@Nullable TaskRecord task, boolean isSystemCaller, int callingUid) {
+        if (mLockTaskModeState == LOCK_TASK_MODE_NONE) {
             return;
         }
 
-        if (isSystemInitiated && mLockTaskModeState == LOCK_TASK_MODE_LOCKED) {
-            // As system can only start app pinning, we also only let it unlock in this mode.
-            showLockTaskToast();
+        if (isSystemCaller) {
+            if (mLockTaskModeState == LOCK_TASK_MODE_PINNED) {
+                clearLockedTasks("stopAppPinning");
+            } else {
+                Slog.e(TAG_LOCKTASK, "Attempted to stop LockTask with isSystemCaller=true");
+                showLockTaskToast();
+            }
+
+        } else {
+            // Ensure calling activity is not null
+            if (task == null) {
+                throw new IllegalArgumentException("can't stop LockTask for null task");
+            }
+
+            // Ensure the same caller for startLockTaskMode and stopLockTaskMode.
+            // It is possible lockTaskMode was started by the system process because
+            // android:lockTaskMode is set to a locking value in the application manifest
+            // instead of the app calling startLockTaskMode. In this case
+            // {@link TaskRecord.mLockTaskUid} will be 0, so we compare the callingUid to the
+            // {@link TaskRecord.effectiveUid} instead. Also caller with
+            // {@link MANAGE_ACTIVITY_STACKS} can stop any lock task.
+            if (callingUid != task.mLockTaskUid
+                    && (task.mLockTaskUid != 0 || callingUid != task.effectiveUid)) {
+                throw new SecurityException("Invalid uid, expected " + task.mLockTaskUid
+                        + " callingUid=" + callingUid + " effectiveUid=" + task.effectiveUid);
+            }
+
+            // We don't care if it's pinned or locked mode; this will stop it anyways.
+            clearLockedTask(task);
+        }
+    }
+
+    /**
+     * Clear all locked tasks and request the end of LockTask mode.
+     *
+     * This method is called by {@link UserController} when starting a new foreground user, and,
+     * unlike {@link #stopLockTaskMode(TaskRecord, boolean, int)}, it doesn't perform the checks.
+     */
+    void clearLockedTasks(String reason) {
+        if (DEBUG_LOCKTASK) Slog.i(TAG_LOCKTASK, "clearLockedTasks: " + reason);
+        if (!mLockTaskModeTasks.isEmpty()) {
+            clearLockedTask(mLockTaskModeTasks.get(0));
+        }
+    }
+
+    /**
+     * Clear one locked task from LockTask mode.
+     *
+     * If the requested task is the root task (see {@link #mLockTaskModeTasks}), then all locked
+     * tasks are cleared. Otherwise, only the requested task is cleared. LockTask mode is stopped
+     * when the last locked task is cleared.
+     *
+     * @param task the task to be cleared from LockTask mode.
+     */
+    void clearLockedTask(final TaskRecord task) {
+        if (task == null || mLockTaskModeTasks.isEmpty()) return;
+
+        if (task == mLockTaskModeTasks.get(0)) {
+            // We're removing the root task while there are other locked tasks. Therefore we should
+            // clear all locked tasks in reverse order.
+            for (int taskNdx = mLockTaskModeTasks.size() - 1; taskNdx > 0; --taskNdx) {
+                clearLockedTask(mLockTaskModeTasks.get(taskNdx));
+            }
+        }
+
+        removeLockedTask(task);
+        if (mLockTaskModeTasks.isEmpty()) {
             return;
         }
-
-        // Ensure the same caller for startLockTaskMode and stopLockTaskMode.
-        // It is possible lockTaskMode was started by the system process because
-        // android:lockTaskMode is set to a locking value in the application manifest
-        // instead of the app calling startLockTaskMode. In this case
-        // {@link TaskRecord.mLockTaskUid} will be 0, so we compare the callingUid to the
-        // {@link TaskRecord.effectiveUid} instead. Also caller with
-        // {@link MANAGE_ACTIVITY_STACKS} can stop any lock task.
-        if (!isSystemInitiated && callingUid != lockTask.mLockTaskUid
-                && (lockTask.mLockTaskUid != 0 || callingUid != lockTask.effectiveUid)) {
-            throw new SecurityException("Invalid uid, expected " + lockTask.mLockTaskUid
-                    + " callingUid=" + callingUid + " effectiveUid=" + lockTask.effectiveUid);
-        }
-
-        clearLockTaskMode("stopLockTask");
+        task.performClearTaskLocked();
+        mSupervisor.resumeFocusedStackTopActivityLocked();
     }
 
     /**
      * Remove the given task from the locked task list. If this was the last task in the list,
      * lock task mode is stopped.
      */
-    void removeLockedTask(final TaskRecord task) {
+    private void removeLockedTask(final TaskRecord task) {
         if (!mLockTaskModeTasks.remove(task)) {
             return;
         }
-        if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK, "removeLockedTask: removed " + task);
+        if (DEBUG_LOCKTASK) Slog.d(TAG_LOCKTASK, "removeLockedTask: removed " + task);
         if (mLockTaskModeTasks.isEmpty()) {
-            // Last one.
             if (DEBUG_LOCKTASK) Slog.d(TAG_LOCKTASK, "removeLockedTask: task=" + task +
                     " last task, reverting locktask mode. Callers=" + Debug.getCallers(3));
             mHandler.post(() -> performStopLockTask(task.userId));
         }
     }
 
-    /**
-     * Remove the topmost task from the locked task list. If this is the last task in the list, it
-     * will result in the end of locked task mode.
-     */
-    void clearLockTaskMode(String reason) {
-        // Take out of lock task mode if necessary
-        final TaskRecord lockedTask = getLockedTask();
-        if (lockedTask != null) {
-            removeLockedTask(lockedTask);
-            if (!mLockTaskModeTasks.isEmpty()) {
-                // There are locked tasks remaining, can only finish this task, not unlock it.
-                if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK,
-                        "setLockTaskMode: Tasks remaining, can't unlock");
-                lockedTask.performClearTaskLocked();
-                mSupervisor.resumeFocusedStackTopActivityLocked();
-                return;
-            }
-        }
-        if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK,
-                "setLockTaskMode: No tasks to unlock. Callers=" + Debug.getCallers(4));
-    }
-
     // This method should only be called on the handler thread
     private void performStopLockTask(int userId) {
         // When lock task ends, we enable the status bars.
@@ -382,17 +443,18 @@
      * Method to start lock task mode on a given task.
      *
      * @param task the task that should be locked.
-     * @param isSystemInitiated indicates whether this request was initiated by the system via
-     *                          {@link ActivityManagerService#startSystemLockTaskMode(int)}.
+     * @param isSystemCaller indicates whether this request was initiated by the system via
+     *                       {@link ActivityManagerService#startSystemLockTaskMode(int)}. If
+     *                       {@code true}, this intends to start pinned mode; otherwise, we look
+     *                       at the calling task's mLockTaskAuth to decide which mode to start.
      * @param callingUid the caller that requested the launch of lock task mode.
      */
-    void startLockTaskMode(@NonNull TaskRecord task, boolean isSystemInitiated,
-            int callingUid) {
-        if (!isSystemInitiated) {
+    void startLockTaskMode(@NonNull TaskRecord task, boolean isSystemCaller, int callingUid) {
+        if (!isSystemCaller) {
             task.mLockTaskUid = callingUid;
             if (task.mLockTaskAuth == LOCK_TASK_AUTH_PINNABLE) {
                 // startLockTask() called by app, but app is not part of lock task whitelist. Show
-                // app pinning request. We will come back here with isSystemInitiated true.
+                // app pinning request. We will come back here with isSystemCaller true.
                 if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK, "Mode default, asking user");
                 StatusBarManagerInternal statusBarManager = LocalServices.getService(
                         StatusBarManagerInternal.class);
@@ -404,8 +466,9 @@
         }
 
         // System can only initiate screen pinning, not full lock task mode
-        if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK, isSystemInitiated ? "Locking pinned" : "Locking fully");
-        setLockTaskMode(task, isSystemInitiated ? LOCK_TASK_MODE_PINNED : LOCK_TASK_MODE_LOCKED,
+        if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK,
+                isSystemCaller ? "Locking pinned" : "Locking fully");
+        setLockTaskMode(task, isSystemCaller ? LOCK_TASK_MODE_PINNED : LOCK_TASK_MODE_LOCKED,
                 "startLockTask", true);
     }
 
@@ -434,12 +497,12 @@
                     task.userId,
                     lockTaskModeState));
         }
-
-        // Add it or move it to the top.
         if (DEBUG_LOCKTASK) Slog.w(TAG_LOCKTASK, "setLockTaskMode: Locking to " + task +
                 " Callers=" + Debug.getCallers(4));
-        mLockTaskModeTasks.remove(task);
-        mLockTaskModeTasks.add(task);
+
+        if (!mLockTaskModeTasks.contains(task)) {
+            mLockTaskModeTasks.add(task);
+        }
 
         if (task.mLockTaskUid == -1) {
             task.mLockTaskUid = task.effectiveUid;
@@ -556,8 +619,7 @@
         }
 
         mLockTaskFeatures.put(userId, flags);
-        TaskRecord lockedTask = getLockedTask();
-        if (lockedTask != null && userId == lockedTask.userId) {
+        if (!mLockTaskModeTasks.isEmpty() && userId == mLockTaskModeTasks.get(0).userId) {
             mHandler.post(() -> {
                 if (mLockTaskModeState == LOCK_TASK_MODE_LOCKED) {
                     setStatusBarState(mLockTaskModeState, userId);
@@ -672,17 +734,6 @@
         return mLockTaskFeatures.get(userId, DevicePolicyManager.LOCK_TASK_FEATURE_NONE);
     }
 
-    /**
-     * @return the topmost locked task
-     */
-    private TaskRecord getLockedTask() {
-        final int top = mLockTaskModeTasks.size() - 1;
-        if (top >= 0) {
-            return mLockTaskModeTasks.get(top);
-        }
-        return null;
-    }
-
     // Should only be called on the handler thread
     @Nullable
     private IStatusBarService getStatusBarService() {
diff --git a/services/core/java/com/android/server/am/TaskRecord.java b/services/core/java/com/android/server/am/TaskRecord.java
index 1b5a1ce..1a4f9d4 100644
--- a/services/core/java/com/android/server/am/TaskRecord.java
+++ b/services/core/java/com/android/server/am/TaskRecord.java
@@ -438,7 +438,7 @@
     }
 
     void removeWindowContainer() {
-        mService.mLockTaskController.removeLockedTask(this);
+        mService.mLockTaskController.clearLockedTask(this);
         mWindowContainerController.removeContainer();
         if (!getWindowConfiguration().persistTaskBounds()) {
             // Reset current bounds for task whose bounds shouldn't be persisted so it uses
diff --git a/services/core/java/com/android/server/am/UserController.java b/services/core/java/com/android/server/am/UserController.java
index 44f83b0..2df5dc9 100644
--- a/services/core/java/com/android/server/am/UserController.java
+++ b/services/core/java/com/android/server/am/UserController.java
@@ -872,9 +872,7 @@
             }
 
             if (foreground) {
-                // TODO: I don't think this does what the caller think it does. Seems to only
-                // remove one locked task and won't work if multiple locked tasks are present.
-                mInjector.clearLockTaskMode("startUser");
+                mInjector.clearAllLockedTasks("startUser");
             }
 
             final UserInfo userInfo = getUserInfo(userId);
@@ -2053,9 +2051,9 @@
             }
         }
 
-        protected void clearLockTaskMode(String reason) {
+        protected void clearAllLockedTasks(String reason) {
             synchronized (mService) {
-                mService.mLockTaskController.clearLockTaskMode(reason);
+                mService.mLockTaskController.clearLockedTasks(reason);
             }
         }
     }