Merge "Change the data structure of restrictions in UM" into rvc-dev

services/core/java/com/android/server/pm/RestrictionsSet.java

@@ -0,0 +1,256 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.pm;
+
+import android.annotation.NonNull;
+import android.annotation.Nullable;
+import android.annotation.UserIdInt;
+import android.os.Bundle;
+import android.os.UserManager;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import org.xmlpull.v1.XmlPullParser;
+import org.xmlpull.v1.XmlPullParserException;
+import org.xmlpull.v1.XmlSerializer;
+
+import java.io.IOException;
+import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Data structure that contains the mapping of users to user restrictions (either the user
+ * restrictions that apply to them, or the user restrictions that they set, depending on the
+ * circumstances).
+ *
+ * @hide
+ */
+public class RestrictionsSet {
+
+    private static final String USER_ID = "user_id";
+    private static final String TAG_RESTRICTIONS = "restrictions";
+    private static final String TAG_RESTRICTIONS_USER = "restrictions_user";
+
+    /**
+     * Mapping of user restrictions.
+     * Only non-empty restriction bundles are stored.
+     * The key is the user id of the user.
+     * userId -> restrictionBundle
+     */
+    private final SparseArray<Bundle> mUserRestrictions = new SparseArray<>(0);
+
+    public RestrictionsSet() {
+    }
+
+    public RestrictionsSet(@UserIdInt int userId, @NonNull Bundle restrictions) {
+        if (restrictions.isEmpty()) {
+            throw new IllegalArgumentException("empty restriction bundle cannot be added.");
+        }
+        mUserRestrictions.put(userId, restrictions);
+    }
+
+    /**
+     * Updates restriction bundle for a given user.
+     * If new bundle is empty, record is removed from the array.
+     *
+     * @return whether restrictions bundle is different from the old one.
+     */
+    public boolean updateRestrictions(@UserIdInt int userId, @Nullable Bundle restrictions) {
+        final boolean changed =
+                !UserRestrictionsUtils.areEqual(mUserRestrictions.get(userId), restrictions);
+        if (!changed) {
+            return false;
+        }
+        if (!UserRestrictionsUtils.isEmpty(restrictions)) {
+            mUserRestrictions.put(userId, restrictions);
+        } else {
+            mUserRestrictions.delete(userId);
+        }
+        return true;
+    }
+
+    /**
+     * Moves a particular restriction from one restriction set to another, e.g. for all users.
+     */
+    public void moveRestriction(@NonNull RestrictionsSet destRestrictions, String restriction) {
+        for (int i = 0; i < mUserRestrictions.size(); i++) {
+            final int userId = mUserRestrictions.keyAt(i);
+            final Bundle from = mUserRestrictions.valueAt(i);
+
+            if (UserRestrictionsUtils.contains(from, restriction)) {
+                from.remove(restriction);
+                Bundle to = destRestrictions.getRestrictions(userId);
+                if (to == null) {
+                    to = new Bundle();
+                    to.putBoolean(restriction, true);
+                    destRestrictions.updateRestrictions(userId, to);
+                } else {
+                    to.putBoolean(restriction, true);
+                }
+                // Don't keep empty bundles.
+                if (from.isEmpty()) {
+                    mUserRestrictions.removeAt(i);
+                    i--;
+                }
+            }
+        }
+    }
+
+    /**
+     * @return whether restrictions set has no restrictions.
+     */
+    public boolean isEmpty() {
+        return mUserRestrictions.size() == 0;
+    }
+
+    /**
+     * Merge all restrictions in restrictions set into one bundle. The original user restrictions
+     * set does not get modified, instead a new bundle is returned.
+     *
+     * @return restrictions bundle containing all user restrictions.
+     */
+    public @NonNull Bundle mergeAll() {
+        final Bundle result = new Bundle();
+        for (int i = 0; i < mUserRestrictions.size(); i++) {
+            UserRestrictionsUtils.merge(result, mUserRestrictions.valueAt(i));
+        }
+        return result;
+    }
+
+    /**
+     * @return list of enforcing users that enforce a particular restriction.
+     */
+    public @NonNull List<UserManager.EnforcingUser> getEnforcingUsers(String restriction,
+            @UserIdInt int deviceOwnerUserId) {
+        final List<UserManager.EnforcingUser> result = new ArrayList<>();
+        for (int i = 0; i < mUserRestrictions.size(); i++) {
+            if (UserRestrictionsUtils.contains(mUserRestrictions.valueAt(i), restriction)) {
+                result.add(getEnforcingUser(mUserRestrictions.keyAt(i), deviceOwnerUserId));
+            }
+        }
+        return result;
+    }
+
+    private UserManager.EnforcingUser getEnforcingUser(@UserIdInt int userId,
+            @UserIdInt int deviceOwnerUserId) {
+        int source = deviceOwnerUserId == userId
+                ? UserManager.RESTRICTION_SOURCE_DEVICE_OWNER
+                : UserManager.RESTRICTION_SOURCE_PROFILE_OWNER;
+        return new UserManager.EnforcingUser(userId, source);
+    }
+
+    /**
+     * @return list of user restrictions for a given user. Null is returned if the user does not
+     * have any restrictions.
+     */
+    public @Nullable Bundle getRestrictions(@UserIdInt int userId) {
+        return mUserRestrictions.get(userId);
+    }
+
+    /**
+     * Removes a given user from the restrictions set, returning true if the user has non-empty
+     * restrictions before removal.
+     */
+    public boolean remove(@UserIdInt int userId) {
+        boolean hasUserRestriction = mUserRestrictions.contains(userId);
+        mUserRestrictions.remove(userId);
+        return hasUserRestriction;
+    }
+
+    /**
+     * Remove list of users and user restrictions.
+     */
+    public void removeAllRestrictions() {
+        mUserRestrictions.clear();
+    }
+
+    /**
+     * Serialize a given {@link RestrictionsSet} to XML.
+     */
+    public void writeRestrictions(@NonNull XmlSerializer serializer, @NonNull String outerTag)
+            throws IOException {
+        serializer.startTag(null, outerTag);
+        for (int i = 0; i < mUserRestrictions.size(); i++) {
+            serializer.startTag(null, TAG_RESTRICTIONS_USER);
+            serializer.attribute(null, USER_ID, String.valueOf(mUserRestrictions.keyAt(i)));
+            UserRestrictionsUtils.writeRestrictions(serializer, mUserRestrictions.valueAt(i),
+                    TAG_RESTRICTIONS);
+            serializer.endTag(null, TAG_RESTRICTIONS_USER);
+        }
+        serializer.endTag(null, outerTag);
+    }
+
+    /**
+     * Read restrictions from XML.
+     */
+    public static RestrictionsSet readRestrictions(@NonNull XmlPullParser parser,
+            @NonNull String outerTag) throws IOException, XmlPullParserException {
+        RestrictionsSet restrictionsSet = new RestrictionsSet();
+        int userId = 0;
+        int type;
+
+        while ((type = parser.next()) != XmlPullParser.END_DOCUMENT) {
+            String tag = parser.getName();
+            if (type == XmlPullParser.END_TAG && outerTag.equals(tag)) {
+                return restrictionsSet;
+            } else if (type == XmlPullParser.START_TAG && TAG_RESTRICTIONS_USER.equals(tag)) {
+                userId = Integer.parseInt(parser.getAttributeValue(null, USER_ID));
+            } else if (type == XmlPullParser.START_TAG && TAG_RESTRICTIONS.equals(tag)) {
+                Bundle restrictions = UserRestrictionsUtils.readRestrictions(parser);
+                restrictionsSet.updateRestrictions(userId, restrictions);
+            }
+        }
+        throw new XmlPullParserException("restrictions cannot be read as xml is malformed.");
+    }
+
+    /**
+     * Dumps {@link RestrictionsSet}.
+     */
+    public void dumpRestrictions(PrintWriter pw, String prefix) {
+        boolean noneSet = true;
+        for (int i = 0; i < mUserRestrictions.size(); i++) {
+            pw.println(prefix + "User Id: " + mUserRestrictions.keyAt(i));
+            UserRestrictionsUtils.dumpRestrictions(pw, prefix + "  ", mUserRestrictions.valueAt(i));
+            noneSet = false;
+        }
+        if (noneSet) {
+            pw.println(prefix + "none");
+        }
+    }
+
+    public boolean containsKey(@UserIdInt int userId) {
+        return mUserRestrictions.contains(userId);
+    }
+
+    @VisibleForTesting
+    public int size() {
+        return mUserRestrictions.size();
+    }
+
+    @VisibleForTesting
+    public int keyAt(int index) {
+        return mUserRestrictions.keyAt(index);
+    }
+
+    @VisibleForTesting
+    public Bundle valueAt(int index) {
+        return mUserRestrictions.valueAt(index);
+    }
+
+}

