Fix NPE during migration.

User restrictions are null for PO parent admin, so we need to ensure
they are initialized before adding.

Bug: 147884539
Test: atest com.android.server.devicepolicy.DevicePolicyManagerServiceMigrationTest#testCompMigrationAffiliated
Change-Id: I5024ac585e30081f944717cfe8d4b920b644213e

3 files changed, 12 insertions, 5 deletions

diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 14bd72b0fd21..af57c29fda82 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -2556,7 +2556,7 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
         }
         for (final String restriction : doAdmin.userRestrictions.keySet()) {
             if (UserRestrictionsUtils.canProfileOwnerOfOrganizationOwnedDeviceChange(restriction)) {
-                parentAdmin.userRestrictions.putBoolean(
+                parentAdmin.ensureUserRestrictions().putBoolean(
                         restriction, doAdmin.userRestrictions.getBoolean(restriction));
             }
         }
diff --git a/services/tests/servicestests/res/raw/comp_policies_primary.xml b/services/tests/servicestests/res/raw/comp_policies_primary.xml
index 1e1a0eff874c..d30f479195e3 100644
--- a/services/tests/servicestests/res/raw/comp_policies_primary.xml
+++ b/services/tests/servicestests/res/raw/comp_policies_primary.xml
@@ -3,5 +3,6 @@
     <admin name="com.android.frameworks.servicestests/com.android.server.devicepolicy.DummyDeviceAdmins$Admin1">
         <policies flags="991"/>
         <password-history-length value="33" />
+        <user-restrictions no_bluetooth="true" />
     </admin>
 </policies>
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java
index 46b83713c159..9574a086c74b 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java
@@ -371,10 +371,16 @@ public class DevicePolicyManagerServiceMigrationTest extends DpmTestBase {
         poContext.binder.callingUid = UserHandle.getUid(COPE_PROFILE_USER_ID, COPE_ADMIN1_APP_ID);
 
         runAsCaller(poContext, dpms, dpm -> {
-            // Check that DO policy is now set on parent instance.
-            assertEquals(33, dpm.getParentProfileInstance(admin1).getPasswordHistoryLength(admin1));
-            // And NOT set on profile instance.
-            assertEquals(0, dpm.getPasswordHistoryLength(admin1));
+            assertEquals("Password history policy wasn't migrated to PO parent instance",
+                    33, dpm.getParentProfileInstance(admin1).getPasswordHistoryLength(admin1));
+            assertEquals("Password history policy was put into non-parent PO instance",
+                    0, dpm.getPasswordHistoryLength(admin1));
+
+            assertTrue("User restriction wasn't migrated to PO parent instance",
+                    dpm.getParentProfileInstance(admin1).getUserRestrictions(admin1)
+                            .containsKey(UserManager.DISALLOW_BLUETOOTH));
+            assertFalse("User restriction was put into non-parent PO instance",
+                    dpm.getUserRestrictions(admin1).containsKey(UserManager.DISALLOW_BLUETOOTH));
 
             // TODO(b/143516163): verify more policies.
         });