Enforce DISALLOW_ADD_MANAGED_PROFILE

Only the device owner should be able to create a managed
profile if that restriction is set

Test: runtest -c com.android.server.devicepolicy.DevicePolicyManagerTest    frameworks-services

Bug: 31952368

Change-Id: Ia5170e54594ccba1e5bcedffaec98c2af42264c0
diff --git a/core/java/android/app/admin/DevicePolicyManager.java b/core/java/android/app/admin/DevicePolicyManager.java
index b641e63..e0b631e 100644
--- a/core/java/android/app/admin/DevicePolicyManager.java
+++ b/core/java/android/app/admin/DevicePolicyManager.java
@@ -385,7 +385,7 @@
             "com.android.server.action.BUGREPORT_SHARING_DECLINED";
 
     /**
-     * Action: Bugreport has been collected and is dispatched to {@link DevicePolicyManagerService}.
+     * Action: Bugreport has been collected and is dispatched to {@code DevicePolicyManagerService}.
      *
      * @hide
      */
@@ -1165,7 +1165,7 @@
     public @interface UserProvisioningState {}
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE},
      * {@link #ACTION_PROVISION_MANAGED_PROFILE}, {@link #ACTION_PROVISION_MANAGED_USER} and
@@ -1176,7 +1176,7 @@
     public static final int CODE_OK = 0;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE} and
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} when the device already has a device
@@ -1187,7 +1187,7 @@
     public static final int CODE_HAS_DEVICE_OWNER = 1;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE},
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} when the user has a profile owner and for
@@ -1198,7 +1198,7 @@
     public static final int CODE_USER_HAS_PROFILE_OWNER = 2;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE} and
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} when the user isn't running.
@@ -1208,7 +1208,7 @@
     public static final int CODE_USER_NOT_RUNNING = 3;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE},
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} if the device has already been setup and
@@ -1233,7 +1233,7 @@
     public static final int CODE_ACCOUNTS_NOT_EMPTY = 6;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE} and
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} if the user is not a system user.
@@ -1243,7 +1243,7 @@
     public static final int CODE_NOT_SYSTEM_USER = 7;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE},
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} and {@link #ACTION_PROVISION_MANAGED_USER}
@@ -1254,7 +1254,7 @@
     public static final int CODE_HAS_PAIRED = 8;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_PROFILE} and
      * {@link #ACTION_PROVISION_MANAGED_USER} on devices which do not support managed users.
@@ -1265,7 +1265,7 @@
     public static final int CODE_MANAGED_USERS_NOT_SUPPORTED = 9;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_USER} if the user is a system user.
      *
@@ -1274,7 +1274,7 @@
     public static final int CODE_SYSTEM_USER = 10;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_PROFILE} when the user cannot have more
      * managed profiles.
@@ -1284,7 +1284,7 @@
     public static final int CODE_CANNOT_ADD_MANAGED_PROFILE = 11;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_USER} and
      * {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE} on devices not running with split system
