Affiliated profile owners can set lock task packages

Mentioned that in the documentation, cleaned up the code
a bit and added unit tests

Bug: 34614754

Test: runtest -c com.android.server.devicepolicy.DevicePolicyManagerTest    frameworks-services
Change-Id: I91232bbe494398015094ab977c6a2adce339811f
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 3c1d274..578e0b3 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -8564,18 +8564,11 @@
         }
     }
 
-    /**
-     * Sets which packages may enter lock task mode.
-     *
-     * <p>This function can only be called by the device owner or alternatively by the profile owner
-     * in case the user is affiliated.
-     *
-     * @param packages The list of packages allowed to enter lock task mode.
-     */
     @Override
     public void setLockTaskPackages(ComponentName who, String[] packages)
             throws SecurityException {
         Preconditions.checkNotNull(who, "ComponentName is null");
+        Preconditions.checkNotNull(packages, "packages is null");
 
         synchronized (this) {
             getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_PROFILE_OWNER);
@@ -8598,48 +8591,51 @@
         updateLockTaskPackagesLocked(packages, userHandle);
     }
 
+    private void maybeClearLockTaskPackagesLocked() {
+        final long ident = mInjector.binderClearCallingIdentity();
+        try {
+            final List<UserInfo> userInfos = mUserManager.getUsers(/*excludeDying=*/ true);
+            for (int i = 0; i < userInfos.size(); i++) {
+                int userId = userInfos.get(i).id;
+                final List<String> lockTaskPackages = getUserData(userId).mLockTaskPackages;
+                if (!lockTaskPackages.isEmpty() &&
+                        !isUserAffiliatedWithDeviceLocked(userId)) {
+                    Slog.d(LOG_TAG,
+                            "User id " + userId + " not affiliated. Clearing lock task packages");
+                    setLockTaskPackagesLocked(userId, Collections.<String>emptyList());
+                }
+            }
+        } finally {
+            mInjector.binderRestoreCallingIdentity(ident);
+        }
+    }
+
     /**
      * This function returns the list of components allowed to start the task lock mode.
      */
     @Override
     public String[] getLockTaskPackages(ComponentName who) {
         Preconditions.checkNotNull(who, "ComponentName is null");
+
+        final int userHandle = mInjector.binderGetCallingUserHandle().getIdentifier();
         synchronized (this) {
-            getActiveAdminForCallerLocked(who, DeviceAdminInfo.USES_POLICY_DEVICE_OWNER);
-            int userHandle = mInjector.binderGetCallingUserHandle().getIdentifier();
-            final List<String> packages = getLockTaskPackagesLocked(userHandle);
+            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.");
+            }
+
+            final List<String> packages = getUserData(userHandle).mLockTaskPackages;
             return packages.toArray(new String[packages.size()]);
         }
     }
 
-    private List<String> getLockTaskPackagesLocked(int userHandle) {
-        final DevicePolicyData policy = getUserData(userHandle);
-        return policy.mLockTaskPackages;
-    }
-
-    /**
-     * This function lets the caller know whether the given package is allowed to start the
-     * lock task mode.
-     * @param pkg The package to check
-     */
     @Override
     public boolean isLockTaskPermitted(String pkg) {
-        // Get current user's devicepolicy
-        int uid = mInjector.binderGetCallingUid();
-        int userHandle = UserHandle.getUserId(uid);
-        DevicePolicyData policy = getUserData(userHandle);
+        final int userHandle = mInjector.userHandleGetCallingUserId();
         synchronized (this) {
-            for (int i = 0; i < policy.mLockTaskPackages.size(); i++) {
-                String lockTaskPackage = policy.mLockTaskPackages.get(i);
-
-                // If the given package equals one of the packages stored our list,
-                // we allow this package to start lock task mode.
-                if (lockTaskPackage.equals(pkg)) {
-                    return true;
-                }
-            }
+            return getUserData(userHandle).mLockTaskPackages.contains(pkg);
         }
-        return false;
     }
 
     @Override
@@ -9848,6 +9844,7 @@
             // but as a result of that other users might become affiliated or un-affiliated.
             maybePauseDeviceWideLoggingLocked();
             maybeResumeDeviceWideLoggingLocked();
+            maybeClearLockTaskPackagesLocked();
         }
     }
 
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 f4e4e08..7df638c 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -3332,6 +3332,75 @@
         MoreAsserts.assertEmpty(targetUsers);
     }
 
+    public void testLockTaskPackagesAllowedForAffiliatedUsers() throws Exception {
+        // Setup a device owner.
+        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
+        setupDeviceOwner();
+        // Lock task packages are updated when loading user data.
+        verify(mContext.iactivityManager)
+                .updateLockTaskPackages(eq(UserHandle.USER_SYSTEM), eq(new String[0]));
+
+        // 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);
+        verify(mContext.iactivityManager)
+                .updateLockTaskPackages(eq(MANAGED_PROFILE_USER_ID), eq(new String[0]));
+
+        // 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(mContext.iactivityManager)
+                .updateLockTaskPackages(eq(UserHandle.USER_SYSTEM), eq(doPackages));
+
+        // Managed profile is unaffiliated - shouldn't be able to setLockTaskPackages.
+        mContext.binder.callingUid = MANAGED_PROFILE_ADMIN_UID;
+        final String[] poPackages = {"poPackage1", "poPackage2"};
+        try {
+            dpm.setLockTaskPackages(adminDifferentPackage, poPackages);
+            fail("Didn't throw expected security exception.");
+        } catch (SecurityException expected) {
+        }
+        try {
+            dpm.getLockTaskPackages(adminDifferentPackage);
+            fail("Didn't throw expected security exception.");
+        } catch (SecurityException expected) {
+        }
+        assertFalse(dpm.isLockTaskPermitted("doPackage1"));
+
+        // Setting same affiliation ids
+        final List<String> userAffiliationIds = Arrays.asList("some-affiliation-id");
+        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
+        dpm.setAffiliationIds(admin1, userAffiliationIds);
+
+        mContext.binder.callingUid = MANAGED_PROFILE_ADMIN_UID;
+        dpm.setAffiliationIds(adminDifferentPackage, userAffiliationIds);
+
+        // Now the managed profile can set lock task packages.
+        dpm.setLockTaskPackages(adminDifferentPackage, poPackages);
+        MoreAsserts.assertEquals(poPackages, dpm.getLockTaskPackages(adminDifferentPackage));
+        assertTrue(dpm.isLockTaskPermitted("poPackage1"));
+        assertFalse(dpm.isLockTaskPermitted("doPackage2"));
+        verify(mContext.iactivityManager)
+                .updateLockTaskPackages(eq(MANAGED_PROFILE_USER_ID), eq(poPackages));
+
+        // Unaffiliate the profile, lock task mode no longer available on the profile.
+        dpm.setAffiliationIds(adminDifferentPackage, Collections.<String>emptyList());
+        assertFalse(dpm.isLockTaskPermitted("poPackage1"));
+        // Lock task packages cleared when loading user data and when the user becomes unaffiliated.
+        verify(mContext.iactivityManager, times(2))
+                .updateLockTaskPackages(eq(MANAGED_PROFILE_USER_ID), eq(new String[0]));
+
+        mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
+        assertTrue(dpm.isLockTaskPermitted("doPackage1"));
+    }
+
     public void testIsDeviceManaged() throws Exception {
         mContext.binder.callingUid = DpmMockContext.CALLER_SYSTEM_USER_UID;
         setupDeviceOwner();