Allow unaffiliated Profile Owners to use lock task

Bug: 64948750
Test: bit FrameworksServicesTests:com.android.server.devicepolicy.DevicePolicyManagerTest
Change-Id: Ib775161dd62c819e7975744eaa074fc60d9a6b1b
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index bad19e3..1dc03474 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -7036,14 +7036,14 @@
      * task. From {@link android.os.Build.VERSION_CODES#M} removing packages from the lock task
      * package list results in locked tasks belonging to those packages to be finished.
      * <p>
-     * This function can only be called by the device owner or by a profile owner of a user/profile
-     * that is affiliated with the device. See {@link #isAffiliatedUser}. Any packages
-     * set via this method will be cleared if the user becomes unaffiliated.
+     * This function can only be called by the device owner, a profile owner of an affiliated user
+     * or profile, or the profile owner when no device owner is set. See {@link #isAffiliatedUser}.
+     * Any package set via this method will be cleared if the user becomes unaffiliated.
      *
      * @param packages The list of packages allowed to enter lock task mode
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with.
-     * @throws SecurityException if {@code admin} is not the device owner, or the profile owner of
-     * an affiliated user or profile.
+     * @throws SecurityException if {@code admin} is not the device owner, the profile owner of an
+     * affiliated user or profile, or the profile owner when no device owner is set.
      * @see #isAffiliatedUser
      * @see Activity#startLockTask()
      * @see DeviceAdminReceiver#onLockTaskModeEntering(Context, Intent, String)
@@ -7065,8 +7065,8 @@
     /**
      * Returns the list of packages allowed to start the lock task mode.
      *
-     * @throws SecurityException if {@code admin} is not the device owner, or the profile owner of
-     * an affiliated user or profile.
+     * @throws SecurityException if {@code admin} is not the device owner, the profile owner of an
+     * affiliated user or profile, or the profile owner when no device owner is set.
      * @see #isAffiliatedUser
      * @see #setLockTaskPackages
      */
@@ -7106,9 +7106,9 @@
      * is in LockTask mode. If this method is not called, none of the features listed here will be
      * enabled.
      * <p>
-     * This function can only be called by the device owner or by a profile owner of a user/profile
-     * that is affiliated with the device. See {@link #isAffiliatedUser}. Any features
-     * set via this method will be cleared if the user becomes unaffiliated.
+     * This function can only be called by the device owner, a profile owner of an affiliated user
+     * or profile, or the profile owner when no device owner is set. See {@link #isAffiliatedUser}.
+     * Any features set via this method will be cleared if the user becomes unaffiliated.
      *
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with.
      * @param flags Bitfield of feature flags:
@@ -7119,9 +7119,10 @@
      *              {@link #LOCK_TASK_FEATURE_RECENTS},
      *              {@link #LOCK_TASK_FEATURE_GLOBAL_ACTIONS},
      *              {@link #LOCK_TASK_FEATURE_KEYGUARD}
-     * @throws SecurityException if {@code admin} is not the device owner, or the profile owner of
-     * an affiliated user or profile.
+     * @throws SecurityException if {@code admin} is not the device owner, the profile owner of an
+     * affiliated user or profile, or the profile owner when no device owner is set.
      * @see #isAffiliatedUser
+     * @throws SecurityException if {@code admin} is not the device owner or the profile owner.
      */
     public void setLockTaskFeatures(@NonNull ComponentName admin, @LockTaskFeature int flags) {
         throwIfParentInstance("setLockTaskFeatures");
@@ -7139,8 +7140,8 @@
      *
      * @param admin Which {@link DeviceAdminReceiver} this request is associated with.
      * @return bitfield of flags. See {@link #setLockTaskFeatures(ComponentName, int)} for a list.
-     * @throws SecurityException if {@code admin} is not the device owner, or the profile owner of
-     * an affiliated user or profile.
+     * @throws SecurityException if {@code admin} is not the device owner, the profile owner of an
+     * affiliated user or profile, or the profile owner when no device owner is set.
      * @see #isAffiliatedUser
      * @see #setLockTaskFeatures
      */
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 4d6989b..8295675 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -7883,6 +7883,37 @@
         enforceSystemUserOrPermission(android.Manifest.permission.INTERACT_ACROSS_USERS_FULL);
     }
 
+    private boolean canUserUseLockTaskLocked(int userId) {
+        if (isUserAffiliatedWithDeviceLocked(userId)) {
+            return true;
+        }
+
+        // Unaffiliated profile owners are not allowed to use lock when there is a device owner.
+        if (mOwners.hasDeviceOwner()) {
+            return false;
+        }
+
+        final ComponentName profileOwner = getProfileOwner(userId);
+        if (profileOwner == null) {
+            return false;
+        }
+
+        // Managed profiles are not allowed to use lock task
+        if (isManagedProfile(userId)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    private void enforceCanCallLockTaskLocked(ComponentName who) {
+        getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
+        final int userId =  mInjector.userHandleGetCallingUserId();
+        if (!canUserUseLockTaskLocked(userId)) {
+            throw new SecurityException("User " + userId + " is not allowed to use lock task");
+        }
+    }
+
     private void ensureCallerPackage(@Nullable String packageName) {
         if (packageName == null) {
             Preconditions.checkState(isCallerWithSystemUid(),
@@ -9589,14 +9620,9 @@
         Preconditions.checkNotNull(packages, "packages is null");
 
         synchronized (this) {
-            getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
+            enforceCanCallLockTaskLocked(who);
             final int userHandle = mInjector.userHandleGetCallingUserId();
-            if (isUserAffiliatedWithDeviceLocked(userHandle)) {
-                setLockTaskPackagesLocked(userHandle, new ArrayList<>(Arrays.asList(packages)));
-            } else {
-                throw new SecurityException("Admin " + who +
-                    " is neither the device owner or affiliated user's profile owner.");
-            }
+            setLockTaskPackagesLocked(userHandle, new ArrayList<>(Arrays.asList(packages)));
         }
     }
 
@@ -9615,12 +9641,7 @@
 
         final int userHandle = mInjector.binderGetCallingUserHandle().getIdentifier();
         synchronized (this) {
-            getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
-            if (!isUserAffiliatedWithDeviceLocked(userHandle)) {
-                throw new SecurityException("Admin " + who +
-                    " is neither the device owner or affiliated user's profile owner.");
-            }
-
+            enforceCanCallLockTaskLocked(who);
             final List<String> packages = getUserData(userHandle).mLockTaskPackages;
             return packages.toArray(new String[packages.size()]);
         }
@@ -9639,11 +9660,7 @@
         Preconditions.checkNotNull(who, "ComponentName is null");
         final int userHandle = mInjector.userHandleGetCallingUserId();
         synchronized (this) {
-            getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
-            if (!isUserAffiliatedWithDeviceLocked(userHandle)) {
-                throw new SecurityException("Admin " + who +
-                        " is neither the device owner or affiliated user's profile owner.");
-            }
+            enforceCanCallLockTaskLocked(who);
             setLockTaskFeaturesLocked(userHandle, flags);
         }
     }
@@ -9660,11 +9677,7 @@
         Preconditions.checkNotNull(who, "ComponentName is null");
         final int userHandle = mInjector.userHandleGetCallingUserId();
         synchronized (this) {
-            getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
-            if (!isUserAffiliatedWithDeviceLocked(userHandle)) {
-                throw new SecurityException("Admin " + who +
-                        " is neither the device owner or affiliated user's profile owner.");
-            }
+            enforceCanCallLockTaskLocked(who);
             return getUserData(userHandle).mLockTaskFeatures;
         }
     }
@@ -9675,7 +9688,7 @@
             final List<UserInfo> userInfos = mUserManager.getUsers(/*excludeDying=*/ true);
             for (int i = userInfos.size() - 1; i >= 0; i--) {
                 int userId = userInfos.get(i).id;
-                if (isUserAffiliatedWithDeviceLocked(userId)) {
+                if (canUserUseLockTaskLocked(userId)) {
                     continue;
                 }
 
@@ -11213,10 +11226,12 @@
             // of a split user device.
             return true;
         }
+
         final ComponentName profileOwner = getProfileOwner(userId);
         if (profileOwner == null) {
             return false;
         }
+        
         final Set<String> userAffiliationIds = getUserData(userId).mAffiliationIds;
         final Set<String> deviceAffiliationIds =
                 getUserData(UserHandle.USER_SYSTEM).mAffiliationIds;
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
index e8d6ed2..741a169 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -3652,15 +3652,47 @@
         MoreAsserts.assertEmpty(targetUsers);
     }
 
+    private void verifyLockTaskState(int userId) throws Exception {
+        verifyLockTaskState(userId, new String[0], DevicePolicyManager.LOCK_TASK_FEATURE_NONE);
+    }
+
+    private void verifyLockTaskState(int userId, String[] packages, int flags) throws Exception {
+        verify(getServices().iactivityManager).updateLockTaskPackages(userId, packages);
+        verify(getServices().iactivityManager).updateLockTaskFeatures(userId, flags);
+    }
+
+    private void verifyCanSetLockTask(int uid, int userId, ComponentName who, String[] packages,
+            int flags) throws Exception {
+        mContext.binder.callingUid = uid;
+        dpm.setLockTaskPackages(who, packages);
+        MoreAsserts.assertEquals(packages, dpm.getLockTaskPackages(who));
+        for (String p : packages) {
+            assertTrue(dpm.isLockTaskPermitted(p));
+        }
+        assertFalse(dpm.isLockTaskPermitted("anotherPackage"));
+        // Test to see if set lock task features can be set
+        dpm.setLockTaskFeatures(who, flags);
+        verifyLockTaskState(userId, packages, flags);
+    }
+
+    private void verifyCanNotSetLockTask(int uid, ComponentName who, String[] packages,
+            int flags) throws Exception {
+        mContext.binder.callingUid = uid;
+        assertExpectException(SecurityException.class, /* messageRegex =*/ null,
+                () -> dpm.setLockTaskPackages(who, packages));
+        assertExpectException(SecurityException.class, /* messageRegex =*/ null,
+                () -> dpm.getLockTaskPackages(who));
+        assertFalse(dpm.isLockTaskPermitted("doPackage1"));
+        assertExpectException(SecurityException.class, /* messageRegex =*/ null,
+                () -> dpm.setLockTaskFeatures(who, flags));
+    }
+
     public void testLockTaskPolicyAllowedForAffiliatedUsers() throws Exception {
         // Setup a device owner.
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         setupDeviceOwner();
         // Lock task policy is updated when loading user data.
-        verify(getServices().iactivityManager).updateLockTaskPackages(
-                UserHandle.USER_SYSTEM, new String[0]);
-        verify(getServices().iactivityManager).updateLockTaskFeatures(
-                UserHandle.USER_SYSTEM, DevicePolicyManager.LOCK_TASK_FEATURE_NONE);
+        verifyLockTaskState(UserHandle.USER_SYSTEM);
 
         // Set up a managed profile managed by different package (package name shouldn't matter)
         final int MANAGED_PROFILE_USER_ID = 15;
@@ -3668,40 +3700,30 @@
         final ComponentName adminDifferentPackage =
                 new ComponentName("another.package", "whatever.class");
         addManagedProfile(adminDifferentPackage, MANAGED_PROFILE_ADMIN_UID, admin2);
-        verify(getServices().iactivityManager).updateLockTaskPackages(
-                MANAGED_PROFILE_USER_ID, new String[0]);
-        verify(getServices().iactivityManager).updateLockTaskFeatures(
-                MANAGED_PROFILE_USER_ID, DevicePolicyManager.LOCK_TASK_FEATURE_NONE);
+        verifyLockTaskState(MANAGED_PROFILE_USER_ID);
+
+        // Setup a PO on the secondary user
+        mContext.binder.callingUid = DpmMockContext.CALLER_UID;
+        setAsProfileOwner(admin3);
+        verifyLockTaskState(DpmMockContext.CALLER_USER_HANDLE);
 
         // The DO can still set lock task packages
-        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         final String[] doPackages = {"doPackage1", "doPackage2"};
-        dpm.setLockTaskPackages(admin1, doPackages);
-        MoreAsserts.assertEquals(doPackages, dpm.getLockTaskPackages(admin1));
-        assertTrue(dpm.isLockTaskPermitted("doPackage1"));
-        assertFalse(dpm.isLockTaskPermitted("anotherPackage"));
-        verify(getServices().iactivityManager).updateLockTaskPackages(
-                UserHandle.USER_SYSTEM, doPackages);
-        // And the DO can still set lock task features
-        final int doFlags = DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS
+        final int flags = DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS
                 | DevicePolicyManager.LOCK_TASK_FEATURE_RECENTS;
-        dpm.setLockTaskFeatures(admin1, doFlags);
-        verify(getServices().iactivityManager).updateLockTaskFeatures(
-                UserHandle.USER_SYSTEM, doFlags);
+        verifyCanSetLockTask(DpmMockContext.CALLER_SYSTEM_USER_UID, UserHandle.USER_SYSTEM, admin1, doPackages, flags);
+
+        final String[] secondaryPoPackages = {"secondaryPoPackage1", "secondaryPoPackage2"};
+        final int secondaryPoFlags = DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS
+                | DevicePolicyManager.LOCK_TASK_FEATURE_RECENTS;
+        verifyCanNotSetLockTask(DpmMockContext.CALLER_UID, admin3, secondaryPoPackages, secondaryPoFlags);
 
         // Managed profile is unaffiliated - shouldn't be able to setLockTaskPackages.
         mContext.binder.callingUid = MANAGED_PROFILE_ADMIN_UID;
         final String[] poPackages = {"poPackage1", "poPackage2"};
-        assertExpectException(SecurityException.class, /* messageRegex =*/ null,
-                () -> dpm.setLockTaskPackages(adminDifferentPackage, poPackages));
-        assertExpectException(SecurityException.class, /* messageRegex =*/ null,
-                () -> dpm.getLockTaskPackages(adminDifferentPackage));
-        assertFalse(dpm.isLockTaskPermitted("doPackage1"));
-        // And it shouldn't be able to setLockTaskFeatures.
         final int poFlags = DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS
                 | DevicePolicyManager.LOCK_TASK_FEATURE_RECENTS;
-        assertExpectException(SecurityException.class, /* messageRegex =*/ null,
-                () -> dpm.setLockTaskFeatures(adminDifferentPackage, poFlags));
+        verifyCanNotSetLockTask(MANAGED_PROFILE_ADMIN_UID, adminDifferentPackage, poPackages, poFlags);
 
         // Setting same affiliation ids
         final Set<String> userAffiliationIds = Collections.singleton("some-affiliation-id");
@@ -3716,12 +3738,9 @@
         MoreAsserts.assertEquals(poPackages, dpm.getLockTaskPackages(adminDifferentPackage));
         assertTrue(dpm.isLockTaskPermitted("poPackage1"));
         assertFalse(dpm.isLockTaskPermitted("doPackage2"));
-        verify(getServices().iactivityManager).updateLockTaskPackages(
-                MANAGED_PROFILE_USER_ID, poPackages);
         // And it can set lock task features.
         dpm.setLockTaskFeatures(adminDifferentPackage, poFlags);
-        verify(getServices().iactivityManager).updateLockTaskFeatures(
-                MANAGED_PROFILE_USER_ID, poFlags);
+        verifyLockTaskState(MANAGED_PROFILE_USER_ID, poPackages, poFlags);
 
         // Unaffiliate the profile, lock task mode no longer available on the profile.
         dpm.setAffiliationIds(adminDifferentPackage, Collections.emptySet());
@@ -3732,8 +3751,38 @@
         verify(getServices().iactivityManager, times(2)).updateLockTaskFeatures(
                 MANAGED_PROFILE_USER_ID, DevicePolicyManager.LOCK_TASK_FEATURE_NONE);
 
+        // Verify that lock task packages were not cleared for the DO
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         assertTrue(dpm.isLockTaskPermitted("doPackage1"));
+
+    }
+
+    public void testLockTaskPolicyForProfileOwner() throws Exception {
+        // Setup a PO
+        mContext.binder.callingUid = DpmMockContext.CALLER_UID;
+        setAsProfileOwner(admin1);
+        verifyLockTaskState(DpmMockContext.CALLER_USER_HANDLE);
+
+        final String[] poPackages = {"poPackage1", "poPackage2"};
+        final int poFlags = DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS
+                | DevicePolicyManager.LOCK_TASK_FEATURE_RECENTS;
+        verifyCanSetLockTask(DpmMockContext.CALLER_UID, DpmMockContext.CALLER_USER_HANDLE, admin1,
+                poPackages, poFlags);
+
+        // Set up a managed profile managed by different package (package name shouldn't matter)
+        final int MANAGED_PROFILE_USER_ID = 15;
+        final int MANAGED_PROFILE_ADMIN_UID = UserHandle.getUid(MANAGED_PROFILE_USER_ID, 20456);
+        final ComponentName adminDifferentPackage =
+                new ComponentName("another.package", "whatever.class");
+        addManagedProfile(adminDifferentPackage, MANAGED_PROFILE_ADMIN_UID, admin2);
+        verifyLockTaskState(MANAGED_PROFILE_USER_ID);
+
+        // Managed profile is unaffiliated - shouldn't be able to setLockTaskPackages.
+        mContext.binder.callingUid = MANAGED_PROFILE_ADMIN_UID;
+        final String[] mpoPackages = {"poPackage1", "poPackage2"};
+        final int mpoFlags = DevicePolicyManager.LOCK_TASK_FEATURE_NOTIFICATIONS
+                | DevicePolicyManager.LOCK_TASK_FEATURE_RECENTS;
+        verifyCanNotSetLockTask(MANAGED_PROFILE_ADMIN_UID, adminDifferentPackage, mpoPackages, mpoFlags);
     }
 
     public void testIsDeviceManaged() throws Exception {