From c83234a0ef4180343af8cc3cb0edcdf92eee2a46 Mon Sep 17 00:00:00 2001
From: Louis Chang <louischang@google.com>
Date: Mon, 22 Mar 2021 21:27:58 +0800
Subject: [PATCH] Detects all activities for whether showing work challenge

Work challenge did not show when a work activity is not on top, but
still visible after screen turns on.

Also show work challenge even if the work activity is behind a top
fullscreen activity of another profile because the user can still
navigate back to the work activity when top activity finishes.

Bug: 168527567
Test: manually test work challenges
Test: RootWindowContainerTests

Change-Id: I5e09b09be547d04fdfd709cb9cd4bcd4a94bbf21
---
 .../server/wm/RootWindowContainer.java        | 43 ++++---------------
 .../server/wm/RootWindowContainerTests.java   | 19 ++++++++
 2 files changed, 27 insertions(+), 35 deletions(-)

diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 72610481ad328..52551ec54b320 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -3363,46 +3363,19 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
     }
 
     /**
-     * Find all visible tasks containing {@param userId} and intercept them with an activity
+     * Find all tasks 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.
      */
     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();
-        } 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>
-     */
-    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);
-        }
+        forAllLeafTasks(task -> {
+            if (task.getActivity(activity -> !activity.finishing && activity.mUserId == userId)
+                    != null) {
+                mService.getTaskChangeNotificationController().notifyTaskProfileLocked(
+                        task.mTaskId, userId);
+            }
+        }, true /* traverseTopToBottom */);
     }
 
     void cancelInitializingActivities() {
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 20bced2527641..0bf237dc6545c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
@@ -70,6 +70,7 @@ import android.content.pm.ApplicationInfo;
 import android.content.pm.ResolveInfo;
 import android.content.res.Configuration;
 import android.content.res.Resources;
+import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
 import android.util.MergedConfiguration;
 import android.util.Pair;
@@ -1026,6 +1027,24 @@ public class RootWindowContainerTests extends WindowTestsBase {
         assertEquals(taskDisplayArea.getTopRootTask(), taskDisplayArea.getRootHomeTask());
     }
 
+    @Test
+    public void testLockAllProfileTasks() {
+        // Make an activity visible with the user id set to 0
+        final Task task = new TaskBuilder(mSupervisor).setCreateActivity(true).build();
+        final int taskId = task.mTaskId;
+        final ActivityRecord activity = task.getTopMostActivity();
+
+        // Create another activity on top and the user id is 1
+        final ActivityRecord topActivity = new ActivityBuilder(mAtm).setTask(task)
+                .setUid(UserHandle.PER_USER_RANGE + 1).build();
+
+        // Make sure the listeners will be notified for putting the task to locked state
+        TaskChangeNotificationController controller = mAtm.getTaskChangeNotificationController();
+        spyOn(controller);
+        mWm.mRoot.lockAllProfileTasks(0);
+        verify(controller).notifyTaskProfileLocked(eq(taskId), eq(0));
+    }
+
     /**
      * Mock {@link RootWindowContainer#resolveHomeActivity} for returning consistent activity
      * info for test cases.