services/core/java/com/android/server/pm/UserManagerService.java

@@ -182,6 +182,8 @@ public class UserManagerService extends IUserManager.Stub {
     private static final String TAG_USER = "user";
     private static final String TAG_RESTRICTIONS = "restrictions";
     private static final String TAG_DEVICE_POLICY_RESTRICTIONS = "device_policy_restrictions";
+    private static final String TAG_DEVICE_POLICY_LOCAL_RESTRICTIONS =
+            "device_policy_local_restrictions";
     private static final String TAG_DEVICE_POLICY_GLOBAL_RESTRICTIONS =
             "device_policy_global_restrictions";
     /** Legacy name for device owner id tag. */
@@ -329,7 +331,7 @@ public class UserManagerService extends IUserManager.Stub {
      * {@link #updateUserRestrictionsInternalLR}.
      */
     @GuardedBy("mRestrictionsLock")
-    private final SparseArray<Bundle> mBaseUserRestrictions = new SparseArray<>();
+    private final RestrictionsSet mBaseUserRestrictions = new RestrictionsSet();
 
     /**
      * Cached user restrictions that are in effect -- i.e. {@link #mBaseUserRestrictions} combined
@@ -344,7 +346,7 @@ public class UserManagerService extends IUserManager.Stub {
      * {@link #updateUserRestrictionsInternalLR}.
      */
     @GuardedBy("mRestrictionsLock")
-    private final SparseArray<Bundle> mCachedEffectiveUserRestrictions = new SparseArray<>();
+    private final RestrictionsSet mCachedEffectiveUserRestrictions = new RestrictionsSet();
 
     /**
      * User restrictions that have already been applied in
@@ -353,7 +355,7 @@ public class UserManagerService extends IUserManager.Stub {
      * {@link #updateUserRestrictionsInternalLR(Bundle, int)} call.
      */
     @GuardedBy("mRestrictionsLock")
-    private final SparseArray<Bundle> mAppliedUserRestrictions = new SparseArray<>();
+    private final RestrictionsSet mAppliedUserRestrictions = new RestrictionsSet();
 
     /**
      * User restrictions set by {@link com.android.server.devicepolicy.DevicePolicyManagerService}
@@ -362,7 +364,7 @@ public class UserManagerService extends IUserManager.Stub {
      * The key is the user id of the user whom the restriction originated from.
      */
     @GuardedBy("mRestrictionsLock")
-    private final SparseArray<Bundle> mDevicePolicyGlobalUserRestrictions = new SparseArray<>();
+    private final RestrictionsSet mDevicePolicyGlobalUserRestrictions = new RestrictionsSet();
 
     /**
      * Id of the user that set global restrictions.
@@ -372,11 +374,15 @@ public class UserManagerService extends IUserManager.Stub {
 
     /**
      * User restrictions set by {@link com.android.server.devicepolicy.DevicePolicyManagerService}
-     * for each user. Only non-empty restriction bundles are stored.
+     * for each user.
-     * The key is the user id of the user whom the restriction originated from.
+     * The key is the user id of the user whom the restrictions are targeting.
+     * The key inside the restrictionsSet is the user id of the user whom the restriction
+     * originated from.
+     * targetUserId -> originatingUserId -> restrictionBundle
      */
     @GuardedBy("mRestrictionsLock")
-    private final SparseArray<Bundle> mDevicePolicyLocalUserRestrictions = new SparseArray<>();
+    private final SparseArray<RestrictionsSet> mDevicePolicyLocalUserRestrictions =
+            new SparseArray<>();
 
     @GuardedBy("mGuestRestrictions")
     private final Bundle mGuestRestrictions = new Bundle();
@@ -1710,9 +1716,6 @@ public class UserManagerService extends IUserManager.Stub {
         }
     }
 
-    /**
-     * See {@link UserManagerInternal#setDevicePolicyUserRestrictions}
-     */
     private void setDevicePolicyUserRestrictionsInner(@UserIdInt int originatingUserId,
             @Nullable Bundle restrictions,
             @UserManagerInternal.OwnerType int restrictionOwnerType) {
@@ -1722,16 +1725,36 @@ public class UserManagerService extends IUserManager.Stub {
         // Sort restrictions into local and global ensuring they don't overlap.
         UserRestrictionsUtils.sortToGlobalAndLocal(restrictions, restrictionOwnerType, global,
                 local);
+        boolean isDeviceOwner = restrictionOwnerType == UserManagerInternal.OWNER_TYPE_DEVICE_OWNER;
 
+        RestrictionsSet localRestrictionsSet;
+        if (UserRestrictionsUtils.isEmpty(local)) {
+            localRestrictionsSet = new RestrictionsSet();
+        } else {
+            localRestrictionsSet = new RestrictionsSet(originatingUserId, local);
+        }
+        setDevicePolicyUserRestrictionsInner(originatingUserId, global, localRestrictionsSet,
+                isDeviceOwner);
+    }
+
+    /**
+     * See {@link UserManagerInternal#setDevicePolicyUserRestrictions}
+     */
+    private void setDevicePolicyUserRestrictionsInner(@UserIdInt int originatingUserId,
+            @NonNull Bundle global, @NonNull RestrictionsSet local,
+            boolean isDeviceOwner) {
         boolean globalChanged, localChanged;
+        List<Integer> updatedLocalTargetUserIds;
         synchronized (mRestrictionsLock) {
             // Update global and local restrictions if they were changed.
-            globalChanged = updateRestrictionsIfNeededLR(
+            globalChanged = mDevicePolicyGlobalUserRestrictions
-                    originatingUserId, global, mDevicePolicyGlobalUserRestrictions);
+                    .updateRestrictions(originatingUserId, global);
-            localChanged = updateRestrictionsIfNeededLR(
+            updatedLocalTargetUserIds = getUpdatedTargetUserIdsFromLocalRestrictions(
-                    originatingUserId, local, mDevicePolicyLocalUserRestrictions);
+                    originatingUserId, local);
+            localChanged = updateLocalRestrictionsForTargetUsersLR(originatingUserId, local,
+                    updatedLocalTargetUserIds);
 
-            if (restrictionOwnerType == UserManagerInternal.OWNER_TYPE_DEVICE_OWNER) {
+            if (isDeviceOwner) {
                 // Remember the global restriction owner userId to be able to make a distinction
                 // in getUserRestrictionSource on who set local policies.
                 mDeviceOwnerUserId = originatingUserId;
@@ -1754,8 +1777,20 @@ public class UserManagerService extends IUserManager.Stub {
         }
         // Don't call them within the mRestrictionsLock.
         synchronized (mPackagesLock) {
-            if (localChanged || globalChanged) {
+            if (globalChanged || localChanged) {
-                writeUserLP(getUserDataNoChecks(originatingUserId));
+                if (updatedLocalTargetUserIds.size() == 1
+                        && updatedLocalTargetUserIds.contains(originatingUserId)) {
+                    writeUserLP(getUserDataNoChecks(originatingUserId));
+                } else {
+                    if (globalChanged) {
+                        writeUserLP(getUserDataNoChecks(originatingUserId));
+                    }
+                    if (localChanged) {
+                        for (int targetUserId : updatedLocalTargetUserIds) {
+                            writeAllTargetUsersLP(targetUserId);
+                        }
+                    }
+                }
             }
         }
 
@@ -1763,44 +1798,88 @@ public class UserManagerService extends IUserManager.Stub {
             if (globalChanged) {
                 applyUserRestrictionsForAllUsersLR();
             } else if (localChanged) {
-                applyUserRestrictionsLR(originatingUserId);
+                for (int targetUserId : updatedLocalTargetUserIds) {
+                    applyUserRestrictionsLR(targetUserId);
+                }
             }
         }
     }
 
     /**
-     * Updates restriction bundle for a given user in a given restriction array. If new bundle is
+     * @return the list of updated target user ids in device policy local restrictions for a
-     * empty, record is removed from the array.
+     * given originating user id.
-     * @return whether restrictions bundle is different from the old one.
      */
-    private boolean updateRestrictionsIfNeededLR(@UserIdInt int userId,
+    private List<Integer> getUpdatedTargetUserIdsFromLocalRestrictions(int originatingUserId,
-            @Nullable Bundle restrictions, SparseArray<Bundle> restrictionsArray) {
+            @NonNull RestrictionsSet local) {
-        final boolean changed =
+        List<Integer> targetUserIds = new ArrayList<>();
-                !UserRestrictionsUtils.areEqual(restrictionsArray.get(userId), restrictions);
+        // Update all the target user ids from the local restrictions set
-        if (changed) {
+        for (int i = 0; i < local.size(); i++) {
-            if (!UserRestrictionsUtils.isEmpty(restrictions)) {
+            targetUserIds.add(local.keyAt(i));
-                restrictionsArray.put(userId, restrictions);
+        }
-            } else {
+        // Update the target user id from device policy local restrictions if the local
-                restrictionsArray.delete(userId);
+        // restrictions set does not contain the target user id.
+        for (int i = 0; i < mDevicePolicyLocalUserRestrictions.size(); i++) {
+            int targetUserId = mDevicePolicyLocalUserRestrictions.keyAt(i);
+            RestrictionsSet restrictionsSet = mDevicePolicyLocalUserRestrictions.valueAt(i);
+            if (!local.containsKey(targetUserId)
+                    && restrictionsSet.containsKey(originatingUserId)) {
+                targetUserIds.add(targetUserId);
+            }
+        }
+        return targetUserIds;
+    }
+
+    /**
+     * Update restrictions for all target users in the restriction set. If a target user does not
+     * exist in device policy local restrictions, remove the restrictions bundle for that target
+     * user originating from the specified originating user.
+     */
+    private boolean updateLocalRestrictionsForTargetUsersLR(int originatingUserId,
+            RestrictionsSet local, List<Integer> updatedTargetUserIds) {
+        boolean changed = false;
+        for (int targetUserId : updatedTargetUserIds) {
+            Bundle restrictions = local.getRestrictions(targetUserId);
+            if (restrictions == null) {
+                restrictions = new Bundle();
+            }
+            if (getDevicePolicyLocalRestrictionsForTargetUserLR(targetUserId)
+                    .updateRestrictions(originatingUserId, restrictions)) {
+                changed = true;
             }
         }
         return changed;
     }
 
+    /**
+     * A new restriction set is created if a restriction set does not already exist for a given
+     * target user.
+     *
+     * @return restrictions set for a given target user.
+     */
+    private @NonNull RestrictionsSet getDevicePolicyLocalRestrictionsForTargetUserLR(
+            int targetUserId) {
+        RestrictionsSet result = mDevicePolicyLocalUserRestrictions.get(targetUserId);
+        if (result == null) {
+            result = new RestrictionsSet();
+            mDevicePolicyLocalUserRestrictions.put(targetUserId, result);
+        }
+        return result;
+    }
+
     @GuardedBy("mRestrictionsLock")
     private Bundle computeEffectiveUserRestrictionsLR(@UserIdInt int userId) {
         final Bundle baseRestrictions =
-                UserRestrictionsUtils.nonNull(mBaseUserRestrictions.get(userId));
+                UserRestrictionsUtils.nonNull(mBaseUserRestrictions.getRestrictions(userId));
-        final Bundle global = UserRestrictionsUtils.mergeAll(mDevicePolicyGlobalUserRestrictions);
+        final Bundle global = mDevicePolicyGlobalUserRestrictions.mergeAll();
-        final Bundle local = mDevicePolicyLocalUserRestrictions.get(userId);
+        final RestrictionsSet local = getDevicePolicyLocalRestrictionsForTargetUserLR(userId);
 
-        if (UserRestrictionsUtils.isEmpty(global) && UserRestrictionsUtils.isEmpty(local)) {
+        if (UserRestrictionsUtils.isEmpty(global) && local.isEmpty()) {
             // Common case first.
             return baseRestrictions;
         }
         final Bundle effective = UserRestrictionsUtils.clone(baseRestrictions);
         UserRestrictionsUtils.merge(effective, global);
-        UserRestrictionsUtils.merge(effective, local);
+        UserRestrictionsUtils.merge(effective, local.mergeAll());
 
         return effective;
     }
@@ -1815,10 +1894,10 @@ public class UserManagerService extends IUserManager.Stub {
 
     private Bundle getEffectiveUserRestrictions(@UserIdInt int userId) {
         synchronized (mRestrictionsLock) {
-            Bundle restrictions = mCachedEffectiveUserRestrictions.get(userId);
+            Bundle restrictions = mCachedEffectiveUserRestrictions.getRestrictions(userId);
             if (restrictions == null) {
                 restrictions = computeEffectiveUserRestrictionsLR(userId);
-                mCachedEffectiveUserRestrictions.put(userId, restrictions);
+                mCachedEffectiveUserRestrictions.updateRestrictions(userId, restrictions);
             }
             return restrictions;
         }
@@ -1921,31 +2000,17 @@ public class UserManagerService extends IUserManager.Stub {
         }
 
         synchronized (mRestrictionsLock) {
-            // Check if it is set by profile owner.
+            // Check if it is set as a local restriction.
-            Bundle profileOwnerRestrictions = mDevicePolicyLocalUserRestrictions.get(userId);
+            result.addAll(getDevicePolicyLocalRestrictionsForTargetUserLR(userId).getEnforcingUsers(
-            if (UserRestrictionsUtils.contains(profileOwnerRestrictions, restrictionKey)) {
+                    restrictionKey, mDeviceOwnerUserId));
-                result.add(getEnforcingUserLocked(userId));
-            }
 
-            // Iterate over all users who enforce global restrictions.
+            // Check if it is set as a global restriction.
-            for (int i = mDevicePolicyGlobalUserRestrictions.size() - 1; i >= 0; i--) {
+            result.addAll(mDevicePolicyGlobalUserRestrictions.getEnforcingUsers(restrictionKey,
-                Bundle globalRestrictions = mDevicePolicyGlobalUserRestrictions.valueAt(i);
+                    mDeviceOwnerUserId));
-                int profileUserId = mDevicePolicyGlobalUserRestrictions.keyAt(i);
-                if (UserRestrictionsUtils.contains(globalRestrictions, restrictionKey)) {
-                    result.add(getEnforcingUserLocked(profileUserId));
-                }
-            }
         }
         return result;
     }
 
-    @GuardedBy("mRestrictionsLock")
-    private EnforcingUser getEnforcingUserLocked(@UserIdInt int userId) {
-        int source = mDeviceOwnerUserId == userId ? UserManager.RESTRICTION_SOURCE_DEVICE_OWNER
-                : UserManager.RESTRICTION_SOURCE_PROFILE_OWNER;
-        return new EnforcingUser(userId, source);
-    }
-
     /**
      * @return UserRestrictions that are in effect currently.  This always returns a new
      * {@link Bundle}.
@@ -1963,7 +2028,7 @@ public class UserManagerService extends IUserManager.Stub {
             return false;
         }
         synchronized (mRestrictionsLock) {
-            Bundle bundle = mBaseUserRestrictions.get(userId);
+            Bundle bundle = mBaseUserRestrictions.getRestrictions(userId);
             return (bundle != null && bundle.getBoolean(restrictionKey, false));
         }
     }
@@ -1978,7 +2043,7 @@ public class UserManagerService extends IUserManager.Stub {
             // Note we can't modify Bundles stored in mBaseUserRestrictions directly, so create
             // a copy.
             final Bundle newRestrictions = UserRestrictionsUtils.clone(
-                    mBaseUserRestrictions.get(userId));
+                    mBaseUserRestrictions.getRestrictions(userId));
             newRestrictions.putBoolean(key, value);
 
             updateUserRestrictionsInternalLR(newRestrictions, userId);
@@ -1997,25 +2062,25 @@ public class UserManagerService extends IUserManager.Stub {
     private void updateUserRestrictionsInternalLR(
             @Nullable Bundle newBaseRestrictions, @UserIdInt int userId) {
         final Bundle prevAppliedRestrictions = UserRestrictionsUtils.nonNull(
-                mAppliedUserRestrictions.get(userId));
+                mAppliedUserRestrictions.getRestrictions(userId));
 
         // Update base restrictions.
         if (newBaseRestrictions != null) {
             // If newBaseRestrictions == the current one, it's probably a bug.
-            final Bundle prevBaseRestrictions = mBaseUserRestrictions.get(userId);
+            final Bundle prevBaseRestrictions = mBaseUserRestrictions.getRestrictions(userId);
 
             Preconditions.checkState(prevBaseRestrictions != newBaseRestrictions);
-            Preconditions.checkState(mCachedEffectiveUserRestrictions.get(userId)
+            Preconditions.checkState(mCachedEffectiveUserRestrictions.getRestrictions(userId)
                     != newBaseRestrictions);
 
-            if (updateRestrictionsIfNeededLR(userId, newBaseRestrictions, mBaseUserRestrictions)) {
+            if (mBaseUserRestrictions.updateRestrictions(userId, newBaseRestrictions)) {
                 scheduleWriteUser(getUserDataNoChecks(userId));
             }
         }
 
         final Bundle effective = computeEffectiveUserRestrictionsLR(userId);
 
-        mCachedEffectiveUserRestrictions.put(userId, effective);
+        mCachedEffectiveUserRestrictions.updateRestrictions(userId, effective);
 
         // Apply the new restrictions.
         if (DBG) {
@@ -2038,7 +2103,7 @@ public class UserManagerService extends IUserManager.Stub {
 
         propagateUserRestrictionsLR(userId, effective, prevAppliedRestrictions);
 
-        mAppliedUserRestrictions.put(userId, new Bundle(effective));
+        mAppliedUserRestrictions.updateRestrictions(userId, new Bundle(effective));
     }
 
     private void propagateUserRestrictionsLR(final int userId,
@@ -2090,7 +2155,7 @@ public class UserManagerService extends IUserManager.Stub {
             debug("applyUserRestrictionsForAllUsersLR");
         }
         // First, invalidate all cached values.
-        mCachedEffectiveUserRestrictions.clear();
+        mCachedEffectiveUserRestrictions.removeAllRestrictions();
 
         // We don't want to call into ActivityManagerService while taking a lock, so we'll call
         // it on a handler.
@@ -2566,7 +2631,7 @@ public class UserManagerService extends IUserManager.Stub {
             synchronized (mRestrictionsLock) {
                 if (!UserRestrictionsUtils.isEmpty(oldGlobalUserRestrictions)
                         && mDeviceOwnerUserId != UserHandle.USER_NULL) {
-                    mDevicePolicyGlobalUserRestrictions.put(
+                    mDevicePolicyGlobalUserRestrictions.updateRestrictions(
                             mDeviceOwnerUserId, oldGlobalUserRestrictions);
                 }
                 // ENSURE_VERIFY_APPS is now enforced globally even if put by profile owner, so move
@@ -2691,7 +2756,8 @@ public class UserManagerService extends IUserManager.Stub {
 
         if (!restrictions.isEmpty()) {
             synchronized (mRestrictionsLock) {
-                mBaseUserRestrictions.append(UserHandle.USER_SYSTEM, restrictions);
+                mBaseUserRestrictions.updateRestrictions(UserHandle.USER_SYSTEM,
+                        restrictions);
             }
         }
 
@@ -2718,6 +2784,16 @@ public class UserManagerService extends IUserManager.Stub {
         }
     }
 
+    private void writeAllTargetUsersLP(int originatingUserId) {
+        for (int i = 0; i < mDevicePolicyLocalUserRestrictions.size(); i++) {
+            int targetUserId = mDevicePolicyLocalUserRestrictions.keyAt(i);
+            RestrictionsSet restrictionsSet = mDevicePolicyLocalUserRestrictions.valueAt(i);
+            if (restrictionsSet.containsKey(originatingUserId)) {
+                writeUserLP(getUserDataNoChecks(targetUserId));
+            }
+        }
+    }
+
     private void writeUserLP(UserData userData) {
         if (DBG) {
             debug("writeUserLP " + userData);
@@ -2802,12 +2878,11 @@ public class UserManagerService extends IUserManager.Stub {
         }
         synchronized (mRestrictionsLock) {
             UserRestrictionsUtils.writeRestrictions(serializer,
-                    mBaseUserRestrictions.get(userInfo.id), TAG_RESTRICTIONS);
+                    mBaseUserRestrictions.getRestrictions(userInfo.id), TAG_RESTRICTIONS);
+            getDevicePolicyLocalRestrictionsForTargetUserLR(userInfo.id).writeRestrictions(
+                    serializer, TAG_DEVICE_POLICY_LOCAL_RESTRICTIONS);
             UserRestrictionsUtils.writeRestrictions(serializer,
-                    mDevicePolicyLocalUserRestrictions.get(userInfo.id),
+                    mDevicePolicyGlobalUserRestrictions.getRestrictions(userInfo.id),
-                    TAG_DEVICE_POLICY_RESTRICTIONS);
-            UserRestrictionsUtils.writeRestrictions(serializer,
-                    mDevicePolicyGlobalUserRestrictions.get(userInfo.id),
                     TAG_DEVICE_POLICY_GLOBAL_RESTRICTIONS);
         }
 
@@ -2937,7 +3012,8 @@ public class UserManagerService extends IUserManager.Stub {
         String seedAccountType = null;
         PersistableBundle seedAccountOptions = null;
         Bundle baseRestrictions = null;
-        Bundle localRestrictions = null;
+        Bundle legacyLocalRestrictions = null;
+        RestrictionsSet localRestrictions = null;
         Bundle globalRestrictions = null;
 
         XmlPullParser parser = Xml.newPullParser();
@@ -3007,7 +3083,10 @@ public class UserManagerService extends IUserManager.Stub {
                 } else if (TAG_RESTRICTIONS.equals(tag)) {
                     baseRestrictions = UserRestrictionsUtils.readRestrictions(parser);
                 } else if (TAG_DEVICE_POLICY_RESTRICTIONS.equals(tag)) {
-                    localRestrictions = UserRestrictionsUtils.readRestrictions(parser);
+                    legacyLocalRestrictions = UserRestrictionsUtils.readRestrictions(parser);
+                } else if (TAG_DEVICE_POLICY_LOCAL_RESTRICTIONS.equals(tag)) {
+                    localRestrictions = RestrictionsSet.readRestrictions(parser,
+                            TAG_DEVICE_POLICY_LOCAL_RESTRICTIONS);
                 } else if (TAG_DEVICE_POLICY_GLOBAL_RESTRICTIONS.equals(tag)) {
                     globalRestrictions = UserRestrictionsUtils.readRestrictions(parser);
                 } else if (TAG_ACCOUNT.equals(tag)) {
@@ -3052,13 +3131,20 @@ public class UserManagerService extends IUserManager.Stub {
 
         synchronized (mRestrictionsLock) {
             if (baseRestrictions != null) {
-                mBaseUserRestrictions.put(id, baseRestrictions);
+                mBaseUserRestrictions.updateRestrictions(id, baseRestrictions);
             }
             if (localRestrictions != null) {
                 mDevicePolicyLocalUserRestrictions.put(id, localRestrictions);
+                if (legacyLocalRestrictions != null) {
+                    Slog.wtf(LOG_TAG, "Seeing both legacy and current local restrictions in xml");
+                }
+            } else if (legacyLocalRestrictions != null) {
+                mDevicePolicyLocalUserRestrictions.put(id,
+                        new RestrictionsSet(id, legacyLocalRestrictions));
             }
             if (globalRestrictions != null) {
-                mDevicePolicyGlobalUserRestrictions.put(id, globalRestrictions);
+                mDevicePolicyGlobalUserRestrictions.updateRestrictions(id,
+                        globalRestrictions);
             }
         }
         return userData;
@@ -3374,7 +3460,7 @@ public class UserManagerService extends IUserManager.Stub {
                 userTypeDetails.addDefaultRestrictionsTo(restrictions);
             }
             synchronized (mRestrictionsLock) {
-                mBaseUserRestrictions.append(userId, restrictions);
+                mBaseUserRestrictions.updateRestrictions(userId, restrictions);
             }
 
             t.traceBegin("PM.onNewUserCreated-" + userId);
@@ -3845,9 +3931,17 @@ public class UserManagerService extends IUserManager.Stub {
             mBaseUserRestrictions.remove(userId);
             mAppliedUserRestrictions.remove(userId);
             mCachedEffectiveUserRestrictions.remove(userId);
-            mDevicePolicyLocalUserRestrictions.remove(userId);
+            // Remove local restrictions affecting user
-            if (mDevicePolicyGlobalUserRestrictions.get(userId) != null) {
+            mDevicePolicyLocalUserRestrictions.delete(userId);
-                mDevicePolicyGlobalUserRestrictions.remove(userId);
+            // Remove local restrictions set by user
+            boolean changed = false;
+            for (int i = 0; i < mDevicePolicyLocalUserRestrictions.size(); i++) {
+                int targetUserId = mDevicePolicyLocalUserRestrictions.keyAt(i);
+                changed |= getDevicePolicyLocalRestrictionsForTargetUserLR(targetUserId)
+                        .remove(userId);
+            }
+            changed |= mDevicePolicyGlobalUserRestrictions.remove(userId);
+            if (changed) {
                 applyUserRestrictionsForAllUsersLR();
             }
         }
@@ -4540,16 +4634,18 @@ public class UserManagerService extends IUserManager.Stub {
                     pw.println("    Restrictions:");
                     synchronized (mRestrictionsLock) {
                         UserRestrictionsUtils.dumpRestrictions(
-                                pw, "      ", mBaseUserRestrictions.get(userInfo.id));
+                                pw, "      ", mBaseUserRestrictions.getRestrictions(userInfo.id));
                         pw.println("    Device policy global restrictions:");
                         UserRestrictionsUtils.dumpRestrictions(
-                                pw, "      ", mDevicePolicyGlobalUserRestrictions.get(userInfo.id));
+                                pw, "      ",
+                                mDevicePolicyGlobalUserRestrictions.getRestrictions(userInfo.id));
                         pw.println("    Device policy local restrictions:");
-                        UserRestrictionsUtils.dumpRestrictions(
+                        getDevicePolicyLocalRestrictionsForTargetUserLR(
-                                pw, "      ", mDevicePolicyLocalUserRestrictions.get(userInfo.id));
+                                userInfo.id).dumpRestrictions(pw, "      ");
                         pw.println("    Effective restrictions:");
                         UserRestrictionsUtils.dumpRestrictions(
-                                pw, "      ", mCachedEffectiveUserRestrictions.get(userInfo.id));
+                                pw, "      ",
+                                mCachedEffectiveUserRestrictions.getRestrictions(userInfo.id));
                     }
 
                     if (userData.account != null) {
@@ -4665,7 +4761,7 @@ public class UserManagerService extends IUserManager.Stub {
         @Override
         public Bundle getBaseUserRestrictions(@UserIdInt int userId) {
             synchronized (mRestrictionsLock) {
-                return mBaseUserRestrictions.get(userId);
+                return mBaseUserRestrictions.getRestrictions(userId);
             }
         }
 
@@ -4673,8 +4769,8 @@ public class UserManagerService extends IUserManager.Stub {
         public void setBaseUserRestrictionsByDpmsForMigration(
                 @UserIdInt int userId, Bundle baseRestrictions) {
             synchronized (mRestrictionsLock) {
-                if (updateRestrictionsIfNeededLR(
+                if (mBaseUserRestrictions.updateRestrictions(userId,
-                        userId, new Bundle(baseRestrictions), mBaseUserRestrictions)) {
+                        new Bundle(baseRestrictions))) {
                     invalidateEffectiveUserRestrictionsLR(userId);
                 }
             }

services/core/java/com/android/server/pm/UserRestrictionsUtils.java

@@ -394,22 +394,6 @@ public class UserRestrictionsUtils {
         }
     }
 
-    /**
-     * Merges a sparse array of restrictions bundles into one.
-     */
-    @Nullable
-    public static Bundle mergeAll(SparseArray<Bundle> restrictions) {
-        if (restrictions.size() == 0) {
-            return null;
-        } else {
-            final Bundle result = new Bundle();
-            for (int i = 0; i < restrictions.size(); i++) {
-                merge(result, restrictions.valueAt(i));
-            }
-            return result;
-        }
-    }
-
     /**
      * @return true if a restriction is settable by device owner.
      */
@@ -864,27 +848,15 @@ public class UserRestrictionsUtils {
     }
 
     /**
-     * Moves a particular restriction from one array of bundles to another, e.g. for all users.
+     * Moves a particular restriction from one array of restrictions sets to a restriction set,
+     * e.g. for all users.
      */
-    public static void moveRestriction(String restrictionKey, SparseArray<Bundle> srcRestrictions,
+    public static void moveRestriction(String restrictionKey,
-            SparseArray<Bundle> destRestrictions) {
+            SparseArray<RestrictionsSet> sourceRestrictionsSets,
-        for (int i = 0; i < srcRestrictions.size(); i++) {
+            RestrictionsSet destRestrictionSet) {
-            final int key = srcRestrictions.keyAt(i);
+        for (int i = 0; i < sourceRestrictionsSets.size(); i++) {
-            final Bundle from = srcRestrictions.valueAt(i);
+            final RestrictionsSet sourceRestrictionsSet = sourceRestrictionsSets.valueAt(i);
-            if (contains(from, restrictionKey)) {
+            sourceRestrictionsSet.moveRestriction(destRestrictionSet, restrictionKey);
-                from.remove(restrictionKey);
-                Bundle to = destRestrictions.get(key);
-                if (to == null) {
-                    to = new Bundle();
-                    destRestrictions.append(key, to);
-                }
-                to.putBoolean(restrictionKey, true);
-                // Don't keep empty bundles.
-                if (from.isEmpty()) {
-                    srcRestrictions.removeAt(i);
-                    i--;
-                }
-            }
         }
     }
 

services/tests/servicestests/src/com/android/server/devicepolicy/DpmTestUtils.java

@@ -25,6 +25,8 @@ import android.test.AndroidTestCase;
 import android.util.Log;
 import android.util.Printer;
 
+import com.android.server.pm.RestrictionsSet;
+
 import libcore.io.Streams;
 
 import com.google.android.collect.Lists;
@@ -35,7 +37,6 @@ import org.junit.Assert;
 
 import java.io.BufferedReader;
 import java.io.File;
-import java.io.FileNotFoundException;
 import java.io.FileOutputStream;
 import java.io.FileWriter;
 import java.io.IOException;
@@ -58,6 +59,13 @@ public class DpmTestUtils extends AndroidTestCase {
         return list == null ? 0 : list.size();
     }
 
+    public static RestrictionsSet newRestrictions(int userId, String... restrictions) {
+        Bundle localRestrictionsBundle = newRestrictions(restrictions);
+        RestrictionsSet localRestrictions = new RestrictionsSet();
+        localRestrictions.updateRestrictions(userId, localRestrictionsBundle);
+        return localRestrictions;
+    }
+
     public static Bundle newRestrictions(String... restrictions) {
         final Bundle ret = new Bundle();
         for (String restriction : restrictions) {
@@ -66,6 +74,17 @@ public class DpmTestUtils extends AndroidTestCase {
         return ret;
     }
 
+    public static void assertRestrictions(RestrictionsSet expected, RestrictionsSet actual) {
+        assertEquals(expected.size(), actual.size());
+
+        for (int i = 0; i < expected.size(); i++) {
+            int originatingUserId = expected.keyAt(i);
+            Bundle actualRestrictions = actual.getRestrictions(originatingUserId);
+            assertFalse(actualRestrictions.isEmpty());
+            assertRestrictions(expected.getRestrictions(originatingUserId), actualRestrictions);
+        }
+    }
+
     public static void assertRestrictions(Bundle expected, Bundle actual) {
         final ArrayList<String> elist;
         if (expected == null) {

services/tests/servicestests/src/com/android/server/pm/RestrictionsSetTest.java

@@ -0,0 +1,191 @@
+/*
+ * Copyright (C) 2020 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.pm;
+
+import static com.android.server.devicepolicy.DpmTestUtils.assertRestrictions;
+import static com.android.server.devicepolicy.DpmTestUtils.newRestrictions;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+
+import android.os.Bundle;
+import android.os.UserHandle;
+import android.os.UserManager;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.util.List;
+
+/** Test for {@link RestrictionsSet}. */
+@RunWith(AndroidJUnit4.class)
+public class RestrictionsSetTest {
+
+    private RestrictionsSet mRestrictionsSet = new RestrictionsSet();
+    private final int originatingUserId = 0;
+
+    @Test
+    public void testUpdateRestrictions_addRestrictions() {
+        Bundle restrictions = newRestrictions(UserManager.ENSURE_VERIFY_APPS);
+
+        assertTrue(mRestrictionsSet.updateRestrictions(originatingUserId, restrictions));
+
+        assertRestrictions(restrictions, mRestrictionsSet.getRestrictions(originatingUserId));
+    }
+
+    @Test
+    public void testUpdateRestrictions_removeRestrictions() {
+        Bundle restrictions = newRestrictions(UserManager.ENSURE_VERIFY_APPS);
+        mRestrictionsSet.updateRestrictions(originatingUserId, restrictions);
+
+        assertTrue(mRestrictionsSet.updateRestrictions(originatingUserId, new Bundle()));
+
+        assertNull(mRestrictionsSet.getRestrictions(originatingUserId));
+    }
+
+    @Test
+    public void testUpdateRestrictions_noChange() {
+        Bundle restrictions = newRestrictions(UserManager.ENSURE_VERIFY_APPS);
+        mRestrictionsSet.updateRestrictions(originatingUserId, restrictions);
+
+        assertFalse(mRestrictionsSet.updateRestrictions(originatingUserId, restrictions));
+    }
+
+    @Test
+    public void testMoveRestriction_containsRestriction() {
+        RestrictionsSet destRestrictionsSet = new RestrictionsSet();
+
+        String restriction = UserManager.DISALLOW_CONFIG_DATE_TIME;
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(restriction));
+
+        mRestrictionsSet.moveRestriction(destRestrictionsSet, restriction);
+
+        assertNull(mRestrictionsSet.getRestrictions(originatingUserId));
+        assertNotNull(destRestrictionsSet.getRestrictions(originatingUserId));
+        assertRestrictions(newRestrictions(restriction),
+                destRestrictionsSet.getRestrictions(originatingUserId));
+    }
+
+    @Test
+    public void testMoveRestriction_doesNotContainRestriction() {
+        RestrictionsSet destRestrictionsSet = new RestrictionsSet();
+
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(UserManager.ENSURE_VERIFY_APPS));
+
+        mRestrictionsSet.moveRestriction(destRestrictionsSet,
+                UserManager.DISALLOW_CONFIG_DATE_TIME);
+
+        assertRestrictions(newRestrictions(UserManager.ENSURE_VERIFY_APPS),
+                mRestrictionsSet.getRestrictions(originatingUserId));
+        assertNull(destRestrictionsSet.getRestrictions(originatingUserId));
+    }
+
+    @Test
+    public void testIsEmpty_noRestrictions() {
+        assertTrue(mRestrictionsSet.isEmpty());
+    }
+
+    @Test
+    public void testIsEmpty_hasRestrictions() {
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(UserManager.ENSURE_VERIFY_APPS,
+                        UserManager.DISALLOW_CONFIG_DATE_TIME));
+
+        assertFalse(mRestrictionsSet.isEmpty());
+    }
+
+    @Test
+    public void testMergeAll_noRestrictions() {
+        assertTrue(mRestrictionsSet.mergeAll().isEmpty());
+    }
+
+    @Test
+    public void testMergeAll_hasRestrictions() {
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(UserManager.ENSURE_VERIFY_APPS,
+                        UserManager.DISALLOW_CONFIG_DATE_TIME));
+        mRestrictionsSet.updateRestrictions(10,
+                newRestrictions(UserManager.DISALLOW_ADD_USER,
+                        UserManager.DISALLOW_AIRPLANE_MODE));
+
+        Bundle actual = mRestrictionsSet.mergeAll();
+        assertRestrictions(newRestrictions(UserManager.ENSURE_VERIFY_APPS,
+                UserManager.DISALLOW_CONFIG_DATE_TIME, UserManager.DISALLOW_ADD_USER,
+                UserManager.DISALLOW_AIRPLANE_MODE), actual);
+    }
+
+    @Test
+    public void testGetEnforcingUsers_hasEnforcingUser() {
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(UserManager.ENSURE_VERIFY_APPS));
+        mRestrictionsSet.updateRestrictions(10,
+                newRestrictions(UserManager.DISALLOW_ADD_USER));
+
+        List<UserManager.EnforcingUser> enforcingUsers = mRestrictionsSet.getEnforcingUsers(
+                UserManager.ENSURE_VERIFY_APPS, originatingUserId);
+
+        UserManager.EnforcingUser enforcingUser1 = enforcingUsers.get(0);
+        assertEquals(UserHandle.of(originatingUserId), enforcingUser1.getUserHandle());
+        assertEquals(UserManager.RESTRICTION_SOURCE_DEVICE_OWNER,
+                enforcingUser1.getUserRestrictionSource());
+    }
+
+    @Test
+    public void testGetEnforcingUsers_hasMultipleEnforcingUsers() {
+        int originatingUserId2 = 10;
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(UserManager.ENSURE_VERIFY_APPS));
+        mRestrictionsSet.updateRestrictions(originatingUserId2,
+                newRestrictions(UserManager.ENSURE_VERIFY_APPS));
+
+        List<UserManager.EnforcingUser> enforcingUsers = mRestrictionsSet.getEnforcingUsers(
+                UserManager.ENSURE_VERIFY_APPS, originatingUserId);
+
+        assertEquals(2, enforcingUsers.size());
+        for (UserManager.EnforcingUser enforcingUser : enforcingUsers) {
+            int userId = enforcingUser.getUserHandle().getIdentifier();
+            assertTrue((userId == originatingUserId) || (userId == originatingUserId2));
+            if (userId == originatingUserId) {
+                assertEquals(UserManager.RESTRICTION_SOURCE_DEVICE_OWNER,
+                        enforcingUser.getUserRestrictionSource());
+            }
+            if (userId == originatingUserId2) {
+                assertEquals(UserManager.RESTRICTION_SOURCE_PROFILE_OWNER,
+                        enforcingUser.getUserRestrictionSource());
+            }
+        }
+    }
+
+    @Test
+    public void testGetEnforcingUsers_noEnforcingUsers() {
+        mRestrictionsSet.updateRestrictions(originatingUserId,
+                newRestrictions(UserManager.DISALLOW_USER_SWITCH));
+
+        List<UserManager.EnforcingUser> enforcingUsers = mRestrictionsSet.getEnforcingUsers(
+                UserManager.ENSURE_VERIFY_APPS, originatingUserId);
+
+        assertTrue(enforcingUsers.isEmpty());
+    }
+
+}

services/tests/servicestests/src/com/android/server/pm/UserRestrictionsUtilsTest.java

@@ -247,60 +247,43 @@ public class UserRestrictionsUtilsTest extends AndroidTestCase {
         assertRestrictions(newRestrictions(UserManager.DISALLOW_CAMERA), local);
     }
 
-    public void testMergeAll() {
-        SparseArray<Bundle> restrictions = new SparseArray<>();
-        assertNull(UserRestrictionsUtils.mergeAll(restrictions));
-
-        restrictions.put(0, newRestrictions(UserManager.DISALLOW_ADJUST_VOLUME));
-        restrictions.put(1, newRestrictions(UserManager.DISALLOW_USB_FILE_TRANSFER));
-        restrictions.put(2, newRestrictions(UserManager.DISALLOW_APPS_CONTROL));
-
-        Bundle result = UserRestrictionsUtils.mergeAll(restrictions);
-        assertRestrictions(
-                newRestrictions(
-                        UserManager.DISALLOW_ADJUST_VOLUME,
-                        UserManager.DISALLOW_USB_FILE_TRANSFER,
-                        UserManager.DISALLOW_APPS_CONTROL),
-                result);
-    }
-
     public void testMoveRestriction() {
-        SparseArray<Bundle> localRestrictions = new SparseArray<>();
+        SparseArray<RestrictionsSet> localRestrictions = new SparseArray<>();
-        SparseArray<Bundle> globalRestrictions = new SparseArray<>();
+        RestrictionsSet globalRestrictions = new RestrictionsSet();
 
         // User 0 has only local restrictions, nothing should change.
-        localRestrictions.put(0, newRestrictions(UserManager.DISALLOW_ADJUST_VOLUME));
+        localRestrictions.put(0, newRestrictions(0, UserManager.DISALLOW_ADJUST_VOLUME));
         // User 1 has a local restriction to be moved to global and some global already. Local
         // restrictions should be removed for this user.
-        localRestrictions.put(1, newRestrictions(UserManager.ENSURE_VERIFY_APPS));
+        localRestrictions.put(1, newRestrictions(1, UserManager.ENSURE_VERIFY_APPS));
-        globalRestrictions.put(1, newRestrictions(UserManager.DISALLOW_ADD_USER));
+        globalRestrictions.updateRestrictions(1,
+                newRestrictions(UserManager.DISALLOW_ADD_USER));
         // User 2 has a local restriction to be moved and one to leave local.
-        localRestrictions.put(2, newRestrictions(
+        localRestrictions.put(2, newRestrictions(2,
-                UserManager.ENSURE_VERIFY_APPS,
+                UserManager.ENSURE_VERIFY_APPS, UserManager.DISALLOW_CONFIG_VPN));
-                UserManager.DISALLOW_CONFIG_VPN));
 
         UserRestrictionsUtils.moveRestriction(
                 UserManager.ENSURE_VERIFY_APPS, localRestrictions, globalRestrictions);
 
         // Check user 0.
         assertRestrictions(
-                newRestrictions(UserManager.DISALLOW_ADJUST_VOLUME),
+                newRestrictions(0, UserManager.DISALLOW_ADJUST_VOLUME),
                 localRestrictions.get(0));
-        assertNull(globalRestrictions.get(0));
+        assertNull(globalRestrictions.getRestrictions(0));
 
         // Check user 1.
-        assertNull(localRestrictions.get(1));
+        assertTrue(localRestrictions.get(1).isEmpty());
         assertRestrictions(
                 newRestrictions(UserManager.ENSURE_VERIFY_APPS, UserManager.DISALLOW_ADD_USER),
-                globalRestrictions.get(1));
+                globalRestrictions.getRestrictions(1));
 
         // Check user 2.
         assertRestrictions(
-                newRestrictions(UserManager.DISALLOW_CONFIG_VPN),
+                newRestrictions(2, UserManager.DISALLOW_CONFIG_VPN),
                 localRestrictions.get(2));
         assertRestrictions(
                 newRestrictions(UserManager.ENSURE_VERIFY_APPS),
-                globalRestrictions.get(2));
+                globalRestrictions.getRestrictions(2));
     }
 
     public void testAreEqual() {