Revert "Revert "Detects all activities for whether showing work challenge""

This reverts commit 2992a1dac025074726364d522cf9eac586f3d2c3.

* This commit is present in android-11.0.0_r40, but not in
  android-11.0.0_r43 since that's on top of android-11.0.0_r39
* Just revert this to match what's shipping in r43

Change-Id: I8fe467db3bbf5aa40c10c4e1311829a733373279

2 files changed, 41 insertions, 32 deletions

diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 93f6cdfed7d6..a0a071172606 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -3391,7 +3391,7 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     }
 
     /**
-     * Find all visible task stacks containing {@param userId} and intercept them with an activity
+     * Find all task stacks containing {@param userId} and intercept them with an activity
      * to block out the contents and possibly start a credential-confirming intent.
      *
      * @param userId user handle for the locked managed profile.
@@ -3399,42 +3399,18 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     void lockAllProfileTasks(@UserIdInt int userId) {
         mService.deferWindowLayout();
         try {
-            final PooledConsumer c = PooledLambda.obtainConsumer(
-                    RootWindowContainer::taskTopActivityIsUser, this, PooledLambda.__(Task.class),
-                    userId);
-            forAllLeafTasks(c, true /* traverseTopToBottom */);
-            c.recycle();
+            forAllLeafTasks(task -> {
+                if (task.getActivity(activity -> !activity.finishing && activity.mUserId == userId)
+                        != null) {
+                    mService.getTaskChangeNotificationController().notifyTaskProfileLocked(
+                            task.mTaskId, userId);
+                }
+            }, true /* traverseTopToBottom */);
         } finally {
             mService.continueWindowLayout();
         }
     }
 
-    /**
-     * Detects whether we should show a lock screen in front of this task for a locked user.
-     * <p>
-     * We'll do this if either of the following holds:
-     * <ul>
-     *   <li>The top activity explicitly belongs to {@param userId}.</li>
-     *   <li>The top activity returns a result to an activity belonging to {@param userId}.</li>
-     * </ul>
-     *
-     * @return {@code true} if the top activity looks like it belongs to {@param userId}.
-     */
-    private void taskTopActivityIsUser(Task task, @UserIdInt int userId) {
-        // To handle the case that work app is in the task but just is not the top one.
-        final ActivityRecord activityRecord = task.getTopNonFinishingActivity();
-        final ActivityRecord resultTo = (activityRecord != null ? activityRecord.resultTo : null);
-
-        // Check the task for a top activity belonging to userId, or returning a
-        // result to an activity belonging to userId. Example case: a document
-        // picker for personal files, opened by a work app, should still get locked.
-        if ((activityRecord != null && activityRecord.mUserId == userId)
-                || (resultTo != null && resultTo.mUserId == userId)) {
-            mService.getTaskChangeNotificationController().notifyTaskProfileLocked(
-                    task.mTaskId, userId);
-        }
-    }
-
     void cancelInitializingActivities() {
         for (int displayNdx = getChildCount() - 1; displayNdx >= 0; --displayNdx) {
             final DisplayContent display = getChildAt(displayNdx);
diff --git a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
index 35d1b17d5822..1aff8a7b5382 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
@@ -25,6 +25,7 @@ import static android.view.WindowManager.LayoutParams.TYPE_NOTIFICATION_SHADE;
 import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_TOAST;
 
+import static com.android.dx.mockito.inline.extended.ExtendedMockito.spyOn;
 import static com.android.server.wm.ActivityStack.ActivityState.FINISHING;
 import static com.android.server.wm.ActivityStack.ActivityState.PAUSED;
 import static com.android.server.wm.ActivityStack.ActivityState.PAUSING;
@@ -36,10 +37,13 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.verify;
 
 import android.app.WindowConfiguration;
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
+import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
 
 import androidx.test.filters.SmallTest;
@@ -169,5 +173,34 @@ public class RootWindowContainerTests extends WindowTestsBase {
         activity.setState(FINISHING, "test FINISHING");
         assertThat(mWm.mRoot.allPausedActivitiesComplete()).isTrue();
     }
+
+    @Test
+    public void testLockAllProfileTasks() {
+        // Make an activity visible with the user id set to 0
+        DisplayContent displayContent = mWm.mRoot.getDisplayContent(DEFAULT_DISPLAY);
+        TaskDisplayArea taskDisplayArea = displayContent.getTaskDisplayAreaAt(0);
+        final ActivityStack stack = createTaskStackOnDisplay(WINDOWING_MODE_FULLSCREEN,
+                ACTIVITY_TYPE_STANDARD, displayContent);
+        final ActivityRecord activity = new ActivityTestsBase.ActivityBuilder(stack.mAtmService)
+                .setStack(stack)
+                .setUid(0)
+                .setCreateTask(true)
+                .build();
+
+        // Create another activity on top and the user id is 1
+        Task task = activity.getTask();
+        final ActivityRecord topActivity = new ActivityTestsBase.ActivityBuilder(mWm.mAtmService)
+                .setStack(stack)
+                .setUid(UserHandle.PER_USER_RANGE + 1)
+                .setTask(task)
+                .build();
+
+        // Make sure the listeners will be notified for putting the task to locked state
+        TaskChangeNotificationController controller =
+                mWm.mAtmService.getTaskChangeNotificationController();
+        spyOn(controller);
+        mWm.mRoot.lockAllProfileTasks(0);
+        verify(controller).notifyTaskProfileLocked(eq(task.mTaskId), eq(0));
+    }
 }