Disallow work profiles on non-initial users in headless

Modified 'canHaveProfile' method to ensure we allow work profiles only on initial user and not additional users- in both headless and non-headless environment.For that we can use isMain() as on non-headless, the system user(initial) is main,and on headless the first secondary user(initial) is main.

Test: atest
Bug: 266053489
Fix: 266053489

Change-Id: Ia959213b03ea7b36d113fd13e23fca187fe3f59d
diff --git a/core/java/android/content/pm/UserInfo.java b/core/java/android/content/pm/UserInfo.java
index e38cb65..333066b 100644
--- a/core/java/android/content/pm/UserInfo.java
+++ b/core/java/android/content/pm/UserInfo.java
@@ -436,11 +436,7 @@
         if (isProfile() || isGuest() || isRestricted()) {
             return false;
         }
-        if (UserManager.isHeadlessSystemUserMode()) {
-            return id != UserHandle.USER_SYSTEM;
-        } else {
-            return id == UserHandle.USER_SYSTEM;
-        }
+        return isMain();
     }
 
     // TODO(b/142482943): Get rid of this (after removing it from all tests) if feasible.
diff --git a/services/tests/servicestests/src/com/android/server/pm/UserManagerServiceUserInfoTest.java b/services/tests/servicestests/src/com/android/server/pm/UserManagerServiceUserInfoTest.java
index 92fddc7..d999aa3 100644
--- a/services/tests/servicestests/src/com/android/server/pm/UserManagerServiceUserInfoTest.java
+++ b/services/tests/servicestests/src/com/android/server/pm/UserManagerServiceUserInfoTest.java
@@ -22,6 +22,7 @@
 import static android.content.pm.UserInfo.FLAG_FULL;
 import static android.content.pm.UserInfo.FLAG_GUEST;
 import static android.content.pm.UserInfo.FLAG_INITIALIZED;
+import static android.content.pm.UserInfo.FLAG_MAIN;
 import static android.content.pm.UserInfo.FLAG_MANAGED_PROFILE;
 import static android.content.pm.UserInfo.FLAG_PROFILE;
 import static android.content.pm.UserInfo.FLAG_RESTRICTED;
@@ -206,6 +207,13 @@
         assertFalse("Switching to a profiles should be disabled", userInfo.supportsSwitchTo());
     }
 
+    /** Test UserInfo.canHaveProfile for main user */
+    @Test
+    public void testCanHaveProfile() throws Exception {
+        UserInfo userInfo = createUser(100, FLAG_MAIN, null);
+        assertTrue("Main users can have profile", userInfo.canHaveProfile());
+    }
+
     /** Tests upgradeIfNecessaryLP (but without locking) for upgrading from version 8 to 9+. */
     @Test
     public void testUpgradeIfNecessaryLP_9() {