Merge "Check userId validity when setting user restrictions"

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

@@ -3044,7 +3044,12 @@ class PackageManagerShellCommand extends ShellCommand {
                 translateUserId(userId, UserHandle.USER_NULL, "runSetUserRestriction");
         final IUserManager um = IUserManager.Stub.asInterface(
                 ServiceManager.getService(Context.USER_SERVICE));
-        um.setUserRestriction(restriction, value, translatedUserId);
+        try {
+            um.setUserRestriction(restriction, value, translatedUserId);
+        } catch (IllegalArgumentException e) {
+            getErrPrintWriter().println(e.getMessage());
+            return 1;
+        }
         return 0;
     }
 

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

@@ -2668,6 +2668,11 @@ public class UserManagerService extends IUserManager.Stub {
         if (!UserRestrictionsUtils.isValidRestriction(key)) {
             return;
         }
+
+        if (!userExists(userId)) {
+            throw new IllegalArgumentException("Cannot set user restriction. "
+                    + "User with this id does not exist");
+        }
         synchronized (mRestrictionsLock) {
             // Note we can't modify Bundles stored in mBaseUserRestrictions directly, so create
             // a copy.
@@ -3251,6 +3256,22 @@ public class UserManagerService extends IUserManager.Stub {
         }
     }
 
+    /**
+     * Checks whether user with a given ID exists.
+     * @param id User id to be checked.
+     */
+    @VisibleForTesting
+    boolean userExists(int id) {
+        synchronized (mUsersLock) {
+            for (int userId : mUserIds) {
+                if (userId == id) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     /**
      * Returns an array of user ids, including pre-created users.
      *
@@ -3774,7 +3795,6 @@ public class UserManagerService extends IUserManager.Stub {
             }
         }
 
-        updateUserIds();
         initDefaultGuestRestrictions();
 
         writeUserLP(userData);
@@ -4971,6 +4991,7 @@ public class UserManagerService extends IUserManager.Stub {
         synchronized (mUsersLock) {
             mUsers.put(userInfo.id, userData);
         }
+        updateUserIds();
         return userData;
     }
 

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

@@ -16,6 +16,7 @@
 
 package com.android.server.pm;
 
+import static android.os.UserManager.DISALLOW_BLUETOOTH;
 import static android.os.UserManager.DISALLOW_USER_SWITCH;
 
 import static com.google.common.truth.Truth.assertThat;
@@ -40,6 +41,7 @@ import androidx.test.runner.AndroidJUnit4;
 import com.android.server.LocalServices;
 
 import org.junit.After;
+import org.junit.Assert;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -173,6 +175,21 @@ public class UserManagerServiceTest {
         assertThat(mUserManagerService.hasUserRestriction(DISALLOW_USER_SWITCH, userId)).isFalse();
     }
 
+    @Test
+    public void testSetUserRestrictionWithIncorrectID() throws Exception {
+        int incorrectId = 1;
+        while (mUserManagerService.userExists(incorrectId)) {
+            incorrectId++;
+        }
+        try {
+            mUserManagerService.setUserRestriction(DISALLOW_BLUETOOTH, true, incorrectId);
+            Assert.fail();
+        } catch (IllegalArgumentException e) {
+            //Exception is expected to be thrown if user ID does not exist.
+            // IllegalArgumentException thrown means this test is successful.
+        }
+    }
+
     @Test
     public void assertIsUserSwitcherEnabledOnMultiUserSettings() throws Exception {
         int userId = ActivityManager.getCurrentUser();