@@ -1295,7 +1295,7 @@
     public static final int CODE_NOT_SYSTEM_USER_SPLIT = 12;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_DEVICE},
      * {@link #ACTION_PROVISION_MANAGED_PROFILE}, {@link #ACTION_PROVISION_MANAGED_USER} and
@@ -1307,7 +1307,7 @@
     public static final int CODE_DEVICE_ADMIN_NOT_SUPPORTED = 13;
 
     /**
-     * Result code for {@link checkProvisioningPreCondition}.
+     * Result code for {@link #checkProvisioningPreCondition}.
      *
      * <p>Returned for {@link #ACTION_PROVISION_MANAGED_PROFILE} when the device the user is a
      * system user on a split system user device.
@@ -1317,7 +1317,17 @@
     public static final int CODE_SPLIT_SYSTEM_USER_DEVICE_SYSTEM_USER = 14;
 
     /**
-     * Result codes for {@link checkProvisioningPreCondition} indicating all the provisioning pre
+     * Result code for {@link #checkProvisioningPreCondition}.
+     *
+     * <p>Returned for {@link #ACTION_PROVISION_MANAGED_PROFILE} when adding a managed profile is
+     * disallowed by {@link UserManager#DISALLOW_ADD_MANAGED_PROFILE}.
+     *
+     * @hide
+     */
+    public static final int CODE_ADD_MANAGED_PROFILE_DISALLOWED = 15;
+
+    /**
+     * Result codes for {@link #checkProvisioningPreCondition} indicating all the provisioning pre
      * conditions.
      *
      * @hide
@@ -1327,7 +1337,7 @@
             CODE_USER_SETUP_COMPLETED, CODE_NOT_SYSTEM_USER, CODE_HAS_PAIRED,
             CODE_MANAGED_USERS_NOT_SUPPORTED, CODE_SYSTEM_USER, CODE_CANNOT_ADD_MANAGED_PROFILE,
             CODE_NOT_SYSTEM_USER_SPLIT, CODE_DEVICE_ADMIN_NOT_SUPPORTED,
-            CODE_SPLIT_SYSTEM_USER_DEVICE_SYSTEM_USER})
+            CODE_SPLIT_SYSTEM_USER_DEVICE_SYSTEM_USER, CODE_ADD_MANAGED_PROFILE_DISALLOWED})
     public @interface ProvisioningPreCondition {}
 
     /**
@@ -6184,34 +6194,40 @@
     }
 
     /**
-     * Returns if provisioning a managed profile or device is possible or not.
+     * Returns whether it is possible for the caller to initiate provisioning of a managed profile
+     * or device, setting itself as the device or profile owner.
+     *
      * @param action One of {@link #ACTION_PROVISION_MANAGED_DEVICE},
      * {@link #ACTION_PROVISION_MANAGED_PROFILE}.
-     * @return if provisioning a managed profile or device is possible or not.
+     * @return whether provisioning a managed profile or device is possible.
      * @throws IllegalArgumentException if the supplied action is not valid.
      */
-    public boolean isProvisioningAllowed(String action) {
+    public boolean isProvisioningAllowed(@NonNull String action) {
         throwIfParentInstance("isProvisioningAllowed");
         try {
-            return mService.isProvisioningAllowed(action);
+            return mService.isProvisioningAllowed(action, mContext.getPackageName());
         } catch (RemoteException re) {
             throw re.rethrowFromSystemServer();
         }
     }
 
     /**
-     * Checks if provisioning a managed profile or device is possible and returns one of the
-     * {@link ProvisioningPreCondition}.
+     * Checks whether it is possible to initiate provisioning a managed device,
+     * profile or user, setting the given package as owner.
      *
      * @param action One of {@link #ACTION_PROVISION_MANAGED_DEVICE},
      *        {@link #ACTION_PROVISION_MANAGED_PROFILE},
      *        {@link #ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE},
      *        {@link #ACTION_PROVISION_MANAGED_USER}
+     * @param packageName The package of the component that would be set as device, user, or profile
+     *        owner.
+     * @return A {@link ProvisioningPreCondition} value indicating whether provisioning is allowed.
      * @hide
      */
-    public @ProvisioningPreCondition int checkProvisioningPreCondition(String action) {
+    public @ProvisioningPreCondition int checkProvisioningPreCondition(
+            String action, @NonNull String packageName) {
         try {
-            return mService.checkProvisioningPreCondition(action);
+            return mService.checkProvisioningPreCondition(action, packageName);
         } catch (RemoteException re) {
             throw re.rethrowFromSystemServer();
         }
@@ -6940,8 +6956,8 @@
      * @hide
      * Force update user setup completed status. This API has no effect on user build.
      * @throws {@link SecurityException} if the caller has no
-     *         {@link android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS} or the caller is
-     *         not {@link UserHandle.SYSTEM_USER}
+     *         {@code android.Manifest.permission.MANAGE_PROFILE_AND_DEVICE_OWNERS} or the caller is
+     *         not {@link UserHandle#SYSTEM_USER}
      */
     public void forceUpdateUserSetupComplete() {
         try {
diff --git a/core/java/android/app/admin/IDevicePolicyManager.aidl b/core/java/android/app/admin/IDevicePolicyManager.aidl
index 3e22825..afb426c 100644
--- a/core/java/android/app/admin/IDevicePolicyManager.aidl
+++ b/core/java/android/app/admin/IDevicePolicyManager.aidl
@@ -269,8 +269,8 @@
     boolean setPermissionGrantState(in ComponentName admin, String packageName,
             String permission, int grantState);
     int getPermissionGrantState(in ComponentName admin, String packageName, String permission);
-    boolean isProvisioningAllowed(String action);
-    int checkProvisioningPreCondition(String action);
+    boolean isProvisioningAllowed(String action, String packageName);
+    int checkProvisioningPreCondition(String action, String packageName);
     void setKeepUninstalledPackages(in ComponentName admin,in List<String> packageList);
     List<String> getKeepUninstalledPackages(in ComponentName admin);
     boolean isManagedProfile(in ComponentName admin);
diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 3281bd6..8f46414 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -18,6 +18,7 @@
 
 import static android.Manifest.permission.MANAGE_CA_CERTIFICATES;
 import static android.app.admin.DevicePolicyManager.CODE_ACCOUNTS_NOT_EMPTY;
+import static android.app.admin.DevicePolicyManager.CODE_ADD_MANAGED_PROFILE_DISALLOWED;
 import static android.app.admin.DevicePolicyManager.CODE_CANNOT_ADD_MANAGED_PROFILE;
 import static android.app.admin.DevicePolicyManager.CODE_DEVICE_ADMIN_NOT_SUPPORTED;
 import static android.app.admin.DevicePolicyManager.CODE_HAS_DEVICE_OWNER;
@@ -1689,9 +1690,9 @@
 
         mSecurityLogMonitor = new SecurityLogMonitor(this);
 
-        mHasFeature = mContext.getPackageManager()
+        mHasFeature = mInjector.getPackageManager()
                 .hasSystemFeature(PackageManager.FEATURE_DEVICE_ADMIN);
-        mIsWatch = mContext.getPackageManager()
+        mIsWatch = mInjector.getPackageManager()
                 .hasSystemFeature(PackageManager.FEATURE_WATCH);
         if (!mHasFeature) {
             // Skip the rest of the initialization
@@ -4460,7 +4461,7 @@
             }
 
             try {
-                int uid = mContext.getPackageManager().getPackageUidAsUser(
+                int uid = mInjector.getPackageManager().getPackageUidAsUser(
                         policy.mDelegatedCertInstallerPackage, userHandle);
                 return uid == callingUid;
             } catch (NameNotFoundException e) {
@@ -6021,6 +6022,14 @@
         }
     }
 
+    private boolean isDeviceOwnerPackage(String packageName, int userId) {
+        synchronized (this) {
+            return mOwners.hasDeviceOwner()
+                    && mOwners.getDeviceOwnerUserId() == userId
+                    && mOwners.getDeviceOwnerPackageName().equals(packageName);
+        }
+    }
+
     public boolean isProfileOwner(ComponentName who, int userId) {
         final ComponentName profileOwner = getProfileOwner(userId);
         return who != null && who.equals(profileOwner);
@@ -6103,7 +6112,7 @@
         Preconditions.checkNotNull(packageName, "packageName is null");
         final int callingUid = mInjector.binderGetCallingUid();
         try {
-            int uid = mContext.getPackageManager().getPackageUidAsUser(packageName,
+            int uid = mInjector.getPackageManager().getPackageUidAsUser(packageName,
                     UserHandle.getUserId(callingUid));
             if (uid != callingUid) {
                 throw new SecurityException("Invalid packageName");
@@ -6847,7 +6856,7 @@
             }
 
             try {
-                int uid = mContext.getPackageManager().getPackageUidAsUser(
+                int uid = mInjector.getPackageManager().getPackageUidAsUser(
                         policy.mApplicationRestrictionsManagingPackage, userHandle);
                 return uid == callingUid;
             } catch (NameNotFoundException e) {
@@ -8635,7 +8644,7 @@
             }
             final String deviceOwnerPackageName = mOwners.getDeviceOwnerComponent()
                     .getPackageName();
-            final String[] pkgs = mContext.getPackageManager().getPackagesForUid(callerUid);
+            final String[] pkgs = mInjector.getPackageManager().getPackagesForUid(callerUid);
 
             for (String pkg : pkgs) {
                 if (deviceOwnerPackageName.equals(pkg)) {
@@ -8672,7 +8681,7 @@
 
             ActivityInfo[] receivers = null;
             try {
-                receivers  = mContext.getPackageManager().getPackageInfo(
+                receivers  = mInjector.getPackageManager().getPackageInfo(
                         deviceOwnerPackage, PackageManager.GET_RECEIVERS).receivers;
             } catch (NameNotFoundException e) {
                 Log.e(LOG_TAG, "Cannot find device owner package", e);
@@ -8728,7 +8737,7 @@
                         < android.os.Build.VERSION_CODES.M) {
                     return false;
                 }
-                final PackageManager packageManager = mContext.getPackageManager();
+                final PackageManager packageManager = mInjector.getPackageManager();
                 switch (grantState) {
                     case DevicePolicyManager.PERMISSION_GRANT_STATE_GRANTED: {
                         mInjector.getPackageManagerInternal().grantRuntimePermission(packageName,
@@ -8763,7 +8772,7 @@
     @Override
     public int getPermissionGrantState(ComponentName admin, String packageName,
             String permission) throws RemoteException {
-        PackageManager packageManager = mContext.getPackageManager();
+        PackageManager packageManager = mInjector.getPackageManager();
 
         UserHandle user = mInjector.binderGetCallingUserHandle();
         synchronized (this) {
@@ -8800,17 +8809,33 @@
     }
 
     @Override
-    public boolean isProvisioningAllowed(String action) {
-        return checkProvisioningPreConditionSkipPermission(action) == CODE_OK;
+    public boolean isProvisioningAllowed(String action, String packageName) {
+        Preconditions.checkNotNull(packageName);
+
+        final int callingUid = mInjector.binderGetCallingUid();
+        final long ident = mInjector.binderClearCallingIdentity();
+        try {
+            final int uidForPackage = mInjector.getPackageManager().getPackageUidAsUser(
+                    packageName, UserHandle.getUserId(callingUid));
+            Preconditions.checkArgument(callingUid == uidForPackage,
+                    "Caller uid doesn't match the one for the provided package.");
+        } catch (NameNotFoundException e) {
+            throw new IllegalArgumentException("Invalid package provided " + packageName, e);
+        } finally {
+            mInjector.binderRestoreCallingIdentity(ident);
+        }
+
+        return checkProvisioningPreConditionSkipPermission(action, packageName) == CODE_OK;
     }
 
     @Override
-    public int checkProvisioningPreCondition(String action) {
+    public int checkProvisioningPreCondition(String action, String packageName) {
+        Preconditions.checkNotNull(packageName);
         enforceCanManageProfileAndDeviceOwners();
-        return checkProvisioningPreConditionSkipPermission(action);
+        return checkProvisioningPreConditionSkipPermission(action, packageName);
     }
 
-    private int checkProvisioningPreConditionSkipPermission(String action) {
+    private int checkProvisioningPreConditionSkipPermission(String action, String packageName) {
         if (!mHasFeature) {
             return CODE_DEVICE_ADMIN_NOT_SUPPORTED;
         }
@@ -8819,7 +8844,7 @@
         if (action != null) {
             switch (action) {
                 case DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE:
-                    return checkManagedProfileProvisioningPreCondition(callingUserId);
+                    return checkManagedProfileProvisioningPreCondition(packageName, callingUserId);
                 case DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE:
                     return checkDeviceOwnerProvisioningPreCondition(callingUserId);
                 case DevicePolicyManager.ACTION_PROVISION_MANAGED_USER:
@@ -8888,7 +8913,7 @@
         }
     }
 
-    private int checkManagedProfileProvisioningPreCondition(int callingUserId) {
+    private int checkManagedProfileProvisioningPreCondition(String packageName, int callingUserId) {
         if (!hasFeatureManagedUsers()) {
             return CODE_MANAGED_USERS_NOT_SUPPORTED;
         }
@@ -8901,24 +8926,25 @@
             // Managed user cannot have a managed profile.
             return CODE_USER_HAS_PROFILE_OWNER;
         }
+
         final long ident = mInjector.binderClearCallingIdentity();
         try {
-             /* STOPSHIP(b/31952368) Reinstate a check similar to this once ManagedProvisioning
-                   uses checkProvisioningPreCondition (see ag/1607846) and passes the packageName
-                   there. In isProvisioningAllowed we should check isCallerDeviceOwner, but for
-                   managed provisioning we need to check the package that is going to be set as PO
-                if (mUserManager.hasUserRestriction(UserManager.DISALLOW_ADD_MANAGED_PROFILE)) {
-                    if (!isCallerDeviceOwner(callingUid)
-                            || isAdminAffectedByRestriction(mOwners.getDeviceOwnerComponent(),
-                                    UserManager.DISALLOW_ADD_MANAGED_PROFILE, callingUserId)) {
+            final UserHandle callingUserHandle = UserHandle.of(callingUserId);
+            if (mUserManager.hasUserRestriction(
+                    UserManager.DISALLOW_ADD_MANAGED_PROFILE, callingUserHandle)) {
+                // The DO can initiate provisioning if the restriction was set by the DO.
+                if (!isDeviceOwnerPackage(packageName, callingUserId)
+                        || isAdminAffectedByRestriction(mOwners.getDeviceOwnerComponent(),
+                                UserManager.DISALLOW_ADD_MANAGED_PROFILE, callingUserId)) {
                     // Caller is not DO or the restriction was set by the system.
-                    return false;
-                    }
-                } */
+                    return CODE_ADD_MANAGED_PROFILE_DISALLOWED;
+                }
+            }
+
             // TODO: Allow it if the caller is the DO? DO could just call removeUser() before
             // provisioning, so not strictly required...
             boolean canRemoveProfile = !mUserManager.hasUserRestriction(
-                        UserManager.DISALLOW_REMOVE_MANAGED_PROFILE, UserHandle.of(callingUserId));
+                        UserManager.DISALLOW_REMOVE_MANAGED_PROFILE, callingUserHandle);
             if (!mUserManager.canAddMoreManagedProfiles(callingUserId, canRemoveProfile)) {
                 return CODE_CANNOT_ADD_MANAGED_PROFILE;
             }
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 48c9853..c35d114 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -77,7 +77,7 @@
 
 /**
  * Tests for DevicePolicyManager( and DevicePolicyManagerService).
- *
+ * You can run them via:
  m FrameworksServicesTests &&
  adb install \
    -r ${ANDROID_PRODUCT_OUT}/data/app/FrameworksServicesTests/FrameworksServicesTests.apk &&
@@ -85,6 +85,9 @@
    -w com.android.frameworks.servicestests/android.support.test.runner.AndroidJUnitRunner
 
  (mmma frameworks/base/services/tests/servicestests/ for non-ninja build)
+ *
+ * , or:
+ * runtest -c com.android.server.devicepolicy.DevicePolicyManagerTest frameworks-services
  */
 @SmallTest
 public class DevicePolicyManagerTest extends DpmTestBase {
@@ -2010,7 +2013,7 @@
         // UnfinishedVerificationException.
     }
 
-    public void setup_DeviceAdminFeatureOff() throws Exception {
+    private void setup_DeviceAdminFeatureOff() throws Exception {
         when(mContext.packageManager.hasSystemFeature(PackageManager.FEATURE_DEVICE_ADMIN))
                 .thenReturn(false);
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
@@ -2026,6 +2029,8 @@
 
     public void testIsProvisioningAllowed_DeviceAdminFeatureOff() throws Exception {
         setup_DeviceAdminFeatureOff();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE, false);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, false);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE,
@@ -2047,7 +2052,7 @@
                 DevicePolicyManager.CODE_DEVICE_ADMIN_NOT_SUPPORTED);
     }
 
-    public void setup_ManagedProfileFeatureOff() throws Exception {
+    private void setup_ManagedProfileFeatureOff() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(false);
         initializeDpms();
@@ -2061,6 +2066,8 @@
 
     public void testIsProvisioningAllowed_ManagedProfileFeatureOff() throws Exception {
         setup_ManagedProfileFeatureOff();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE, true);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, false);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE,
@@ -2102,7 +2109,7 @@
                 DevicePolicyManager.CODE_MANAGED_USERS_NOT_SUPPORTED);
     }
 
-    public void setup_nonSplitUser_firstBoot_primaryUser() throws Exception {
+    private void setup_nonSplitUser_firstBoot_primaryUser() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(false);
@@ -2115,6 +2122,8 @@
 
     public void testIsProvisioningAllowed_nonSplitUser_firstBoot_primaryUser() throws Exception {
         setup_nonSplitUser_firstBoot_primaryUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE, true);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE,
@@ -2138,7 +2147,7 @@
                 DevicePolicyManager.CODE_NOT_SYSTEM_USER_SPLIT);
     }
 
-    public void setup_nonSplitUser_afterDeviceSetup_primaryUser() throws Exception {
+    private void setup_nonSplitUser_afterDeviceSetup_primaryUser() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(false);
@@ -2152,6 +2161,8 @@
     public void testIsProvisioningAllowed_nonSplitUser_afterDeviceSetup_primaryUser()
             throws Exception {
         setup_nonSplitUser_afterDeviceSetup_primaryUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE,
                 false/* because of completed device setup */);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
@@ -2176,7 +2187,88 @@
                 DevicePolicyManager.CODE_NOT_SYSTEM_USER_SPLIT);
     }
 
-    public void setup_splitUser_firstBoot_systemUser() throws Exception {
+    public void testIsProvisioningAllowed_nonSplitUser_withDo_primaryUser() throws Exception {
+        setDeviceOwner();
+        setup_nonSplitUser_afterDeviceSetup_primaryUser();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
+        mContext.packageName = admin1.getPackageName();
+
+        // COMP mode is allowed.
+        assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
+
+        when(mContext.userManager.hasUserRestriction(
+                eq(UserManager.DISALLOW_ADD_MANAGED_PROFILE),
+                eq(UserHandle.getUserHandleForUid(mContext.binder.callingUid))))
+                .thenReturn(true);
+
+        // The DO should be allowed to initiate provisioning if it set the restriction itself.
+        when(mContext.userManager.getUserRestrictionSource(
+                eq(UserManager.DISALLOW_ADD_MANAGED_PROFILE),
+                eq(UserHandle.getUserHandleForUid(mContext.binder.callingUid))))
+                .thenReturn(UserManager.RESTRICTION_SOURCE_DEVICE_OWNER);
+        assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
+
+        // The DO should not be allowed to initiate provisioning if the restriction is set by
+        // another entity.
+        when(mContext.userManager.getUserRestrictionSource(
+                eq(UserManager.DISALLOW_ADD_MANAGED_PROFILE),
+                eq(UserHandle.getUserHandleForUid(mContext.binder.callingUid))))
+                .thenReturn(UserManager.RESTRICTION_SOURCE_SYSTEM);
+        assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, false);
+    }
+
+    public void
+    testCheckProvisioningPreCondition_nonSplitUser_withDo_primaryUser() throws Exception {
+        setDeviceOwner();
+        setup_nonSplitUser_afterDeviceSetup_primaryUser();
+        mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        assertCheckProvisioningPreCondition(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE,
+                DevicePolicyManager.CODE_HAS_DEVICE_OWNER);
+
+        // COMP mode is allowed.
+        assertCheckProvisioningPreCondition(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
+                DevicePolicyManager.CODE_OK);
+
+        // And other DPCs can also provisioning a managed profile (DO + BYOD case).
+        assertCheckProvisioningPreCondition(
+                DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
+                "some.other.dpc.package.name",
+                DevicePolicyManager.CODE_OK);
+
+        when(mContext.userManager.hasUserRestriction(
+                eq(UserManager.DISALLOW_ADD_MANAGED_PROFILE),
+                eq(UserHandle.getUserHandleForUid(mContext.binder.callingUid))))
+                .thenReturn(true);
+
+        // The DO should be allowed to initiate provisioning if it set the restriction itself, but
+        // other packages should be forbidden.
+        when(mContext.userManager.getUserRestrictionSource(
+                eq(UserManager.DISALLOW_ADD_MANAGED_PROFILE),
+                eq(UserHandle.getUserHandleForUid(mContext.binder.callingUid))))
+                .thenReturn(UserManager.RESTRICTION_SOURCE_DEVICE_OWNER);
+        assertCheckProvisioningPreCondition(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
+                DevicePolicyManager.CODE_OK);
+        assertCheckProvisioningPreCondition(
+                DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
+                "some.other.dpc.package.name",
+                DevicePolicyManager.CODE_ADD_MANAGED_PROFILE_DISALLOWED);
+
+        // The DO should not be allowed to initiate provisioning if the restriction is set by
+        // another entity.
+        when(mContext.userManager.getUserRestrictionSource(
+                eq(UserManager.DISALLOW_ADD_MANAGED_PROFILE),
+                eq(UserHandle.getUserHandleForUid(mContext.binder.callingUid))))
+                .thenReturn(UserManager.RESTRICTION_SOURCE_SYSTEM);
+        assertCheckProvisioningPreCondition(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
+                DevicePolicyManager.CODE_ADD_MANAGED_PROFILE_DISALLOWED);
+                assertCheckProvisioningPreCondition(
+                DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
+                "some.other.dpc.package.name",
+                DevicePolicyManager.CODE_ADD_MANAGED_PROFILE_DISALLOWED);
+    }
+
+    private void setup_splitUser_firstBoot_systemUser() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(true);
@@ -2189,6 +2281,8 @@
 
     public void testIsProvisioningAllowed_splitUser_firstBoot_systemUser() throws Exception {
         setup_splitUser_firstBoot_systemUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE, true);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
                 false /* because canAddMoreManagedProfiles returns false */);
@@ -2213,7 +2307,7 @@
                 DevicePolicyManager.CODE_SYSTEM_USER);
     }
 
-    public void setup_splitUser_afterDeviceSetup_systemUser() throws Exception {
+    private void setup_splitUser_afterDeviceSetup_systemUser() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(true);
@@ -2226,6 +2320,8 @@
 
     public void testIsProvisioningAllowed_splitUser_afterDeviceSetup_systemUser() throws Exception {
         setup_splitUser_afterDeviceSetup_systemUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE,
                 true/* it's undefined behavior. Can be changed into false in the future */);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
@@ -2251,7 +2347,7 @@
                 DevicePolicyManager.CODE_SYSTEM_USER);
     }
 
-    public void setup_splitUser_firstBoot_primaryUser() throws Exception {
+    private void setup_splitUser_firstBoot_primaryUser() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(true);
@@ -2264,6 +2360,8 @@
 
     public void testIsProvisioningAllowed_splitUser_firstBoot_primaryUser() throws Exception {
         setup_splitUser_firstBoot_primaryUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE, true);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_SHAREABLE_DEVICE,
@@ -2286,7 +2384,7 @@
                 DevicePolicyManager.CODE_OK);
     }
 
-    public void setup_splitUser_afterDeviceSetup_primaryUser() throws Exception {
+    private void setup_splitUser_afterDeviceSetup_primaryUser() throws Exception {
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(true);
@@ -2300,6 +2398,8 @@
     public void testIsProvisioningAllowed_splitUser_afterDeviceSetup_primaryUser()
             throws Exception {
         setup_splitUser_afterDeviceSetup_primaryUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_DEVICE,
                 true/* it's undefined behavior. Can be changed into false in the future */);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
@@ -2324,7 +2424,7 @@
                 DevicePolicyManager.CODE_USER_SETUP_COMPLETED);
     }
 
-    public void setup_provisionManagedProfileWithDeviceOwner_systemUser() throws Exception {
+    private void setup_provisionManagedProfileWithDeviceOwner_systemUser() throws Exception {
         setDeviceOwner();
 
         when(mContext.ipackageManager.hasSystemFeature(PackageManager.FEATURE_MANAGED_USERS, 0))
@@ -2340,6 +2440,8 @@
     public void testIsProvisioningAllowed_provisionManagedProfileWithDeviceOwner_systemUser()
             throws Exception {
         setup_provisionManagedProfileWithDeviceOwner_systemUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
                 false /* can't provision managed profile on system user */);
     }
@@ -2368,6 +2470,8 @@
     public void testIsProvisioningAllowed_provisionManagedProfileWithDeviceOwner_primaryUser()
             throws Exception {
         setup_provisionManagedProfileWithDeviceOwner_primaryUser();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
+        mContext.packageName = admin1.getPackageName();
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, true);
     }
 
@@ -2375,6 +2479,8 @@
             throws Exception {
         setup_provisionManagedProfileWithDeviceOwner_primaryUser();
         mContext.callerPermissions.add(permission.MANAGE_PROFILE_AND_DEVICE_OWNERS);
+
+        // COMP mode is allowed.
         assertCheckProvisioningPreCondition(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE,
                 DevicePolicyManager.CODE_OK);
     }
@@ -2386,8 +2492,8 @@
                 .thenReturn(true);
         when(mContext.userManagerForMock.isSplitSystemUser()).thenReturn(true);
         when(mContext.userManager.hasUserRestriction(
-                UserManager.DISALLOW_REMOVE_MANAGED_PROFILE,
-                UserHandle.of(DpmMockContext.CALLER_USER_HANDLE)))
+                eq(UserManager.DISALLOW_REMOVE_MANAGED_PROFILE),
+                eq(UserHandle.of(DpmMockContext.CALLER_USER_HANDLE))))
                 .thenReturn(true);
         when(mContext.userManager.canAddMoreManagedProfiles(DpmMockContext.CALLER_USER_HANDLE,
                 false /* we can't remove a managed profile */)).thenReturn(false);
@@ -2401,6 +2507,8 @@
     public void testIsProvisioningAllowed_provisionManagedProfileCantRemoveUser_primaryUser()
             throws Exception {
         setup_provisionManagedProfileCantRemoveUser_primaryUser();
+        mContext.packageName = admin1.getPackageName();
+        setUpPackageManagerForAdmin(admin1, mContext.binder.callingUid);
         assertProvisioningAllowed(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, false);
     }
 
@@ -2415,7 +2523,8 @@
     public void testCheckProvisioningPreCondition_permission() {
         // GIVEN the permission MANAGE_PROFILE_AND_DEVICE_OWNERS is not granted
         try {
-            dpm.checkProvisioningPreCondition(DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE);
+            dpm.checkProvisioningPreCondition(
+                    DevicePolicyManager.ACTION_PROVISION_MANAGED_PROFILE, "some.package");
             fail("Didn't throw SecurityException");
         } catch (SecurityException expected) {
         }
@@ -2820,8 +2929,14 @@
     }
 
     private void assertCheckProvisioningPreCondition(String action, int provisioningCondition) {
-        assertEquals("checkProvisioningPreCondition(" + action + ") returning unexpected result",
-                provisioningCondition, dpm.checkProvisioningPreCondition(action));
+        assertCheckProvisioningPreCondition(action, admin1.getPackageName(), provisioningCondition);
+    }
+
+    private void assertCheckProvisioningPreCondition(
+            String action, String packageName, int provisioningCondition) {
+        assertEquals("checkProvisioningPreCondition("
+                        + action + ", " + packageName + ") returning unexpected result",
+                provisioningCondition, dpm.checkProvisioningPreCondition(action, packageName));
     }
 
     /**
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DpmTestBase.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DpmTestBase.java
index db27f72..8a11976 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DpmTestBase.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DpmTestBase.java
@@ -82,6 +82,10 @@
                 eq(packageName),
                 eq(0),
                 eq(userId));
+
+        doReturn(ai.uid).when(mMockContext.packageManager).getPackageUidAsUser(
+                eq(packageName),
+                eq(userId));
     }
 
     protected void setUpPackageManagerForAdmin(ComponentName admin, int packageUid)